diff --git a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussScan.scala b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussScan.scala index 96def8ee814..17a4b22d20f 100644 --- a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/FlussScan.scala +++ b/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, diff --git a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/SplitPlanner.scala b/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/SplitPlanner.scala index 31f687be218..bd8c66fc44f 100644 --- a/fluss-spark/fluss-spark-common/src/main/scala/org/apache/fluss/spark/read/SplitPlanner.scala +++ b/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 } diff --git a/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/SparkRuntimeFilteringTest.scala b/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/SparkRuntimeFilteringTest.scala new file mode 100644 index 00000000000..ebc7314a0fc --- /dev/null +++ b/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")) + } + } + } +} diff --git a/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/lake/SparkLakeRuntimeFilteringTest.scala b/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/lake/SparkLakeRuntimeFilteringTest.scala new file mode 100644 index 00000000000..cdc77ee968c --- /dev/null +++ b/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/lake/SparkLakeRuntimeFilteringTest.scala @@ -0,0 +1,117 @@ +/* + * 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.lake + +import org.apache.fluss.config.{ConfigOptions, Configuration} +import org.apache.fluss.metadata.DataLakeFormat +import org.apache.fluss.spark.read.FlussScan +import org.apache.fluss.spark.read.lake.{FlussLakeInputPartition, FlussLakeUpsertInputPartition} + +import org.apache.spark.sql.Row +import org.apache.spark.sql.execution.datasources.v2.DataSourceV2ScanRelation +import org.assertj.core.api.Assertions.assertThat + +import java.nio.file.Files + +import scala.collection.JavaConverters._ + +/** Verifies runtime partition pruning across lake snapshots and Fluss log tails. */ +@SparkLakeTest +class SparkLakeRuntimeFilteringTest extends SparkLakeTableReadTestBase { + + override protected def dataLakeFormat: DataLakeFormat = DataLakeFormat.PAIMON + + override protected def flussConf: Configuration = { + val conf = super.flussConf + conf.setString("datalake.format", DataLakeFormat.PAIMON.toString) + conf.setString("datalake.paimon.metastore", "filesystem") + conf.setString("datalake.paimon.cache-enabled", "false") + warehousePath = Files.createTempDirectory("fluss-runtime-lake").resolve("warehouse").toString + conf.setString("datalake.paimon.warehouse", warehousePath) + conf + } + + override protected def lakeCatalogConf: Configuration = { + val conf = new Configuration() + conf.setString("metastore", "filesystem") + conf.setString("warehouse", warehousePath) + conf + } + + for (primaryKey <- Seq(false, true)) { + test(s"runtime filtering prunes lake-only and lake-union reads: primaryKey=$primaryKey") { + val tableName = if (primaryKey) "runtime_lake_pk" else "runtime_lake_log" + withTable(tableName) { + val primaryKeyProperty = if (primaryKey) "'primary.key'='id,dt'," else "" + sql(s"""CREATE TABLE $tableName (id INT, dt STRING) PARTITIONED BY (dt) + |TBLPROPERTIES($primaryKeyProperty 'bucket.num'='1', + |'${ConfigOptions.TABLE_DATALAKE_ENABLED.key()}'='true', + |'${ConfigOptions.TABLE_DATALAKE_FRESHNESS.key()}'='1s')""".stripMargin) + sql(s"INSERT INTO $tableName VALUES (10, '0'), (11, '1'), (12, '2')") + tierToLake(tableName) + spark + .range(5) + .selectExpr("cast(id as string) as dt", "id % 2 as flag") + .createOrReplaceTempView("runtime_lake_keys") + try { + withSQLConf( + "spark.sql.adaptive.enabled" -> "false", + "spark.sql.optimizer.dynamicPartitionPruning.reuseBroadcastOnly" -> "false", + "spark.sql.optimizer.dynamicPartitionPruning.useStats" -> "false", + "spark.sql.optimizer.dynamicPartitionPruning.fallbackFilterRatio" -> "1.0" + ) { + for (union <- Seq(false, true)) { + if (union) { + // Tail on an existing lake partition plus a partition found only in Fluss. + sql(s"INSERT INTO $tableName VALUES (20, '0'), (14, '4')") + } + val query = s"""SELECT /*+ BROADCAST(d) */ f.id, f.dt FROM $tableName f + |JOIN runtime_lake_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 scan = df.queryExecution.optimizedPlan.collect { + case DataSourceV2ScanRelation(_, s: FlussScan, _, _, _) => s + }.head + val initial = scan.toBatch.planInputPartitions() + assertThat(initial.exists { + case _: FlussLakeInputPartition | _: FlussLakeUpsertInputPartition => true + case _ => false + }).isTrue() + val rows = df.collect().toSeq + assertThat(rows.toArray).containsExactlyInAnyOrderElementsOf(baseline.asJava) + val expected = Seq(Row(10, "0"), Row(12, "2")) ++ + (if (union) Seq(Row(20, "0"), Row(14, "4")) else Seq.empty) + assertThat(rows.toArray).containsExactlyInAnyOrderElementsOf(expected.asJava) + val filtered = scan.toBatch.planInputPartitions() + assertThat(filtered.length).isLessThan(initial.length) + assertThat(filtered).isSubsetOf(initial.toSeq.asJava) + assertThat(scan.description()).contains("RuntimePredicates") + } + } + } + } finally { + spark.catalog.dropTempView("runtime_lake_keys") + } + } + } + } +} diff --git a/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/read/FlussRuntimeFilteringTest.scala b/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/read/FlussRuntimeFilteringTest.scala new file mode 100644 index 00000000000..f8fb73841dc --- /dev/null +++ b/fluss-spark/fluss-spark-ut/src/test/scala/org/apache/fluss/spark/read/FlussRuntimeFilteringTest.scala @@ -0,0 +1,288 @@ +/* + * 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.read + +import org.apache.fluss.config.Configuration +import org.apache.fluss.lake.source.LakeSplit +import org.apache.fluss.metadata.{Schema, TableBucket, TableInfo, TablePath} +import org.apache.fluss.spark.read.lake.{FlussLakeInputPartition, FlussLakeUpsertInputPartition} +import org.apache.fluss.types.DataTypes + +import org.apache.spark.sql.connector.expressions.{Expression, Expressions, Literal} +import org.apache.spark.sql.connector.expressions.filter.Predicate +import org.apache.spark.sql.connector.read.{Batch, InputPartition, PartitionReaderFactory} +import org.apache.spark.sql.types.{DataType, IntegerType, StringType, StructType} +import org.apache.spark.sql.util.CaseInsensitiveStringMap +import org.assertj.core.api.Assertions.assertThat +import org.mockito.Mockito.{mock, times, verify, when} +import org.scalatest.funsuite.AnyFunSuite + +import java.util.Collections + +import scala.collection.JavaConverters._ + +/** Verifies partition metadata, runtime filter semantics and the batch scan lifecycle. */ +class FlussRuntimeFilteringTest extends AnyFunSuite { + + private val path = TablePath.of("db", "t") + private val options = new CaseInsensitiveStringMap(Collections.emptyMap[String, String]()) + private val partitionedInfo = tableInfo(Seq("dt", "region")) + private val values = + Map(10L -> Seq("a", "east"), 20L -> Seq("b", "west"), 30L -> Seq("a", "west")) + private val window = Some(FlussTimeRange(100L, 200L)) + private val append = Array[InputPartition]( + FlussAppendInputPartition(bucket(10L), 4L, 9L, window), + FlussAppendInputPartition(bucket(20L), 1L, 7L, window), + FlussAppendInputPartition(bucket(30L), 2L, 8L, window), + FlussAppendInputPartition(bucket(10L), 9L, 12L, window) + ) + + test("filter attributes use projected partition columns and quote special names") { + val special = tableInfo(Seq("a.b", "a`b", "with space")) + val scan = testScan(special, Array.empty[InputPartition], Map.empty) + assertThat(scan.filterAttributes().map(_.fieldNames().toSeq)) + .containsExactly(Seq("a.b"), Seq("a`b"), Seq("with space")) + val projected = new TestScan( + special, + Some(new StructType().add("a`b", StringType)), + mock(classOf[AppendPlanner]), + mock(classOf[Batch])) + assertThat(projected.filterAttributes().map(_.fieldNames().toSeq)) + .containsExactly(Seq("a`b")) + assertThat(testScan(tableInfo(Seq.empty), append, Map.empty).filterAttributes().length) + .isZero() + val emptyProjection = + new TestScan( + partitionedInfo, + Some(new StructType()), + mock(classOf[AppendPlanner]), + mock(classOf[Batch])) + assertThat(emptyProjection.filterAttributes().length).isZero() + scan.filter(Array(in("`a.b`", "a"))) + assertThat(scan.description()).contains("RuntimePredicates") + } + + test("runtime predicates intersect while retaining initial split order and boundaries") { + val scan = testScan(partitionedInfo, append, values) + val batch = scan.toBatch + val initial = batch.planInputPartitions() + // A caller must not be able to mutate the cached plan through the returned array. + initial(0) = append(1) + scan.filter(Array(in("dt", "a"), in("region", "east", "west"))) + assertPartitions(scan, Seq(append(0), append(2), append(3))) + scan.filter(Array(in("region", "east"))) + assertPartitions(scan, Seq(append(0), append(3))) + scan.filter(Array(in("dt", "b"))) + assertPartitions(scan, Seq.empty) + scan.filter(Array(in("dt", "a"))) + assertPartitions(scan, Seq.empty) + assertThat(scan.toBatch).isSameAs(batch) + verify(scan.source, times(1)).planInputPartitions() + verify(scan.planner, times(1)).partitionValuesById + } + + test("filter before planning captures one plan and ignores subsequent planner changes") { + val scan = testScan(partitionedInfo, append, values) + scan.filter(Array(in("dt", "a"))) + when(scan.source.planInputPartitions()).thenReturn(Array(append(1))) + when(scan.planner.partitionValuesById).thenReturn(Map(10L -> Seq("b", "east"))) + assertPartitions(scan, Seq(append(0), append(2), append(3))) + scan.filter(Array(in("region", "east"))) + assertPartitions(scan, Seq(append(0), append(3))) + verify(scan.source, times(1)).planInputPartitions() + } + + test("empty and unsupported filters are no-ops, empty partition IN removes all splits") { + val scan = testScan(partitionedInfo, append, values) + scan.filter(Array.empty[Predicate]) + scan.filter(Array(in("id"), in("missing"), new Predicate("UNSUPPORTED", Array.empty))) + verify(scan.source, times(0)).planInputPartitions() + assertPartitions(scan, append.toSeq) + scan.filter(Array(in("dt"))) + assertPartitions(scan, Seq.empty) + val unpartitioned = testScan(tableInfo(Seq.empty), append, Map.empty) + unpartitioned.filter(Array(in("id"))) + assertPartitions(unpartitioned, append.toSeq) + } + + test("null IN literals follow SQL semantics and static pruning stays in effect") { + val scan = testScan(partitionedInfo, Array(append(1), append(2)), values) + scan.filter(Array(in("dt", "a", null))) + assertPartitions(scan, Seq(append(2))) + scan.filter(Array(in("dt", null))) + assertPartitions(scan, Seq.empty) + val nullPartition = testScan(partitionedInfo, Array(append(0)), Map(10L -> Seq(null, "east"))) + nullPartition.filter( + Array( + new Predicate("<=>", Array[Expression](Expressions.column("dt"), lit(null, StringType))))) + assertPartitions(nullPartition, Seq(append(0))) + } + + test("lake-only, Fluss-only and combined splits use their own partition metadata") { + val lakeA = lakeSplit(Seq("a", "east")) + val lakeB = lakeSplit(Seq("b", "west")) + val unknown = lakeSplit(Seq("a")) + val splits = Array[InputPartition]( + FlussLakeInputPartition(bucket(-1L), lakeA), + FlussLakeInputPartition(bucket(-1L), lakeB), + FlussLakeUpsertInputPartition(bucket(10L), Seq(lakeA).asJava, 5L, 9L), + FlussLakeUpsertInputPartition(bucket(20L), null, 2L, 8L), + FlussLakeUpsertInputPartition(bucket(-1L), Seq(lakeB).asJava, 5L, 5L), + FlussLakeInputPartition(bucket(-1L), unknown), + FlussUpsertInputPartition(bucket(30L), 42L, 6L, 11L, window), + FlussAppendInputPartition(bucket(99L), 2L, 7L, window) + ) + val scan = testScan(partitionedInfo, splits, values) + scan.filter(Array(in("dt", "a"))) + assertPartitions(scan, Seq(splits(0), splits(2), splits(5), splits(6), splits(7))) + // Even incomplete metadata cannot match an empty runtime key set. + scan.filter(Array(in("region"))) + assertPartitions(scan, Seq.empty) + } + + test("incomplete lake metadata and missing Fluss partition ids are retained conservatively") { + val splits = Array[InputPartition]( + FlussLakeInputPartition(bucket(-1L), lakeSplit(null)), + FlussLakeInputPartition(bucket(-1L), lakeSplit(Seq.empty)), + FlussLakeInputPartition(bucket(-1L), lakeSplit(Seq("a", "east", "extra"))), + FlussLakeUpsertInputPartition(bucket(99L), Seq(lakeSplit(Seq("b", "west"))).asJava, 1L, 2L), + new InputPartition {} + ) + val scan = testScan(partitionedInfo, splits, values) + scan.filter(Array(in("dt", "a"))) + assertPartitions(scan, splits.toSeq) + val numericInfo = tableInfo(Seq("dt"), numeric = true) + val invalid = testScan(numericInfo, Array(append(0)), Map(10L -> Seq("invalid"))) + invalid.filter( + Array( + new Predicate( + "IN", + Array[Expression](Expressions.column("dt"), lit(Integer.valueOf(1), IntegerType))))) + assertPartitions(invalid, Seq(append(0))) + } + + test("append and upsert scans reuse their batches and original reader factories") { + val appendPlanner = mock(classOf[AppendPlanner]) + when(appendPlanner.plan()).thenReturn(append) + when(appendPlanner.partitionValuesById).thenReturn(values) + val appendScan = FlussAppendScan( + path, + partitionedInfo, + None, + None, + None, + Seq.empty, + None, + options, + new Configuration(), + appendPlanner) + val upsertPlanner = mock(classOf[UpsertPlanner]) + val upsert = FlussUpsertInputPartition(bucket(10L), 42L, 6L, 11L, window) + when(upsertPlanner.plan()).thenReturn(Array[InputPartition](upsert)) + when(upsertPlanner.partitionValuesById).thenReturn(values) + val upsertScan = FlussUpsertScan( + path, + partitionedInfo, + None, + None, + None, + Seq.empty, + None, + options, + new Configuration(), + upsertPlanner) + Seq(appendScan, upsertScan).foreach { + scan => + val batch = scan.toBatch + val factory = batch.createReaderFactory() + batch.planInputPartitions() + scan.filter(Array(in("dt", "a"))) + assertThat(scan.toBatch).isSameAs(batch) + assertThat(scan.toBatch.createReaderFactory()).isSameAs(factory) + assertThat(scan.toBatch.planInputPartitions().length).isPositive() + } + assertThat(upsertScan.toBatch.planInputPartitions()(0)).isSameAs(upsert) + verify(appendPlanner, times(1)).plan() + verify(upsertPlanner, times(1)).plan() + } + + private def tableInfo(keys: Seq[String], numeric: Boolean = false): TableInfo = { + val schema = Schema.newBuilder().column("id", DataTypes.INT()) + keys.foreach(k => schema.column(k, if (numeric) DataTypes.INT() else DataTypes.STRING())) + new TableInfo( + path, + 1L, + 1, + schema.build(), + Collections.emptyList[String](), + keys.asJava, + 1, + new Configuration(), + new Configuration(), + null, + null, + 0L, + 0L) + } + + private def bucket(id: Long): TableBucket = new TableBucket(1L, id, 0) + + private def in(column: String, values: String*): Predicate = + new Predicate( + "IN", + Array[Expression](Expressions.column(column)) ++ + values.map(v => lit(v, StringType))) + + private def lit[T](v: T, tpe: DataType): Literal[T] = new Literal[T] { + override def value(): T = v + override def dataType(): DataType = tpe + } + + private def lakeSplit(values: Seq[String]): LakeSplit = { + val split = mock(classOf[LakeSplit]) + when(split.partition()).thenReturn(if (values == null) null else values.asJava) + split + } + + private def testScan( + partitionedInfo: TableInfo, + partitions: Array[InputPartition], + values: Map[Long, Seq[String]]): TestScan = { + val planner = mock(classOf[AppendPlanner]) + when(planner.partitionValuesById).thenReturn(values) + val source = mock(classOf[Batch]) + when(source.planInputPartitions()).thenReturn(partitions) + when(source.createReaderFactory()).thenReturn(mock(classOf[PartitionReaderFactory])) + new TestScan(partitionedInfo, None, planner, source) + } + + private def assertPartitions(scan: FlussScan, expected: Seq[InputPartition]): Unit = + assertThat(scan.toBatch.planInputPartitions()) + .containsExactlyElementsOf(expected.asJava) + + private class TestScan( + val tableInfo: TableInfo, + val requiredSchema: Option[StructType], + val planner: AppendPlanner, + val source: Batch) + extends FlussScan { + override def tablePath: TablePath = path + override protected def scanType: String = "Test" + override protected def createBatch(): Batch = source + } +}