Skip to content

Instantly share code, notes, and snippets.

@IvanVergiliev

IvanVergiliev/tree_fold.scala

Created June 17, 2021 20:39
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save IvanVergiliev/4b648a476aec4387b076f308fa294bcb to your computer and use it in GitHub Desktop.
Save IvanVergiliev/4b648a476aec4387b076f308fa294bcb to your computer and use it in GitHub Desktop.
Download ZIP
Balanced fold
Raw
tree_fold.scala
import org.apache.spark.sql.Column
def treeFold(
  columns: List[Column],
  initialValue: Column,
  op: (Column, Column) => Column
): Column = {
  if (columns.isEmpty) {
    initialValue
  } else if (columns.length == 1) {
    columns.head
  } else {
    val (left, right) = columns.splitAt(columns.length / 2)
    op(
      treeFold(left, initialValue, op),
      treeFold(right, initialValue, op)
    )
  }
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment