[WIP] Window functions and column sorting

core/src/main/scala/frameless/CatalystRowOrdered.scala

@@ -0,0 +1,45 @@
+package frameless
+
+import shapeless._
+
+import scala.annotation.implicitNotFound
+
+/** Types that can be used to sort a dataset by Catalyst. */
+@implicitNotFound("Cannot order by columns of type ${A}.")
+trait CatalystRowOrdered[A]
+
+object CatalystRowOrdered {
+  private[this] val theInstance = new CatalystRowOrdered[Any] {}
+  private[this] def of[A]: CatalystRowOrdered[A] = theInstance.asInstanceOf[CatalystRowOrdered[A]]
+
+  /*
+  The following are sortable by spark:
+  see [[org.apache.spark.sql.catalyst.expressions.RowOrdering.isOrderable]]
+    AtomicType
+    StructType containing only orderable types
+    ArrayType containing only orderable types
+    UserDefinedType containing only orderable types
+
+  MapType can't be used in order!
+  TODO: UDF
+   */
+
+  implicit def orderedEvidence[A](implicit catalystOrdered: CatalystOrdered[A]): CatalystRowOrdered[A] = of[A]
+
+  implicit def arrayEv[A](implicit catalystOrdered: CatalystRowOrdered[A]): CatalystRowOrdered[Array[A]] = of[Array[A]]
+
+  implicit def collectionEv[C[X] <: Seq[X], A](implicit catalystOrdered: CatalystRowOrdered[A]): CatalystRowOrdered[C[A]] = of[C[A]]
+
+  implicit def optionEv[A](implicit catalystOrdered: CatalystRowOrdered[A]): CatalystRowOrdered[Option[A]] = of[Option[A]]
+
+  implicit def recordEv[A, G <: HList](implicit i0: Generic.Aux[A, G], i1: HasRowOrdered[G]): CatalystRowOrdered[A] = of[A]
+
+  trait HasRowOrdered[T <: HList]
+  object HasRowOrdered {
+    implicit def deriveOrderHNil[H](implicit catalystRowOrdered: CatalystRowOrdered[H]): HasRowOrdered[H :: HNil] =
+      new HasRowOrdered[H :: HNil] {}
+
+    implicit def deriveOrderHCons[H, T <: HList](implicit head: CatalystRowOrdered[H], tail: HasRowOrdered[T]): HasRowOrdered[H :: T] =
+      new HasRowOrdered[H :: T] {}
+  }
+}

dataset/src/main/scala/frameless/TypedColumn.scala

@@ -278,6 +278,60 @@ sealed class TypedColumn[T, U](
     */
   def /(u: U)(implicit n: CatalystNumeric[U]): TypedColumn[T, Double] = self.untyped.divide(u).typed
 
+  /** Returns a descending ordering used in sorting
+    *
+    * apache/spark
+    */
+  def desc(implicit catalystRowOrdering: CatalystRowOrdered[U]): TypedSortedColumn[T, U] =
+    new TypedSortedColumn[T, U](withExpr {
+      SortOrder(expr, Descending)
+    })
+
+  /** Returns a descending ordering used in sorting where None values appear before non-None values
+    *
+    * apache/spark
+    */
+  def descNullsFirst(implicit isOption: U <:< Option[_], catalystRowOrdering: CatalystRowOrdered[U]): TypedSortedColumn[T, U] =
+    new TypedSortedColumn[T, U](withExpr {
+      SortOrder(expr, Descending, NullsFirst, Set.empty)
+    })
+
+  /** Returns a descending ordering used in sorting where None values appear after non-None values
+    *
+    * apache/spark
+    */
+  def descNullsLast(implicit isOption: U <:< Option[_], catalystRowOrdering: CatalystRowOrdered[U]): TypedSortedColumn[T, U] =
+    new TypedSortedColumn[T, U](withExpr {
+      SortOrder(expr, Descending, NullsLast, Set.empty)
+    })
+
+  /** Returns an ascending ordering used in sorting
+    *
+    * apache/spark
+    */
+  def asc(implicit catalystRowOrdering: CatalystRowOrdered[U]): TypedSortedColumn[T, U] =
+    new TypedSortedColumn[T, U](withExpr {
+      SortOrder(expr, Ascending)
+    })
+
+  /** Returns an ascending ordering used in sorting where None values appear before non-None values
+    *
+    * apache/spark
+    */
+  def ascNullsFirst(implicit isOption: U <:< Option[_], catalystRowOrdering: CatalystRowOrdered[U]): TypedSortedColumn[T, U] =
+    new TypedSortedColumn[T, U](withExpr {
+      SortOrder(expr, Ascending, NullsFirst, Set.empty)
+    })
+
+  /** Returns an ascending ordering used in sorting where None values appear after non-None values
+    *
+    * apache/spark
+    */
+  def ascNullsLast(implicit isOption: U <:< Option[_], catalystRowOrdering: CatalystRowOrdered[U]): TypedSortedColumn[T, U] =
+    new TypedSortedColumn[T, U](withExpr {
+      SortOrder(expr, Ascending, NullsLast, Set.empty)
+    })
+
   /**
     * Bitwise AND this expression and another expression.
     * {{{
@@ -485,6 +539,28 @@ sealed class TypedAggregate[T, U](val expr: Expression)(
   }
 }
 
+sealed class TypedSortedColumn[T, U](val expr: Expression)(
+  implicit
+  val uencoder: TypedEncoder[U]
+) extends UntypedExpression[T] {
+
+  def this(column: Column)(implicit e: TypedEncoder[U]) {
+    this(FramelessInternals.expr(column))
+  }
+
+  def untyped: Column = new Column(expr)
+}
+
+object TypedSortedColumn {
+  implicit def defaultAscending[T, U : CatalystRowOrdered](typedColumn: TypedColumn[T, U]): TypedSortedColumn[T, U] =
+    new TypedSortedColumn[T, U](new Column(SortOrder(typedColumn.expr, Ascending)))(typedColumn.uencoder)
+
+  object defaultAscendingPoly extends Poly1 {
+    implicit def caseTypedColumn[T, U : CatalystRowOrdered] = at[TypedColumn[T, U]](c => defaultAscending(c))
+    implicit def caseTypeSortedColumn[T, U] = at[TypedSortedColumn[T, U]](identity)
+  }
+}
+
 object TypedColumn {
   /**
     * Evidence that type `T` has column `K` with type `V`.

dataset/src/main/scala/frameless/TypedDataset.scala

@@ -9,7 +9,7 @@ import org.apache.spark.sql.catalyst.plans.{Inner, LeftOuter}
 import org.apache.spark.sql._
 import shapeless._
 import shapeless.labelled.FieldType
-import shapeless.ops.hlist.{Diff, IsHCons, Prepend, ToTraversable, Tupler}
+import shapeless.ops.hlist.{Diff, IsHCons, Mapper, Prepend, ToTraversable, Tupler}
 import shapeless.ops.record.{Keys, Remover, Values}
 
 /** [[TypedDataset]] is a safer interface for working with `Dataset`.
@@ -605,6 +605,44 @@ class TypedDataset[T] protected[frameless](val dataset: Dataset[T])(implicit val
       }
   }
 
+  /** Sort each partition in the dataset by the given column expressions
+    * {{{
+    *   d.sortWithinPartitions(d('a).asc, d('b).desc)
+    * }}}
+    */
+  object sortWithinPartitions extends ProductArgs {
+    def applyProduct[U <: HList, O <: HList](columns: U)
+      (implicit
+        i0: Mapper.Aux[TypedSortedColumn.defaultAscendingPoly.type, U, O],
+        i1: ToTraversable.Aux[O, List, TypedSortedColumn[T, _]]
+      ): TypedDataset[T] = {
+      val sorted = dataset.toDF()
+        .sortWithinPartitions(i0(columns).toList[TypedSortedColumn[T, _]].map(c => new Column(c.expr)):_*)
+        .as[T](TypedExpressionEncoder[T])
+
+      TypedDataset.create[T](sorted)
+    }
+  }
+
+  /** Sort the dataset by the given column expressions
+    * {{{
+    *   d.sort(d('a).asc, d('b).desc)
+    * }}}
+    */
+  object sort extends ProductArgs {
+    def applyProduct[U <: HList, O <: HList](columns: U)
+      (implicit
+        i0: Mapper.Aux[TypedSortedColumn.defaultAscendingPoly.type, U, O],
+        i1: ToTraversable.Aux[O, List, TypedSortedColumn[T, _]]
+      ): TypedDataset[T] = {
+      val sorted = dataset.toDF()
+        .sort(i0(columns).toList[TypedSortedColumn[T, _]].map(c => new Column(c.expr)):_*)
+        .as[T](TypedExpressionEncoder[T])
+
+      TypedDataset.create[T](sorted)
+    }
+  }
+
   /** Returns a new Dataset as a tuple with the specified
     * column dropped.
     * Does not allow for dropping from a single column TypedDataset

dataset/src/main/scala/frameless/TypedWindow.scala

@@ -0,0 +1,87 @@
+package frameless
+
+import org.apache.spark.sql.Column
+import org.apache.spark.sql.catalyst.expressions.{ UnspecifiedFrame, WindowFrame }
+import org.apache.spark.sql.expressions.{ Window, WindowSpec }
+import shapeless.ops.hlist.{ Mapper, ToTraversable }
+import shapeless.{ HList, ProductArgs }
+
+trait OrderedWindow
+trait PartitionedWindow
+
+class TypedWindow[T, A] private (
+  partitionSpec: Seq[UntypedExpression[T]],
+  orderSpec: Seq[UntypedExpression[T]],
+  frame: WindowFrame //TODO. Really a rows or range between
+) {
+
+  def untyped: WindowSpec = Window
+    .partitionBy(partitionSpec.map(e => new Column(e.expr)):_*)
+    .orderBy(orderSpec.map(e => new Column(e.expr)):_*)
+  //TODO: frame
+
+
+  /* TODO: Do we want single column versions like we do for agg for better type inference?
+  def partitionBy[U](column: TypedColumn[T, U]): TypedWindow[T, A with PartitionedWindow] =
+    new TypedWindow[T, A with PartitionedWindow](
+      partitionSpec = Seq(column),
+      orderSpec = orderSpec,
+      frame = frame
+    )
+
+  def orderBy[U](column: TypedSortedColumn[T, U]): TypedWindow[T, A with OrderedWindow] =
+    new TypedWindow[T, A with OrderedWindow](
+      partitionSpec = partitionSpec,
+      orderSpec = Seq(column),
+      frame = frame
+    )
+    */
+
+  object partitionBy extends ProductArgs {
+    def applyProduct[U <: HList](columns: U)
+      (implicit
+        i1: ToTraversable.Aux[U, List, TypedColumn[T, _]]
+      ): TypedWindow[T, A with PartitionedWindow] = {
+      new TypedWindow[T, A with PartitionedWindow](
+        partitionSpec = columns.toList[TypedColumn[T, _]],
+        orderSpec = orderSpec,
+        frame = frame
+      )
+    }
+  }
+
+  object orderBy extends ProductArgs {
+    def applyProduct[U <: HList, O <: HList](columns: U)
+      (implicit
+        i0: Mapper.Aux[TypedSortedColumn.defaultAscendingPoly.type, U, O],
+        i1: ToTraversable.Aux[O, List, TypedSortedColumn[T, _]]
+      ): TypedWindow[T, A with OrderedWindow] = {
+      new TypedWindow[T, A with OrderedWindow](
+        partitionSpec = partitionSpec,
+        orderSpec = i0(columns).toList[TypedSortedColumn[T, _]],
+        frame = frame
+      )
+    }
+  }
+}
+
+object TypedWindow {
+
+  //TODO: Multiple columns.
+  def partitionBy[T](column: TypedColumn[T, _]): TypedWindow[T, PartitionedWindow] = {
+    new TypedWindow[T, PartitionedWindow](
+      partitionSpec = Seq(column),
+      orderSpec = Seq.empty,
+      frame = UnspecifiedFrame
+    )
+  }
+
+  def orderBy[T](column: TypedSortedColumn[T, _]): TypedWindow[T, OrderedWindow] = {
+    new TypedWindow[T, OrderedWindow](
+      partitionSpec = Seq.empty,
+      orderSpec = Seq(column),
+      frame = UnspecifiedFrame
+    )
+  }
+}
+

dataset/src/main/scala/frameless/functions/WindowFunctions.scala

@@ -0,0 +1,26 @@
+package frameless.functions
+
+import frameless.{ OrderedWindow, TypedAggregate, TypedColumn, TypedWindow }
+import org.apache.spark.sql.{ functions => untyped }
+
+trait WindowFunctions {
+  import WindowFunctionsHelpers.dense_rankObj
+
+  def dense_rank() = dense_rankObj
+
+}
+
+//TODO: Move these to the other funcs?
+object WindowFunctions extends WindowFunctions
+
+object WindowFunctionsHelpers {
+  //TODO: Hide this obj so that it doesn't show to users
+  private[functions] object dense_rankObj {
+    //TODO: TypedAggregate version that can be used in `agg`
+    // whose specs are all either aggs or in the groupBy. Not sure how to do the latter one
+    def over[T, A <: OrderedWindow](window: TypedWindow[T, A]): TypedColumn[T, Int] = {
+      new TypedColumn[T, Int](untyped.dense_rank().over(window.untyped))
+    }
+  }
+
+}

dataset/src/test/scala/frameless/SortTests.scala

@@ -0,0 +1,40 @@
+package frameless
+
+import org.apache.spark.sql.{ functions => sfunc }
+import shapeless.test.illTyped
+
+object SortTests {
+  case class Wack(w: Int)
+  case class Foo(a: String, b: Option[Int], c: Array[String], d: Map[String, Int], wack: Wack)
+}
+
+//TODO:
+class SortTests extends TypedDatasetSuite {
+  import SortTests._
+
+  test("sorting") {
+    val seq = Seq(
+      Foo("a", Some(2), Array("a", "b"), Map("world" -> 2), Wack(1)),
+      Foo("b", Some(1), Array("b", "a"), Map("world" -> 2), Wack(2))
+    )
+
+    val ds = TypedDataset.create(seq)
+
+    assert(ds.sort(ds('a).asc).collect().run().map(_.a) === ds.dataset.sort(sfunc.col("a").asc).collect().map(_.a))
+    assert(ds.sort(ds('a).desc).collect().run().map(_.a) === ds.dataset.sort(sfunc.col("a").desc).collect().map(_.a))
+
+    assert(ds.sort(ds('b).asc).collect().run().map(_.a) === ds.dataset.sort(sfunc.col("b").asc).collect().map(_.a))
+    assert(ds.sort(ds('b).desc).collect().run().map(_.a) === ds.dataset.sort(sfunc.col("b").desc).collect().map(_.a))
+
+    assert(ds.sort(ds('b).ascNullsFirst).collect().run().map(_.a) === ds.dataset.sort(sfunc.col("b").asc_nulls_first).collect().map(_.a))
+    assert(ds.sort(ds('a), ds('b).desc).collect().run().map(_.a) === ds.dataset.sort(sfunc.col("a"), sfunc.col("b").desc).collect().map(_.a))
+
+    illTyped {
+      //Maps aren't allow
+      """ds.sort(ds('d).desc)"""
+    }
+
+    assert(ds.sort(ds('wack).desc).collect().run().map(_.a) === ds.dataset.sort(sfunc.col("wack").desc).collect().map(_.a))
+  }
+
+}