Make Literals foldable, ensure Parquet predicates pushdown

build.sbt

@@ -10,6 +10,7 @@ val shapeless = "2.3.10"
 val scalacheck = "1.17.0"
 val scalacheckEffect = "1.0.4"
 val refinedVersion = "0.10.3"
+val nakedFSVersion = "0.1.0"
 
 val Scala212 = "2.12.17"
 val Scala213 = "2.13.10"
@@ -192,7 +193,9 @@ lazy val datasetSettings = framelessSettings ++ framelessTypedDatasetREPL ++ Seq
       dmm("org.apache.spark.sql.FramelessInternals.column")
     )
   },
-  coverageExcludedPackages := "org.apache.spark.sql.reflection"
+  coverageExcludedPackages := "org.apache.spark.sql.reflection",
+
+  libraryDependencies += "com.globalmentor" % "hadoop-bare-naked-local-fs" % nakedFSVersion % Test exclude("org.apache.hadoop", "hadoop-commons")
 )
 
 lazy val refinedSettings = framelessSettings ++ framelessTypedDatasetREPL ++ Seq(

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

@@ -8,8 +8,8 @@ import org.apache.spark.sql.types.DataType
 private[frameless] case class Lit[T <: AnyVal](
     dataType: DataType,
     nullable: Boolean,
-    toCatalyst: CodegenContext => ExprCode,
-    show: () => String
+    show: () => String,
+    convertedExpr: Expression // must be the same toCatalyst as the toCatalyst function
 ) extends Expression with NonSQLExpression {
   override def toString: String = s"FramelessLit(${show()})"
 
@@ -53,7 +53,7 @@ private[frameless] case class Lit[T <: AnyVal](
 
   def children: Seq[Expression] = Nil
 
-  protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = toCatalyst(ctx)
+  protected def doGenCode(ctx: CodegenContext, ev: ExprCode): ExprCode = convertedExpr.genCode(ctx)
 
   protected def withNewChildrenInternal(newChildren: IndexedSeq[Expression]): Expression = this
 }

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

@@ -45,8 +45,8 @@ package object functions extends Udf with UnaryFunctions {
         Lit(
           dataType = encoder.catalystRepr,
           nullable = encoder.nullable,
-          toCatalyst = encoder.toCatalyst(expr).genCode(_),
-          show = () => value.toString
+          show = () => value.toString,
+          encoder.toCatalyst(expr)
         )
       )
     }
@@ -84,8 +84,8 @@ package object functions extends Udf with UnaryFunctions {
       Lit(
         dataType = i7.catalystRepr,
         nullable = i7.nullable,
-        toCatalyst = i7.toCatalyst(expr).genCode(_),
-        show = () => value.toString
+        show = () => value.toString,
+        i7.toCatalyst(expr)
       )
     )
   }
@@ -127,8 +127,8 @@ package object functions extends Udf with UnaryFunctions {
       Lit(
         dataType = i7.catalystRepr,
         nullable = true,
-        toCatalyst = i7.toCatalyst(expr).genCode(_),
-        show = () => value.toString
+        show = () => value.toString,
+        i7.toCatalyst(expr)
       )
     )
   }

dataset/src/main/scala/frameless/optimiser/LiteralRule.scala

@@ -0,0 +1,17 @@
+package frameless.optimiser
+
+import frameless.functions.Lit
+import org.apache.spark.sql.catalyst.expressions.Literal
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.rules.Rule
+
+object LiteralRule extends Rule[LogicalPlan] {
+  override def apply(plan: LogicalPlan): LogicalPlan = plan.transformExpressions {
+
+    /*
+     * replace all literals to allow constant folding and push down
+     */
+    case Lit(dataType, _, _, convertedValue) =>
+      Literal(convertedValue.eval(), dataType)
+  }
+}

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

@@ -1,20 +1,33 @@
 package frameless
 
+import com.globalmentor.apache.hadoop.fs.BareLocalFileSystem
+import org.apache.hadoop.fs.local.StreamingFS
 import org.apache.spark.{SparkConf, SparkContext}
 import org.apache.spark.sql.{SQLContext, SparkSession}
 import org.scalactic.anyvals.PosZInt
 import org.scalatest.BeforeAndAfterAll
 import org.scalatestplus.scalacheck.Checkers
 import org.scalacheck.Prop
 import org.scalacheck.Prop._
+
 import scala.util.{Properties, Try}
 import org.scalatest.funsuite.AnyFunSuite
 
 trait SparkTesting { self: BeforeAndAfterAll =>
 
   val appID: String = new java.util.Date().toString + math.floor(math.random * 10E4).toLong.toString
 
-  val conf: SparkConf = new SparkConf()
+  /**
+   * Allows bare naked to be used instead of winutils for testing / dev
+   */
+  def registerFS(sparkConf: SparkConf): SparkConf =
+    if (System.getProperty("os.name").startsWith("Windows"))
+      sparkConf.set("spark.hadoop.fs.file.impl", classOf[BareLocalFileSystem].getName).
+        set("spark.hadoop.fs.AbstractFileSystem.file.impl", classOf[StreamingFS].getName)
+    else
+      sparkConf
+
+  val conf: SparkConf = registerFS(new SparkConf())
     .setMaster("local[*]")
     .setAppName("test")
     .set("spark.ui.enabled", "false")

dataset/src/test/scala/frameless/optimiser/LitTests.scala

@@ -0,0 +1,44 @@
+package frameless.optimiser
+
+import frameless._
+import org.apache.spark.sql.execution.FileSourceScanExec
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec
+
+class LitTests extends TypedDatasetSuite {
+  test("sqlTimestamp pushdown") {
+    val ms = SQLTimestamp(System.currentTimeMillis())
+    TypedDataset.create(Seq(X1(ms))).write.mode("overwrite").parquet("./target/testData")
+    val dataset = TypedDataset.createUnsafe[X1[SQLTimestamp]](session.read.parquet("./target/testData"))
+
+    {
+      val pushDowns = getPushDowns(dataset.filter(dataset('a) >= ms))
+      assert(!pushDowns.exists(_.contains("GreaterThanOrEqual")), "Should NOT have had a push down predicate for the underlying long as no optimiser is yet used")
+    }
+
+    val orig = session.sqlContext.experimental.extraOptimizations
+    try {
+      session.sqlContext.experimental.extraOptimizations ++= Seq(LiteralRule)
+
+      val pushDowns = getPushDowns(dataset.filter(dataset('a) >= ms))
+      assert(pushDowns.exists(_.contains("GreaterThanOrEqual")), "Should have had a push down predicate for the underlying long")
+    } finally {
+      session.sqlContext.experimental.extraOptimizations = orig
+    }
+  }
+
+  def getPushDowns(dataset: TypedDataset[_]): Seq[String] = {
+    val sparkPlan = dataset.queryExecution.executedPlan
+
+    (if (sparkPlan.children.isEmpty)
+    // assume it's AQE
+      sparkPlan match {
+        case aq: AdaptiveSparkPlanExec => aq.initialPlan
+        case _ => sparkPlan
+      }
+    else
+      sparkPlan).collect {
+      case fs: FileSourceScanExec =>
+        fs.metadata.collect { case ("PushedFilters", value) if value != "[]" => value }
+    }.flatten
+  }
+}

dataset/src/test/scala/org/apache/hadoop/fs/local/StreamingFS.scala

@@ -0,0 +1,7 @@
+package org.apache.hadoop.fs.local
+
+import com.globalmentor.apache.hadoop.fs.BareLocalFileSystem
+import org.apache.hadoop.fs.DelegateToFileSystem
+
+class StreamingFS(uri: java.net.URI, conf: org.apache.hadoop.conf.Configuration) extends
+  DelegateToFileSystem(uri, new BareLocalFileSystem(), conf, "file", false) {}