Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Optimize distinctBy implementation for non-empty collections #4711

Merged
merged 4 commits into from
Feb 1, 2025
+42 −30
Merged

Optimize distinctBy implementation for non-empty collections #4711

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
572fbfa
Optimize distinctBy implementation for NonEmptyList, NonEmptySeq and …
Jan 31, 2025
44b076d
Improve optimizations
Feb 1, 2025
b02647f
Fix compatibility problem
Feb 1, 2025
19f2d92
Move Ordering implicit into else
Feb 1, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 8 additions & 9 deletions core/src/main/scala/cats/data/NonEmptyList.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -345,16 +345,15 @@ final case class NonEmptyList[+A](head: A, tail: List[A]) extends NonEmptyCollec
override def distinctBy[B](f: A => B)(implicit O: Order[B]): NonEmptyList[A] = {
implicit val ord: Ordering[B] = O.toOrdering
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

actually this could be moved inside the else branch to avoid it on singletons.


val buf = ListBuffer.empty[A]
tail.foldLeft(TreeSet(f(head): B)) { (elementsSoFar, a) =>
val b = f(a)
if (elementsSoFar(b)) elementsSoFar
else {
buf += a; elementsSoFar + b
}
val bldr = List.newBuilder[A]
val seen = mutable.TreeSet.empty[B]
val it = iterator
while (it.hasNext) {
val next = it.next()
if (seen.add(f(next)))
bldr += next
}

NonEmptyList(head, buf.toList)
NonEmptyList.fromListUnsafe(bldr.result())
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it would be slightly better to do:

var rest = tail
seen.add(f(head))
while (rest.nonEmpty) {
  val next = rest.head
  rest = rest.tail
  if (seen.add(f(next)) {
    bldr += next
  }
}

NonEmptyList(head, bldr.result())

the main changes:

  1. we don't have to build an iterator for the list, we just directly iterate on List
  2. we don't have to allocate the outermost List just to destructure it inside NonEmptyList.fromListUnsafe.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'd also suggest to check if the source NeList has 1 item only (tail.isEmpty or use sizeCompare) and simply return this in that case.
It would help to avoid allocating unnecessary structures: the builder and tree set and recreating 1-item NeList.

This small optimization could also be applied to the other two collections.

}

/**
Expand Down
17 changes: 8 additions & 9 deletions core/src/main/scala/cats/data/NonEmptySeq.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -241,16 +241,15 @@ final class NonEmptySeq[+A] private (val toSeq: Seq[A]) extends AnyVal with NonE
override def distinctBy[B](f: A => B)(implicit O: Order[B]): NonEmptySeq[A] = {
implicit val ord: Ordering[B] = O.toOrdering

val buf = Seq.newBuilder[A]
tail.foldLeft(TreeSet(f(head): B)) { (elementsSoFar, a) =>
val b = f(a)
if (elementsSoFar(b)) elementsSoFar
else {
buf += a; elementsSoFar + b
}
val bldr = Seq.newBuilder[A]
val seen = mutable.TreeSet.empty[B]
val it = iterator
while (it.hasNext) {
val next = it.next()
if (seen.add(f(next)))
bldr += next
}

NonEmptySeq(head, buf.result())
NonEmptySeq.fromSeqUnsafe(bldr.result())
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this one is different from NonEmptyList, because it is a wrapper on the whole Seq, unlike NonEmptyList which is a class.

Given that, this implementation here seems good to me.

}

/**
Expand Down
17 changes: 8 additions & 9 deletions core/src/main/scala/cats/data/NonEmptyVector.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -251,16 +251,15 @@ final class NonEmptyVector[+A] private (val toVector: Vector[A])
override def distinctBy[B](f: A => B)(implicit O: Order[B]): NonEmptyVector[A] = {
implicit val ord: Ordering[B] = O.toOrdering

val buf = Vector.newBuilder[A]
tail.foldLeft(TreeSet(f(head): B)) { (elementsSoFar, a) =>
val b = f(a)
if (elementsSoFar(b)) elementsSoFar
else {
buf += a; elementsSoFar + b
}
val bldr = Vector.newBuilder[A]
val seen = mutable.TreeSet.empty[B]
val it = iterator
while (it.hasNext) {
val next = it.next()
if (seen.add(f(next)))
bldr += next
}

NonEmptyVector(head, buf.result())
NonEmptyVector.fromVectorUnsafe(bldr.result())
}

/**
Expand Down
Loading