Scala groupBy and partition, in depth: the whole splitting family, what each one traverses and allocates, lazy collections, key traps and a worked import pipeline

Two questions pick the method

Ask two questions. Does every element go to one of two places, or one of many? And is the destination decided by a predicate, by position, or by a computed key? Most splitting bugs come from picking a method from the wrong row below.

partition: two sides, one pass on strict collections

partition(p) returns every element that satisfies p on the left and every element that does not on the right, each in original order. It is defined to equal (filter(p), filterNot(p)), and the generic implementation in IterableOps really is that: two filtered views, so the source is traversed twice and the predicate runs twice per element. Strict collections such as List, Vector and ArrayBuffer mix in StrictOptimizedIterableOps, whose override walks the input once and pushes each element into one of two builders. For everyday code on strict collections, then, partition is one pass and there is no reason to hand-roll a fold.

Lazy collections are different, and the difference is easy to miss. LazyList.partition is literally (filter(p), filterNot(p)): nothing runs until you consume a side, each side tests the elements it walks over, and because a LazyList memoises its cells, holding both sides keeps the source alive. Iterator.partition is built on duplicate: consuming one side buffers the elements the other side has not yet seen. If you drain the left side of a large iterator first, the right side's entire contents sit in memory.

val xs = List(5, 12, 7, 30, 1)

// strict: one traversal, the predicate runs once per element
val (big, small) = xs.partition(_ >= 10)      // (List(12, 30), List(5, 7, 1))

// LazyList: defined as (filter(p), filterNot(p)), so nothing runs yet,
// and each side re-tests elements it walks over
val ll = LazyList.from(1).map { n => println(s"produce $n"); n }
val (evens, odds) = ll.partition(_ % 2 == 0)
evens.take(2).toList     // forces only what the even side needs

// Iterator: built on duplicate, which buffers what one side has seen
// and the other has not; do not use the original iterator afterwards
val (a, b) = Iterator.range(0, 1_000_000).partition(_ < 10)
a.toList                 // walks the whole source, buffering 999,990 Ints for b

partitionMap: split and convert in one step

A predicate only tells you which side an element belongs on. Often you also want to change the element on the way, for example turn a raw row into either an error or a parsed record. partitionMap takes a function returning Either, puts the Left values on the left and the Right values on the right, and on strict collections does it in a single pass. It replaces mapping to Either, partitioning on isLeft and unwrapping both sides.

final case class RawRow(line: Int, customer: String, sku: String, qty: String)
final case class Order(customer: String, sku: String, qty: Int)
final case class RowError(line: Int, reason: String)

def parse(r: RawRow): Either[RowError, Order] =
  r.qty.toIntOption match
    case None                  => Left(RowError(r.line, s"qty '${r.qty}' is not a number"))
    case Some(q) if q <= 0     => Left(RowError(r.line, s"qty $q must be positive"))
    case Some(_) if r.sku.isEmpty => Left(RowError(r.line, "missing sku"))
    case Some(q)               => Right(Order(r.customer, r.sku, q))

val (errors, orders) = rows.partitionMap(parse)   // (List[RowError], List[Order]), one pass

This is the right tool when you want to keep every error, not stop at the first one. If one failure should abort the whole batch, use traverse-style sequencing instead; the trade-off between accumulating and failing fast is covered in Scala Either in depth.

span and splitAt: position, not membership

span(p) returns the longest prefix whose elements all satisfy p, and everything after it, including later elements that would have matched. On strict collections it is one pass and stops calling the predicate at the first failure. splitAt(n) cuts at an index and never looks at the values. Both are defined as (takeWhile(p), dropWhile(p)) and (take(n), drop(n)) respectively.

val readings = List(3, 4, 4, 9, 2, 1)            // already ordered by time

readings.span(_ < 5)       // (List(3, 4, 4), List(9, 2, 1))  prefix, then everything after
readings.partition(_ < 5)  // (List(3, 4, 4, 2, 1), List(9))   every match, anywhere
readings.splitAt(2)        // (List(3, 4), List(4, 9, 2, 1))   by index, no predicate

span is correct only when the input is ordered so that every matching element comes first: log lines before a cutoff time, sorted scores above a threshold, a header block before data. Used on unordered input it silently returns a short prefix. Comment why the input is ordered wherever you use it.

groupBy: from two outputs to many

groupBy(k) computes a key for every element and returns an immutable.Map from each distinct key to the elements that produced it. The library implementation walks the input once, keeps a mutable.Map from key to a builder of the source collection type, appends each element to its key's builder, and finally copies the results into an immutable.HashMap. Three consequences follow.

• Element order inside each group is preserved. Builders append in traversal order, so groupBy on a time-sorted list gives time-sorted groups.
• Key order is not. The result is a hash map, so iterating it gives an order that depends on hash codes and map size. Sort explicitly or convert to SortedMap when order matters.
• The values have the source's type. A List gives List groups, a Vector gives Vector groups, a String gives Map[K, String], and a Map gives a map of maps.

The 2.13 additions exist because most callers do not want whole elements. groupMap(k)(f) stores f(elem) instead of the element. groupMapReduce(k)(f)(op) keeps only one accumulated value per key: the first element contributes f(elem) and each later one is combined as op(acc, f(elem)), strictly in traversal order. That means op does not need to be commutative here, but it should be if you ever expect the same code to run on a parallel or distributed collection.

// groupBy keeps whole elements; the value type matches the source collection
val byCustomer: Map[String, List[Order]] = orders.groupBy(_.customer)

// groupMap keeps only what you need per element
val skusByCustomer: Map[String, List[String]] = orders.groupMap(_.customer)(_.sku)

// groupMapReduce never builds the per-key lists at all
val unitsByCustomer: Map[String, Int] = orders.groupMapReduce(_.customer)(_.qty)(_ + _)

// a stable, readable order is something you ask for explicitly
import scala.collection.immutable.SortedMap
val report = SortedMap.from(unitsByCustomer)     // keys in sorted order

grouped and sliding: chunks and windows

grouped(n) cuts a collection into consecutive chunks of n elements; the last chunk is smaller if the size does not divide evenly. sliding(n, step) produces overlapping windows of size n that advance by step, with a default step of 1. Both return an Iterator, so nothing is computed until you consume it and the result can be traversed only once. Store .toList if you need it twice.

Use grouped for batching side effects: database inserts, API calls with a page limit, messages to a broker. Use sliding(2) for pairwise work such as deltas between consecutive readings: readings.sliding(2).collect { case List(a, b) => b - a }. Matching on the window's shape matters, because a collection shorter than the window size yields one short window instead of none, and code that assumes exactly two elements will throw on it.

Worked example: an import pipeline

Suppose a nightly job loads order lines from a partner's CSV. Each row may be malformed, the same customer and SKU sometimes appear on several lines, the business wants units per customer in a stable order, and the database rejects transactions above a few thousand rows. Each requirement corresponds to one member of the family, and the whole job stays a sequence of readable steps.

import scala.collection.immutable.SortedMap

final case class ImportReport(
    errors: List[RowError],
    unitsByCustomer: SortedMap[String, Int],
    duplicateLines: Map[(String, String), List[Int]],
    batches: Int)

def runImport(rows: List[RawRow], insert: Seq[Order] => Unit): ImportReport =
  // 1. split good from bad in one pass, keeping both
  val (errors, orders) = rows.partitionMap(parse)

  // 2. detect the same (customer, sku) appearing on several lines
  val linesByKey = rows
    .filter(r => parse(r).isRight)
    .groupMap(r => (r.customer, r.sku))(_.line)
  val duplicates = linesByKey.filter((_, lines) => lines.sizeIs > 1)

  // 3. totals per customer without intermediate lists
  val units = SortedMap.from(orders.groupMapReduce(_.customer)(_.qty)(_ + _))

  // 4. write in fixed-size chunks so one transaction never holds 200,000 rows
  var n = 0
  orders.grouped(500).foreach { chunk => insert(chunk); n += 1 }

  ImportReport(errors, units, duplicates, n)

The code uses Scala 3 syntax; on 2.13, write the match with braces and the map filter as filter { case (_, lines) => lines.sizeIs > 1 }.

Trace it with 200,000 rows, 1,200 malformed. partitionMap makes one pass and yields 1,200 errors and 198,800 orders. groupMap stores only line numbers per key, not whole rows. groupMapReduce holds one Int per customer, so ten thousand customers cost ten thousand map entries rather than 198,800 list cells. grouped(500) produces 398 insert calls, the last with 300 rows. The SortedMap makes the report diffable night to night.

Memory, cost and big data

groupBy materialises the entire input as groups, so peak memory is roughly the input plus one builder per key plus the final map. When you only need an aggregate, groupMapReduce is the single biggest win, because no per-key collection is ever built. When the data does not fit in memory at all, stream it: fold an Iterator into a mutable.HashMap of accumulators, or move the job to a distributed engine. For measured collection costs, see Scala collections performance.

The same reasoning carries into Spark, where the names look familiar but the cost is network traffic. An RDD groupByKey ships every value for a key across the cluster before you can aggregate, while reduceByKey combines on the map side first, which is the distributed analogue of choosing groupMapReduce over groupBy. Spark shuffle in depth explains what that shuffle costs.

Failure modes

These are the bugs that reach production most often.

// Array keys use reference equality, so equal contents land in different groups
List(Array(1, 2), Array(1, 2)).groupBy(identity).size        // 2, not 1
List(Array(1, 2), Array(1, 2)).groupBy(_.toList).size        // 1: key on an immutable copy

// groupMap on a Set builds Set values, so equal mapped values collapse
case class Emp(name: String, dept: String, salary: Int)
val staff = Set(Emp("a", "ops", 100), Emp("b", "ops", 100))
staff.groupMap(_.dept)(_.salary)                             // Map(ops -> Set(100))
staff.toList.groupMap(_.dept)(_.salary)                      // Map(ops -> List(100, 100))

• Relying on map order: printing, serialising or comparing a groupBy result by iteration order. Fix: sort or use SortedMap at the boundary.
• Keys with identity equality: arrays and classes without equals group by reference. Fix: key on case classes, tuples or immutable copies.
• Over-fine keys: grouping by a timestamp that includes seconds when you meant the day yields one group per event. Fix: truncate in the key function and test with two events on the same day.
• Set semantics leaking in: groupMap on a Set deduplicates values. Fix: convert to a List first if duplicates count.
• Missing groups: groupBy(p)(true) throws when nothing matched. Fix: use partition for two-way splits, or getOrElse(key, Nil).
• Lazy double work: an expensive or effectful predicate on LazyList.partition runs once per side. Fix: compute the test once with map to a pair, or materialise first.

Trade-offs and a decision rule

The general rule: keep as little per group as the question needs, ask for order explicitly, and know whether your collection is strict before you rely on a pass count.

What to do next

1. Search your codebase for groupBy( followed by .map or .mapValues and replace each with groupMap or groupMapReduce.
2. Find every place a groupBy result is printed, serialised or compared, and add an explicit sort or SortedMap.
3. Replace boolean groupBy calls with partition, and map-then-partition-on-isLeft with partitionMap.
4. Add a comment to every span stating why the input is ordered, and a test that fails if the order changes.
5. Check that no grouping key is an array, a mutable object or a class without structural equality.
6. Batch every bulk side effect with grouped(n) and pick n from the target system's limits.
7. Before using partition on a LazyList or Iterator, confirm the predicate is cheap and pure, or materialise first.