[spark] Implement SupportsRuntimeV2Filtering

fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussScan.scala

@@ -18,19 +18,25 @@
 package org.apache.fluss.spark.read
 
 import org.apache.fluss.config.Configuration
-import org.apache.fluss.metadata.{TableInfo, TablePath}
+import org.apache.fluss.metadata.{TableBucket, TableInfo, TablePath}
 import org.apache.fluss.predicate.{Predicate => FlussPredicate}
 import org.apache.fluss.spark.SparkConversions
+import org.apache.fluss.spark.read.lake.{FlussLakeInputPartition, FlussLakeUpsertInputPartition}
+import org.apache.fluss.spark.utils.SparkPartitionPredicate
 
+import org.apache.spark.sql.connector.expressions.{Expressions, NamedReference}
 import org.apache.spark.sql.connector.expressions.filter.Predicate
 import org.apache.spark.sql.connector.metric.CustomMetric
-import org.apache.spark.sql.connector.read.{Batch, Scan}
+import org.apache.spark.sql.connector.read.{Batch, InputPartition, PartitionReaderFactory, Scan, SupportsRuntimeV2Filtering}
 import org.apache.spark.sql.connector.read.streaming.MicroBatchStream
 import org.apache.spark.sql.types.StructType
 import org.apache.spark.sql.util.CaseInsensitiveStringMap
 
+import scala.collection.JavaConverters._
+import scala.util.control.NonFatal
+
 /** An interface that extends from Spark [[Scan]]. */
-trait FlussScan extends Scan {
+trait FlussScan extends Scan with SupportsRuntimeV2Filtering {
   def tableInfo: TableInfo
 
   def tablePath: TablePath
@@ -46,8 +52,106 @@ trait FlussScan extends Scan {
 
   def timeRange: Option[FlussTimeRange] = None
 
+  /** The planner used to capture partition metadata together with the initial splits. */
+  def planner: SplitPlanner
+
+  /** Creates the batch whose original splits and reader factory are reused for runtime pruning. */
+  protected def createBatch(): Batch
+
   protected def scanType: String
 
+  private var runtimePredicates: Seq[Predicate] = Seq.empty
+  private var filteredPartitions: Option[Array[InputPartition]] = None
+  private var initialPartitionValues: Map[Long, Seq[String]] = Map.empty
+
+  private lazy val originalBatch: Batch = createBatch()
+
+  // Spark calls toBatch again after runtime filtering. Never re-plan offsets or snapshots.
+  private lazy val initialPartitions: Array[InputPartition] = {
+    val partitions = originalBatch.planInputPartitions()
+    initialPartitionValues = planner.partitionValuesById
+    partitions
+  }
+
+  private lazy val batch: Batch = new Batch {
+    private lazy val readerFactory = originalBatch.createReaderFactory()
+
+    override def planInputPartitions(): Array[InputPartition] =
+      filteredPartitions.getOrElse(initialPartitions).clone()
+
+    override def createReaderFactory(): PartitionReaderFactory = readerFactory
+  }
+
+  override def toBatch: Batch = batch
+
+  override def filterAttributes(): Array[NamedReference] =
+    tableInfo.getPartitionKeys.asScala
+      .filter(key => readSchema().fields.count(_.name == key) == 1)
+      .map(key => Expressions.column(s"`${key.replace("`", "``")}`"))
+      .toArray
+
+  override def filter(predicates: Array[Predicate]): Unit = {
+    if (predicates.isEmpty || !tableInfo.isPartitioned) {
+      return
+    }
+    val (unsupported, partitionFilter) =
+      SparkPartitionPredicate.extract(tableInfo, predicates.toSeq)
+    // The shared converter deliberately rejects empty IN. Runtime IN on a partition key,
+    // however, means the join produced no keys, so even splits with unknown metadata can go.
+    val emptyIn = predicates.filter {
+      p =>
+        p.name() == "IN" && (p.children() match {
+          case Array(ref: NamedReference) =>
+            ref.fieldNames().length == 1 &&
+            tableInfo.getPartitionKeys.contains(ref.fieldNames()(0))
+          case _ => false
+        })
+    }
+    if (partitionFilter.isEmpty && emptyIn.isEmpty) {
+      return
+    }
+    val current = filteredPartitions.getOrElse(initialPartitions)
+    filteredPartitions = Some(
+      if (emptyIn.nonEmpty) Array.empty[InputPartition]
+      else current.filter(matchesPartition(_, partitionFilter)))
+    runtimePredicates ++= predicates.filterNot(unsupported.contains) ++ emptyIn
+  }
+
+  private def matchesPartition(
+      partition: InputPartition,
+      predicate: Option[FlussPredicate]): Boolean = {
+    def matches(values: Seq[String]): Boolean = {
+      if (values.size != tableInfo.getPartitionKeys.size()) {
+        true
+      } else {
+        try {
+          SparkPartitionPredicate.matchesPartition(tableInfo, values, predicate)
+        } catch {
+          case NonFatal(_) => true
+        }
+      }
+    }
+
+    def matchesBucket(bucket: TableBucket): Boolean =
+      Option(bucket.getPartitionId)
+        .flatMap(id => initialPartitionValues.get(id.longValue()))
+        .forall(matches)
+
+    partition match {
+      case p: FlussAppendInputPartition => matchesBucket(p.tableBucket)
+      case p: FlussUpsertInputPartition => matchesBucket(p.tableBucket)
+      case p: FlussLakeInputPartition =>
+        matches(Option(p.lakeSplit.partition()).map(_.asScala.toSeq).getOrElse(Seq.empty))
+      case p: FlussLakeUpsertInputPartition =>
+        val lakeMatches = Option(p.lakeSplits).exists(_.asScala.exists {
+          split => matches(Option(split.partition()).map(_.asScala.toSeq).getOrElse(Seq.empty))
+        })
+        val hasLogTail = p.logStartingOffset < p.logStoppingOffset
+        lakeMatches || (hasLogTail && matchesBucket(p.tableBucket))
+      case _ => true
+    }
+  }
+
   override def readSchema(): StructType = {
     requiredSchema.getOrElse(SparkConversions.toSparkDataType(tableInfo.getRowType))
   }
@@ -67,9 +171,12 @@ trait FlussScan extends Scan {
       case Some(r) => s"$withPartition [TimeRange: [${r.startMs}, ${r.endMs})]"
       case None => withPartition
     }
+    val withRuntime =
+      if (runtimePredicates.isEmpty) withTimeRange
+      else s"$withTimeRange [RuntimePredicates: ${runtimePredicates.mkString("[", ", ", "]")}]"
     limit match {
-      case Some(l) => s"$withTimeRange [Limit: $l]"
-      case None => withTimeRange
+      case Some(l) => s"$withRuntime [Limit: $l]"
+      case None => withRuntime
     }
   }
 
@@ -100,7 +207,7 @@ case class FlussAppendScan(
 
   override def timeRange: Option[FlussTimeRange] = planner.timeRange
 
-  override def toBatch: Batch = {
+  override protected def createBatch(): Batch = {
     new FlussAppendBatch(
       tablePath,
       tableInfo,
@@ -146,7 +253,7 @@ case class FlussUpsertScan(
 
   override def timeRange: Option[FlussTimeRange] = planner.timeRange
 
-  override def toBatch: Batch = {
+  override protected def createBatch(): Batch = {
     new FlussUpsertBatch(
       tablePath,
       tableInfo,

fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/SplitPlanner.scala

@@ -68,6 +68,9 @@ sealed trait SplitPlanner extends AutoCloseable {
 
   def plan(): Array[InputPartition]
 
+  /** Partition values captured by the initial plan; accessing this map performs no RPCs. */
+  def partitionValuesById: Map[Long, Seq[String]]
+
   /**
    * Whether this planner will union its Fluss plan with a lake snapshot. Backed by the memoized
    * probe, so the answer is deterministic and stable across `plan()` invocations and repeated reads
@@ -135,6 +138,10 @@ abstract class AbstractSplitPlanner(
     admin0
   }
 
+  private var plannedPartitionValues: Map[Long, Seq[String]] = Map.empty
+
+  override def partitionValuesById: Map[Long, Seq[String]] = plannedPartitionValues
+
   protected lazy val partitionInfos: util.List[PartitionInfo] = {
     val infos = admin.listPartitionInfos(tablePath).get()
     // Fail fast if any partition's bucket count differs from the table-level count.
@@ -149,6 +156,10 @@ abstract class AbstractSplitPlanner(
               s"${info.getBucketCount} but the table-level count is $tableBucketCount.")
         }
     }
+    plannedPartitionValues = infos.asScala.map {
+      info =>
+        info.getPartitionId -> info.getResolvedPartitionSpec.getPartitionValues.asScala.toVector
+    }.toMap
     infos
   }
 

fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/SparkRuntimeFilteringTest.scala

@@ -0,0 +1,125 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.fluss.spark
+
+import org.apache.fluss.spark.read.FlussScan
+
+import org.apache.spark.sql.Row
+import org.apache.spark.sql.connector.expressions.{Expression, Expressions}
+import org.apache.spark.sql.connector.expressions.filter.Predicate
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec
+import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceV2ScanRelation}
+import org.assertj.core.api.Assertions.assertThat
+
+import scala.collection.JavaConverters._
+
+/** Verifies join-driven partition pruning and fixed read boundaries against a Fluss cluster. */
+class SparkRuntimeFilteringTest extends FlussSparkTestBase {
+
+  for (
+    primaryKey <- Seq(false, true); adaptive <- Seq(false, true); incremental <- Seq(false, true)
+  ) {
+    test(
+      s"partition join applies DPP: primaryKey=$primaryKey, AQE=$adaptive, incremental=$incremental") {
+      withTable("runtime_fact") {
+        val properties =
+          if (primaryKey) "'primary.key'='id,dt', 'bucket.num'='1'"
+          else "'bucket.num'='1'"
+        sql(s"""CREATE TABLE runtime_fact (id INT, dt STRING) PARTITIONED BY (dt)
+               |TBLPROPERTIES($properties)""".stripMargin)
+        val startMs = System.currentTimeMillis()
+        sql("INSERT INTO runtime_fact VALUES (10, '0'), (11, '1'), (12, '2')")
+        spark
+          .range(3)
+          .selectExpr("cast(id as string) as dt", "id % 2 as flag")
+          .createOrReplaceTempView("runtime_keys")
+        try {
+          withSQLConf(
+            "spark.sql.adaptive.enabled" -> adaptive.toString,
+            "spark.sql.optimizer.dynamicPartitionPruning.reuseBroadcastOnly" -> "false",
+            "spark.sql.optimizer.dynamicPartitionPruning.useStats" -> "false",
+            "spark.sql.optimizer.dynamicPartitionPruning.fallbackFilterRatio" -> "1.0"
+          ) {
+            val endMs = System.currentTimeMillis() + 1L
+            val fact = if (incremental) {
+              s"fluss_incremental_between_timestamp('$DEFAULT_DATABASE.runtime_fact', '$startMs', '$endMs')"
+            } else "runtime_fact"
+            val query = s"""SELECT /*+ BROADCAST(d) */ f.id, f.dt FROM $fact f
+                           |JOIN runtime_keys d ON f.dt = d.dt WHERE d.flag = 0""".stripMargin
+            var baseline = Seq.empty[Row]
+            withSQLConf("spark.sql.optimizer.dynamicPartitionPruning.enabled" -> "false") {
+              baseline = sql(query).collect().toSeq
+            }
+            withSQLConf("spark.sql.optimizer.dynamicPartitionPruning.enabled" -> "true") {
+              val df = sql(query)
+              val rows = df.collect().toSeq
+              assertThat(rows.toArray).containsExactlyInAnyOrderElementsOf(baseline.asJava)
+              assertThat(rows.toArray).containsExactlyInAnyOrder(Row(10, "0"), Row(12, "2"))
+              val plan = df.queryExecution.executedPlan match {
+                case a: AdaptiveSparkPlanExec => a.executedPlan
+                case p => p
+              }
+              val scans = plan.collect { case b: BatchScanExec => b }
+              assertThat(scans.size).isEqualTo(1)
+              val batchScan = scans.head
+              assertThat(batchScan.runtimeFilters.nonEmpty).isTrue()
+              val scan = batchScan.scan.asInstanceOf[FlussScan]
+              assertThat(scan.description()).contains("RuntimePredicates")
+              assertThat(scan.timeRange.isDefined).isEqualTo(incremental)
+              val initial = batchScan.inputPartitions
+              val filtered = scan.toBatch.planInputPartitions()
+              assertThat(initial.size).isEqualTo(3)
+              assertThat(filtered.length).isEqualTo(2)
+              assertThat(filtered).isSubsetOf(initial.asJava)
+            }
+          }
+        } finally {
+          spark.catalog.dropTempView("runtime_keys")
+        }
+      }
+    }
+  }
+
+  for (primaryKey <- Seq(false, true)) {
+    test(
+      s"data written after planning retains runtime-filtered boundaries: primaryKey=$primaryKey") {
+      withTable("runtime_boundary") {
+        val properties =
+          if (primaryKey) "'primary.key'='id,dt', 'bucket.num'='1'"
+          else "'bucket.num'='1'"
+        sql(s"""CREATE TABLE runtime_boundary (id INT, dt STRING) PARTITIONED BY (dt)
+               |TBLPROPERTIES($properties)""".stripMargin)
+        sql("INSERT INTO runtime_boundary VALUES (1, 'a'), (2, 'b')")
+        val df = sql("SELECT * FROM runtime_boundary WHERE dt = 'a'")
+        val scan = df.queryExecution.optimizedPlan.collect {
+          case DataSourceV2ScanRelation(_, s: FlussScan, _, _, _) => s
+        }.head
+        val initial = scan.toBatch.planInputPartitions()
+        sql("INSERT INTO runtime_boundary VALUES (3, 'a'), (4, 'c')")
+        scan.filter(
+          Array(
+            new Predicate(
+              "IN",
+              Array[Expression](Expressions.column("dt"), Expressions.literal("a")))))
+        assertThat(scan.toBatch.planInputPartitions())
+          .containsExactlyElementsOf(initial.toSeq.asJava)
+        assertThat(df.collect()).containsExactly(Row(1, "a"))
+      }
+    }
+  }
+}