Make Literals foldable, ensure Parquet predicates pushdown

build.sbt

@@ -74,7 +74,7 @@ lazy val dataset = project
 lazy val `dataset-spark33` = project
   .settings(name := "frameless-dataset-spark33")
   .settings(sourceDirectory := (dataset / sourceDirectory).value)
-  .settings(Compile / unmanagedSourceDirectories += (dataset / baseDirectory).value / "src" / "main" / "spark-3")
+  .settings(Compile / unmanagedSourceDirectories += (dataset / baseDirectory).value / "src" / "main" / "spark-33")
   .settings(datasetSettings)
   .settings(sparkDependencies(spark33Version))
   .settings(spark33Settings)
@@ -83,7 +83,7 @@ lazy val `dataset-spark33` = project
 lazy val `dataset-spark32` = project
   .settings(name := "frameless-dataset-spark32")
   .settings(sourceDirectory := (dataset / sourceDirectory).value)
-  .settings(Compile / unmanagedSourceDirectories += (dataset / baseDirectory).value / "src" / "main" / "spark-3")
+  .settings(Compile / unmanagedSourceDirectories += (dataset / baseDirectory).value / "src" / "main" / "spark-32")
   .settings(datasetSettings)
   .settings(sparkDependencies(spark32Version))
   .settings(spark32Settings)

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

@@ -1,5 +1,6 @@
 package frameless.functions
 
+import org.apache.spark.sql.FramelessInternals
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.expressions.codegen._
 import org.apache.spark.sql.catalyst.expressions.{Expression, NonSQLExpression}
@@ -59,7 +60,8 @@ private[frameless] case class Lit[T <: AnyVal](
 
   protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]): Expression = this
 
-  // see https://github.com/typelevel/frameless/pull/721#issuecomment-1581137730 (InvokeLike <3.3.1 SPARK-40380)
+  // see https://github.com/typelevel/frameless/pull/721#issuecomment-1581137730 (InvokeLike <3.3.1 SPARK-40380, ConditionalExpression SPARK-39106)
   // for why this does not push down on 3.2, 3.3.1 and higher _do_ pushdown
-  override val foldable: Boolean = catalystExpr.foldable
+  // TODO remove the compat layer once 3.2 is no longer supported
+  override val foldable: Boolean = FramelessInternals.foldableCompat(catalystExpr)
 }

dataset/src/main/scala/org/apache/spark/sql/FramelessInternals.scala

@@ -1,5 +1,6 @@
 package org.apache.spark.sql
 
+import frameless.FoldableImpl
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.codegen._
 import org.apache.spark.sql.catalyst.expressions.{Alias, CreateStruct}
@@ -10,6 +11,7 @@ import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, Project}
 import org.apache.spark.sql.execution.QueryExecution
 import org.apache.spark.sql.types._
 import org.apache.spark.sql.types.ObjectType
+
 import scala.reflect.ClassTag
 
 object FramelessInternals {
@@ -70,4 +72,11 @@ object FramelessInternals {
     protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = tagged.genCode(ctx)
     protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]): Expression = copy(newChildren.head)
   }
+
+  /**
+   * With 3.2 it's a backport of SPARK-40380, higher versions are simply expression.foldable
+   * @param expression
+   * @return
+   */
+  def foldableCompat(expression: Expression): Boolean = FoldableImpl.foldableCompat(expression)
 }

dataset/src/main/spark-3.4+/frameless/FoldableImpl.scala

@@ -0,0 +1,7 @@
+package frameless
+
+import org.apache.spark.sql.catalyst.expressions.Expression
+
+object FoldableImpl {
+  def foldableCompat(expression: Expression): Boolean = expression.foldable
+}

dataset/src/main/spark-32/frameless/FoldableImpl.scala

@@ -0,0 +1,54 @@
+package frameless
+
+import org.apache.spark.sql.FramelessSpark32Internals
+import org.apache.spark.sql.catalyst.expressions.codegen.{CodegenContext, ExprCode}
+import org.apache.spark.sql.catalyst.expressions.{CaseWhen, Coalesce, Expression, If, NaNvl, UnaryExpression}
+import org.apache.spark.sql.catalyst.expressions.objects.InvokeLike
+import org.apache.spark.sql.types.{DataType, ObjectType}
+
+object FoldableImpl {
+
+  trait ExpressionProxy {
+
+    def child: Expression
+
+    protected def withNewChildInternal(newChild: Expression): Expression = ???
+
+    protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = ???
+
+    def dataType: DataType = child.dataType
+
+  }
+
+  // backported from SPARK-40380
+  case class InvokeLikeImpl(child: InvokeLike) extends UnaryExpression with ExpressionProxy {
+    // Returns true if we can trust all values of the given DataType can be serialized.
+    def trustedSerializable(dt: DataType): Boolean = {
+      // Right now we conservatively block all ObjectType (Java objects) regardless of
+      // serializability, because the type-level info with java.io.Serializable and
+      // java.io.Externalizable marker interfaces are not strong guarantees.
+      // This restriction can be relaxed in the future to expose more optimizations.
+      !FramelessSpark32Internals.existsRecursively(dt)(_.isInstanceOf[ObjectType])
+    }
+
+    override def foldable =
+      child.children.forall(_.foldable) && deterministic && trustedSerializable(dataType)
+  }
+
+  // foldable not implemented in 3.2, is in 3.3 (SPARK-39106)
+  case class ConditionalExpressionImpl(child: Expression) extends UnaryExpression with ExpressionProxy {
+    override def foldable =
+      child.children.forall(_.foldable)
+  }
+
+  // needed as we cannot test foldable on any parent expression if they have Invoke
+  // but similarly we cannot assume the parent is foldable - so we replace InvokeLike
+  def replaced(expression: Expression): Expression = expression transformUp {
+    case il: InvokeLike => InvokeLikeImpl(il)
+    case e@( _: If | _: CaseWhen | _: Coalesce | _: NaNvl ) =>
+      ConditionalExpressionImpl(e)
+  }
+
+  def foldableCompat(expression: Expression): Boolean =
+    replaced(expression).foldable
+}

dataset/src/main/spark-32/org/apache/spark/sql/FramelessSpark32Internals.scala

@@ -0,0 +1,11 @@
+package org.apache.spark.sql
+
+import org.apache.spark.sql.types.DataType
+
+object FramelessSpark32Internals {
+
+  /**
+   * Returns true if any `DataType` of this DataType tree satisfies the given function `f`.
+   */
+  def existsRecursively(dt: DataType)(f: (DataType) => Boolean): Boolean = dt.existsRecursively(f)
+}

dataset/src/main/spark-33/frameless/FoldableImpl.scala

@@ -0,0 +1,7 @@
+package frameless
+
+import org.apache.spark.sql.catalyst.expressions.Expression
+
+object FoldableImpl {
+  def foldableCompat(expression: Expression): Boolean = expression.foldable
+}

dataset/src/main/spark-33/frameless/MapGroups.scala

@@ -0,0 +1,14 @@
+package frameless
+
+import org.apache.spark.sql.Encoder
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, MapGroups => SMapGroups}
+
+object MapGroups {
+  def apply[K: Encoder, T: Encoder, U: Encoder](
+    func: (K, Iterator[T]) => TraversableOnce[U],
+    groupingAttributes: Seq[Attribute],
+    dataAttributes: Seq[Attribute],
+    child: LogicalPlan
+  ): LogicalPlan = SMapGroups(func, groupingAttributes, dataAttributes, child)
+}

dataset/src/test/scala/frameless/sql/rules/FramelessLitPushDownTests.scala

@@ -3,60 +3,51 @@ package frameless.sql.rules
 import frameless._
 import frameless.functions.Lit
 import org.apache.spark.sql.catalyst.util.DateTimeUtils.{currentTimestamp, microsToInstant}
-import org.apache.spark.sql.sources.{GreaterThanOrEqual, IsNotNull, EqualTo}
+import org.apache.spark.sql.sources.{EqualTo, GreaterThanOrEqual, IsNotNull}
 import org.apache.spark.sql.catalyst.expressions
 import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema
-
 import java.time.Instant
 
 class FramelessLitPushDownTests extends SQLRulesSuite {
   private val now: Long = currentTimestamp()
 
-  val invokeShouldFoldOnHigherThan3_2 = not3_2 _
-
-  test("java.sql.Timestamp push-down") (
-    invokeShouldFoldOnHigherThan3_2 {
-      val expected = java.sql.Timestamp.from(microsToInstant(now))
-      val expectedStructure = X1(SQLTimestamp(now))
-      val expectedPushDownFilters = List(IsNotNull("a"), GreaterThanOrEqual("a", expected))
-
-      predicatePushDownTest[SQLTimestamp](
-        expectedStructure,
-        expectedPushDownFilters,
-        { case e @ expressions.GreaterThanOrEqual(_, _: Lit[_]) => e },
-        _ >= expectedStructure.a
-      )
-    }
-  )
-
-  test("java.time.Instant push-down") (
-    invokeShouldFoldOnHigherThan3_2 {
-      val expected = java.sql.Timestamp.from(microsToInstant(now))
-      val expectedStructure = X1(microsToInstant(now))
-      val expectedPushDownFilters = List(IsNotNull("a"), GreaterThanOrEqual("a", expected))
-
-      predicatePushDownTest[Instant](
-        expectedStructure,
-        expectedPushDownFilters,
-        { case e @ expressions.GreaterThanOrEqual(_, _: Lit[_]) => e },
-        _ >= expectedStructure.a
-      )
-    }
-  )
-
-  test("struct push-down") (
-    invokeShouldFoldOnHigherThan3_2 {
-      type Payload = X4[Int, Int, Int, Int]
-      val expectedStructure = X1(X4(1, 2, 3, 4))
-      val expected = new GenericRowWithSchema(Array(1, 2, 3, 4), TypedExpressionEncoder[Payload].schema)
-      val expectedPushDownFilters = List(IsNotNull("a"), EqualTo("a", expected))
-
-      predicatePushDownTest[Payload](
-        expectedStructure,
-        expectedPushDownFilters,
-        { case e @ expressions.EqualTo(_, _: Lit[_]) => e },
-        _ === expectedStructure.a
-      )
-    }
-  )
+  test("java.sql.Timestamp push-down") {
+    val expected = java.sql.Timestamp.from(microsToInstant(now))
+    val expectedStructure = X1(SQLTimestamp(now))
+    val expectedPushDownFilters = List(IsNotNull("a"), GreaterThanOrEqual("a", expected))
+
+    predicatePushDownTest[SQLTimestamp](
+      expectedStructure,
+      expectedPushDownFilters,
+      { case e @ expressions.GreaterThanOrEqual(_, _: Lit[_]) => e },
+      _ >= expectedStructure.a
+    )
+  }
+
+  test("java.time.Instant push-down") {
+    val expected = java.sql.Timestamp.from(microsToInstant(now))
+    val expectedStructure = X1(microsToInstant(now))
+    val expectedPushDownFilters = List(IsNotNull("a"), GreaterThanOrEqual("a", expected))
+
+    predicatePushDownTest[Instant](
+      expectedStructure,
+      expectedPushDownFilters,
+      { case e @ expressions.GreaterThanOrEqual(_, _: Lit[_]) => e },
+      _ >= expectedStructure.a
+    )
+  }
+
+  test("struct push-down") {
+    type Payload = X4[Int, Int, Int, Int]
+    val expectedStructure = X1(X4(1, 2, 3, 4))
+    val expected = new GenericRowWithSchema(Array(1, 2, 3, 4), TypedExpressionEncoder[Payload].schema)
+    val expectedPushDownFilters = List(IsNotNull("a"), EqualTo("a", expected))
+
+    predicatePushDownTest[Payload](
+      expectedStructure,
+      expectedPushDownFilters,
+      { case e @ expressions.EqualTo(_, _: Lit[_]) => e },
+      _ === expectedStructure.a
+    )
+  }
 }

dataset/src/test/scala/frameless/sql/rules/SQLRulesSuite.scala

@@ -2,7 +2,6 @@ package frameless.sql.rules
 
 import frameless._
 import frameless.sql._
-import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.expressions.Expression
 import org.apache.spark.sql.sources.Filter
 import org.apache.spark.sql.catalyst.plans.logical
@@ -17,27 +16,6 @@ trait SQLRulesSuite extends TypedDatasetSuite with Matchers { self =>
     s"$tmpDir/${self.getClass.getName}"
   }
 
-  lazy val sparkFullVersion = {
-    val pos = classOf[Expression].getPackage.getSpecificationVersion
-    if (pos eq null) // DBR is always null
-      SparkSession.active.version // taking a running spark version string, hence lazy
-    else
-      pos
-  }
-
-  lazy val sparkVersion = {
-    sparkFullVersion.split('.').take(2).mkString(".")
-  }
-
-  /**
-   * Don't run this test on 3.2
-   */
-  def not3_2[T](thunk: => T): Any =
-    if (sparkVersion != "3.2")
-      thunk
-    else
-      ()
-
   def withDataset[A: TypedEncoder: CatalystOrdered](payload: A)(f: TypedDataset[A] => Assertion): Assertion = {
     TypedDataset.create(Seq(payload)).write.mode("overwrite").parquet(path)
     f(TypedDataset.createUnsafe[A](session.read.parquet(path)))