diff --git a/core/src/main/scala/flatgraph/traversal/Language.scala b/core/src/main/scala/flatgraph/traversal/Language.scala index 0e78e13d..2a5bf847 100644 --- a/core/src/main/scala/flatgraph/traversal/Language.scala +++ b/core/src/main/scala/flatgraph/traversal/Language.scala @@ -3,6 +3,7 @@ package flatgraph.traversal import flatgraph.help.{Doc, Traversal} import flatgraph.{Accessors, Edge, GNode, MultiPropertyKey, OptionalPropertyKey, PropertyKey, Schema, SinglePropertyKey} +import java.util import scala.annotation.implicitNotFound import scala.collection.immutable.ArraySeq import scala.collection.{Iterator, mutable} @@ -57,7 +58,24 @@ class GenericSteps[A](iterator: Iterator[A]) extends AnyVal { counts.to(Map) } - def groupBy[K](f: A => K): Map[K, List[A]] = l.groupBy(f) + /** Execute the traversal and group elements by a given transformation function, ignoring the iterator order. Use is discouraged. */ + @Doc(info = + "Execute the traversal and group elements by a given transformation function, ignoring the iterator order. Use is discouraged, because iteration order is not reproducible, which tends to produce very bad bugs." + ) + def groupBy[K](f: A => K): Map[K, List[A]] = l.groupBy(f) + + /** Execute the traversal and group elements by a given transformation function, respecting the order of the iterator */ + @Doc(info = "Execute the traversal and group elements by a given transformation function, respecting the order of the iterator.") + def groupByStable[K](f: A => K): scala.collection.SeqMap[K, scala.collection.Seq[A]] = { + val res = mutable.LinkedHashMap[K, mutable.ArrayBuffer[A]]() + while (iterator.hasNext) { + val item = iterator.next + val key = f(item) + res.getOrElseUpdate(key, mutable.ArrayBuffer[A]()).addOne(item) + } + res + } + def groupMap[K, B](key: A => K)(f: A => B): Map[K, List[B]] = l.groupMap(key)(f) def groupMapReduce[K, B](key: A => K)(f: A => B)(reduce: (B, B) => B): Map[K, B] = l.groupMapReduce(key)(f)(reduce) diff --git a/core/src/test/scala/flatgraph/traversal/IterableOnceExtensionTests.scala b/core/src/test/scala/flatgraph/traversal/IterableOnceExtensionTests.scala index 3108c581..34d81656 100644 --- a/core/src/test/scala/flatgraph/traversal/IterableOnceExtensionTests.scala +++ b/core/src/test/scala/flatgraph/traversal/IterableOnceExtensionTests.scala @@ -47,4 +47,11 @@ class IterableOnceExtensionTests extends AnyWordSpec with Matchers { Seq(1, 2).loneElementOption shouldBe None } + "groupBy respects iteration order" in { + Seq(5, 4, 3, 2, 1, 0).groupByStable(x => x % 3).valuesIterator.map { _.l }.l shouldBe List(List(5, 2), List(4, 1), List(3, 0)) + Seq(5, 4, 3, 2, 1, 0).groupByStable(x => (x + 1) % 3).valuesIterator.map { _.l }.l shouldBe List(List(5, 2), List(4, 1), List(3, 0)) + Seq(5, 4, 3, 2, 1, 0).groupByStable(x => (x + 2) % 3).valuesIterator.map { _.l }.l shouldBe List(List(5, 2), List(4, 1), List(3, 0)) + + } + }