Support optional Value class as literal

typelevel · Mar 30, 2022 · 025dd9b · 025dd9b
1 parent 29d8d15
commit 025dd9b
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 0 deletions.
diff --git a/dataset/src/main/scala/frameless/functions/package.scala b/dataset/src/main/scala/frameless/functions/package.scala
@@ -88,4 +88,47 @@ package object functions extends Udf with UnaryFunctions {
       )
     )
   }
+
+  /** Creates a [[frameless.TypedColumn]] of literal value 
+    * for an optional Value class `A`.
+    * 
+    * @tparam A the value class
+    * @tparam T the row type
+    */
+  def litValue[A : IsValueClass, T, G <: ::[_, HNil], H <: ::[_ <: FieldType[_ <: Symbol, _], HNil], K <: Symbol, V, KS <: ::[_ <: Symbol, HNil], VS <: HList](value: Option[A])(
+    implicit
+      i0: LabelledGeneric.Aux[A, G],
+      i1: DropUnitValues.Aux[G, H],
+      i2: IsHCons.Aux[H, _ <: FieldType[K, V], HNil],
+      i3: Keys.Aux[H, KS],
+      i4: Values.Aux[H, VS],
+      i5: IsHCons.Aux[KS, K, HNil],
+      i6: IsHCons.Aux[VS, V, HNil],
+      i7: TypedEncoder[V],
+      i8: ClassTag[A]
+  ): TypedColumn[T, Option[A]] = {
+    val expr = value match {
+      case Some(some) => {
+        val field: H = i1(i0.to(some))
+        val v: V = i6.head(i4(field))
+
+        new Literal(v, i7.jvmRepr)
+      }
+
+      case _ =>
+        Literal.create(null, i7.jvmRepr)
+    }
+
+    implicit val enc: TypedEncoder[A] =
+      RecordFieldEncoder.valueClass[A, G, H, K, V, KS].encoder
+
+    new TypedColumn[T, Option[A]](
+      Lit(
+        dataType = i7.catalystRepr,
+        nullable = true,
+        toCatalyst = i7.toCatalyst(expr).genCode(_),
+        show = () => value.toString
+      )
+    )
+  }
 }
diff --git a/dataset/src/test/scala/frameless/LitTests.scala b/dataset/src/test/scala/frameless/LitTests.scala
@@ -69,6 +69,23 @@ class LitTests extends TypedDatasetSuite with Matchers {
       collect.run() shouldBe initial.map(_.copy(name = lorem))
   }
 
+  test("support optional value class") {
+    val initial = Seq(
+      R(name = "Foo", id = 1, alias = None),
+      R(name = "Bar", id = 2, alias = Some(new Name("Lorem"))))
+    val ds = TypedDataset.create(initial)
+
+    ds.collect.run() shouldBe initial
+
+    val someIpsum: Option[Name] = Some(new Name("Ipsum"))
+
+    ds.withColumnReplaced('alias, functions.litValue(someIpsum)).
+      collect.run() shouldBe initial.map(_.copy(alias = someIpsum))
+
+    ds.withColumnReplaced('alias, functions.litValue(Option.empty[Name])).
+      collect.run() shouldBe initial.map(_.copy(alias = None))
+  }
+
   test("#205: comparing literals encoded using Injection") {
     import org.apache.spark.sql.catalyst.util.DateTimeUtils
     implicit val dateAsInt: Injection[java.sql.Date, Int] =
@@ -85,3 +102,5 @@ class LitTests extends TypedDatasetSuite with Matchers {
 final case class P(i: Int, d: java.sql.Date)
 
 final case class Q(id: Int, name: Name)
+
+final case class R(id: Int, name: String, alias: Option[Name])