Field TypedEncoder

Author: cchantep

dataset/src/main/scala/frameless/RecordEncoder.scala

@@ -187,7 +187,18 @@ final class RecordFieldEncoder[T](
   private[frameless] val jvmRepr: DataType,
   private[frameless] val fromCatalyst: Expression => Expression,
   private[frameless] val toCatalyst: Expression => Expression
-) extends Serializable
+) extends Serializable { self =>
+  private[frameless] def toTypedEncoder = new TypedEncoder[T]()(encoder.classTag) {
+    def nullable: Boolean = encoder.nullable
+
+    def jvmRepr: DataType = self.jvmRepr
+    def catalystRepr: DataType = encoder.catalystRepr
+
+    def fromCatalyst(path: Expression): Expression = self.fromCatalyst(path)
+
+    def toCatalyst(path: Expression): Expression = self.toCatalyst(path)
+  }
+}
 
 object RecordFieldEncoder extends RecordFieldEncoderLowPriority {
 

dataset/src/main/scala/frameless/TypedEncoder.scala

@@ -756,5 +756,14 @@ object TypedEncoder {
     }
   }
 
+  /**
+   * In case a type `T` encoding is supported in derivation (as a struct field),
+   * then this allows to resolve the corresponding `TypedEncoder[T]`,
+   * so the field can be handled invidiually.
+   */
+  def usingFieldEncoder[T](
+    implicit fieldEncoder: shapeless.Lazy[RecordFieldEncoder[T]]):
+      TypedEncoder[T] = fieldEncoder.value.toTypedEncoder
+
   object injections extends InjectionEnum
 }

dataset/src/test/scala/frameless/ColumnTests.scala

@@ -615,4 +615,31 @@ final class ColumnTests extends TypedDatasetSuite with Matchers {
     // we should be able to block the following as well...
     "ds.col(_.a.toInt)" shouldNot typeCheck
   }
+
+  test("col through record encoder (for Value class)") {
+    import RecordEncoderTests.{ Name, Person }
+
+    val bar = new Name("bar")
+    val foo = new Name("foo")
+
+    val ds: TypedDataset[Person] =
+      TypedDataset.create(Seq(Person(bar, 23), Person(foo, 11)))
+
+    a[org.apache.spark.sql.AnalysisException] should be thrownBy {
+      // TypedEncoder[Name] is resolved using case class derivation,
+      // which is not compatible to the way such Value class
+      // is encoded as a field in another class,
+      // which leads to encoding/analysis error.
+
+      ds.select(ds.col[Name](Symbol("name"))).
+        collect.run().toSeq shouldEqual Seq[Name](bar, foo)
+    }
+
+    {
+      implicit def enc: TypedEncoder[Name] = TypedEncoder.usingFieldEncoder[Name]
+
+      ds.select(ds.col[Name](Symbol("name"))).
+        collect.run().toSeq shouldEqual Seq[Name](bar, foo)
+    }
+  }
 }

dataset/src/test/scala/frameless/CreateTests.scala

@@ -130,7 +130,7 @@ class CreateTests extends TypedDatasetSuite with Matchers {
     val v = Vector(X2(1,2))
     val df = TypedDataset.create(v).dataset.toDF()
 
-    a [IllegalStateException] should be thrownBy {
+    a[IllegalStateException] should be thrownBy {
       TypedDataset.createUnsafe[X1[Int]](df).show().run()
     }
   }

dataset/src/test/scala/frameless/TypedDatasetSuite.scala

@@ -79,13 +79,17 @@ class TypedDatasetSuite extends AnyFunSuite with Checkers with BeforeAndAfterAll
     val da = numeric.toDouble(a)
     val db = numeric.toDouble(b)
     val epsilon = 1E-6
+
     // Spark has a weird behaviour concerning expressions that should return Inf
     // Most of the time they return NaN instead, for instance stddev of Seq(-7.827553978923477E227, -5.009124275715786E153)
-    if((da.isNaN || da.isInfinity) && (db.isNaN || db.isInfinity)) proved
-    else if (
+    if ((da.isNaN || da.isInfinity) && (db.isNaN || db.isInfinity)) {
+      proved
+    } else if (
       (da - db).abs < epsilon ||
-      (da - db).abs < da.abs / 100)
+      (da - db).abs < da.abs / 100) {
         proved
-    else falsified :| s"Expected $a but got $b, which is more than 1% off and greater than epsilon = $epsilon."
+    } else {
+      falsified :| s"Expected $a but got $b, which is more than 1% off and greater than epsilon = $epsilon."
+    }
   }
 }