Add a typed col function for creating column references

Resolves #186.
typelevel · Sep 22, 2017 · 897e499 · 897e499
1 parent 68aa838
commit 897e499
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/dataset/src/main/scala/frameless/functions/package.scala b/dataset/src/main/scala/frameless/functions/package.scala
@@ -2,6 +2,8 @@ package frameless
 
 import org.apache.spark.sql.catalyst.ScalaReflection
 import org.apache.spark.sql.catalyst.expressions.Literal
+import org.apache.spark.sql.functions.{ col => sparkCol }
+import shapeless.Witness
 
 package object functions extends Udf with UnaryFunctions {
   object aggregate extends AggregateFunctions
@@ -17,4 +19,12 @@ package object functions extends Udf with UnaryFunctions {
       new TypedColumn(expr)
     }
   }
+
+  def col[T, A](column: Witness.Lt[Symbol])(
+    implicit
+    exists: TypedColumn.Exists[T, column.T, A],
+    encoder: TypedEncoder[A]): TypedColumn[T, A] = {
+    val untypedExpr = sparkCol(column.value.name).as[A](TypedExpressionEncoder[A])
+    new TypedColumn[T, A](untypedExpr)
+  }
 }
diff --git a/dataset/src/test/scala/frameless/SelectTests.scala b/dataset/src/test/scala/frameless/SelectTests.scala
@@ -18,9 +18,10 @@ class SelectTests extends TypedDatasetSuite {
       val A = dataset.col[A]('a)
 
       val dataset2 = dataset.select(A).collect().run().toVector
+      val symDataset2 = dataset.select(functions.col('a)).collect().run().toVector
       val data2 = data.map { case X4(a, _, _, _) => a }
 
-      dataset2 ?= data2
+      (dataset2 ?= data2) && (symDataset2 ?= data2)
     }
 
     check(forAll(prop[Int, Int, Int, Int] _))