diff --git a/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergScanSupport.scala b/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergScanSupport.scala index 31d877930..ccff369a1 100644 --- a/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergScanSupport.scala +++ b/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergScanSupport.scala @@ -27,10 +27,11 @@ import org.apache.iceberg.expressions.{And => IcebergAnd, BoundPredicate, Expres import org.apache.iceberg.spark.source.AuronIcebergSourceUtil import org.apache.spark.internal.Logging import org.apache.spark.sql.auron.{NativeConverters, Shims} -import org.apache.spark.sql.catalyst.expressions.{And => SparkAnd, AttributeReference, EqualTo, Expression => SparkExpression, GreaterThan, GreaterThanOrEqual, In, IsNaN, IsNotNull, IsNull, LessThan, LessThanOrEqual, Literal, Not => SparkNot, Or => SparkOr} +import org.apache.spark.sql.catalyst.expressions.{And => SparkAnd, AttributeReference, DynamicPruningExpression, EqualTo, Expression => SparkExpression, GreaterThan, GreaterThanOrEqual, In, IsNaN, IsNotNull, IsNull, LessThan, LessThanOrEqual, Literal, Not => SparkNot, Or => SparkOr} +import org.apache.spark.sql.catalyst.plans.physical.KeyGroupedPartitioning import org.apache.spark.sql.catalyst.trees.TreeNodeTag import org.apache.spark.sql.connector.read.{InputPartition, Scan} -import org.apache.spark.sql.execution.datasources.v2.BatchScanExec +import org.apache.spark.sql.execution.datasources.v2.{BatchScanExec, DataSourceRDDPartition} import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{BinaryType, DataType, DecimalType, StringType, StructField, StructType} @@ -53,7 +54,8 @@ final case class IcebergScanPlan( fileSchema: StructType, partitionSchema: StructType, pruningPredicates: Seq[pb.PhysicalExprNode], - fieldIdsByName: Map[String, Int]) + fieldIdsByName: Map[String, Int], + groupedScanTasks: Option[Seq[Seq[IcebergNativeScanTask]]] = None) object IcebergScanSupport extends Logging { private val scanPlanTag: TreeNodeTag[Option[IcebergScanPlan]] = TreeNodeTag( @@ -154,7 +156,11 @@ object IcebergScanSupport extends Logging { case None => return None } - val partitions = inputPartitions(exec, useRuntimeFilters) + val (partitions, partitionGroups) = + plannedInputPartitions(exec, useRuntimeFilters) match { + case Some(plannedPartitions) => plannedPartitions + case None => return None + } // Empty scan (e.g. empty table) should still build a plan to return no rows. if (partitions.isEmpty) { logWarning(s"Native Iceberg scan planned with empty partitions for $scanClassName.") @@ -166,10 +172,14 @@ object IcebergScanSupport extends Logging { fileSchema, partitionSchema, Seq.empty, - fieldIdsByName)) + fieldIdsByName, + partitionGroups.map(_.map(_ => Seq.empty)))) } - val icebergPartitions = partitions.flatMap(icebergPartition) + val icebergPartitionGroups = + partitionGroups.map(_.map(_.flatMap(icebergPartition))) + val icebergPartitions = + icebergPartitionGroups.map(_.flatten).getOrElse(partitions.flatMap(icebergPartition)) // All partitions must be Iceberg SparkInputPartition with file scan tasks; otherwise fallback. if (icebergPartitions.size != partitions.size) { return None @@ -203,7 +213,14 @@ object IcebergScanSupport extends Logging { } val pruningPredicates = collectPruningPredicates(scan.asInstanceOf[AnyRef], readSchema) - val nativeTasks = fileTasks.map(task => toNativeScanTask(task, partitionSchema)) + val groupedNativeTasks = icebergPartitionGroups.map(_.map { group => + group + .flatMap(_.tasks) + .collect { case task: FileScanTask => toNativeScanTask(task, partitionSchema) } + }) + val nativeTasks = groupedNativeTasks + .map(_.flatten) + .getOrElse(fileTasks.map(task => toNativeScanTask(task, partitionSchema))) Some( IcebergScanPlan( nativeTasks, @@ -212,7 +229,8 @@ object IcebergScanSupport extends Logging { fileSchema, partitionSchema, pruningPredicates, - fieldIdsByName)) + fieldIdsByName, + groupedNativeTasks)) } private def planChangelogScan( @@ -236,7 +254,11 @@ object IcebergScanSupport extends Logging { case None => return None } - val partitions = inputPartitions(exec, useRuntimeFilters) + val (partitions, partitionGroups) = + plannedInputPartitions(exec, useRuntimeFilters) match { + case Some(plannedPartitions) => plannedPartitions + case None => return None + } if (partitions.isEmpty) { return Some( IcebergScanPlan( @@ -246,10 +268,14 @@ object IcebergScanSupport extends Logging { fileSchema, partitionSchema, Seq.empty, - fieldIdsByName)) + fieldIdsByName, + partitionGroups.map(_.map(_ => Seq.empty)))) } - val icebergPartitions = partitions.flatMap(icebergPartition) + val icebergPartitionGroups = + partitionGroups.map(_.map(_.flatMap(icebergPartition))) + val icebergPartitions = + icebergPartitionGroups.map(_.flatten).getOrElse(partitions.flatMap(icebergPartition)) if (icebergPartitions.size != partitions.size) { return None } @@ -291,7 +317,14 @@ object IcebergScanSupport extends Logging { } val pruningPredicates = collectPruningPredicates(scan.asInstanceOf[AnyRef], readSchema) - val nativeTasks = addedRowsTasks.map(task => toNativeScanTask(task, partitionSchema)) + val groupedNativeTasks = icebergPartitionGroups.map(_.map { group => + group + .flatMap(_.tasks) + .collect { case task: AddedRowsScanTask => toNativeScanTask(task, partitionSchema) } + }) + val nativeTasks = groupedNativeTasks + .map(_.flatten) + .getOrElse(addedRowsTasks.map(task => toNativeScanTask(task, partitionSchema))) Some( IcebergScanPlan( nativeTasks, @@ -300,7 +333,8 @@ object IcebergScanSupport extends Logging { fileSchema, partitionSchema, pruningPredicates, - fieldIdsByName)) + fieldIdsByName, + groupedNativeTasks)) } private def inspectFieldIdSupport( @@ -391,6 +425,59 @@ object IcebergScanSupport extends Logging { private def deletesEmpty(deletes: java.util.List[_]): Boolean = deletes == null || deletes.isEmpty + private def plannedInputPartitions(exec: BatchScanExec, useRuntimeFilters: Boolean) + : Option[(Seq[InputPartition], Option[Seq[Seq[InputPartition]]])] = { + exec.outputPartitioning match { + case partitioning: KeyGroupedPartitioning => + // Runtime filtering can change the final groups after static planning. Keep this + // combination on Spark until native execution can preserve those dynamic groups. + val hasEffectiveRuntimeFilters = exec.runtimeFilters + .exists(_ != DynamicPruningExpression(Literal.TrueLiteral)) + if (hasEffectiveRuntimeFilters) { + None + } else { + keyGroupedInputPartitions(exec, partitioning) + .map(groups => groups.flatten -> Some(groups)) + } + case _ => + Some(inputPartitions(exec, useRuntimeFilters) -> None) + } + } + + private def keyGroupedInputPartitions( + exec: BatchScanExec, + partitioning: KeyGroupedPartitioning): Option[Seq[Seq[InputPartition]]] = { + try { + // BatchScanExec.inputRDD contains Spark's final partition groups, including empty groups + // introduced while aligning both sides of a storage-partitioned join. + val rddPartitions = exec.inputRDD.partitions.toSeq + val dataSourcePartitions = rddPartitions.collect { case partition: DataSourceRDDPartition => + partition + } + if (dataSourcePartitions.size != rddPartitions.size) { + logWarning( + s"Expected DataSourceRDDPartition for every key-grouped Iceberg partition in " + + s"${exec.getClass.getName}.") + return None + } + + val groups = dataSourcePartitions.map(_.inputPartitions.toSeq) + if (groups.size != partitioning.numPartitions) { + logWarning( + s"Key-grouped Iceberg partition count mismatch: planned ${groups.size}, " + + s"declared ${partitioning.numPartitions}.") + return None + } + Some(groups) + } catch { + case NonFatal(t) => + logWarning( + s"Failed to obtain final key-grouped input partitions for ${exec.getClass.getName}.", + t) + None + } + } + private def inputPartitions( exec: BatchScanExec, useRuntimeFilters: Boolean): Seq[InputPartition] = { diff --git a/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeIcebergTableScanExec.scala b/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeIcebergTableScanExec.scala index ae608630c..907b42d39 100644 --- a/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeIcebergTableScanExec.scala +++ b/thirdparty/auron-iceberg/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeIcebergTableScanExec.scala @@ -34,7 +34,7 @@ import org.apache.spark.sql.auron.{EmptyNativeRDD, NativeConverters, NativeHelpe import org.apache.spark.sql.auron.iceberg.{IcebergNativeScanTask, IcebergScanPlan, IcebergScanSupport} import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Expression, GenericInternalRow, Literal} -import org.apache.spark.sql.catalyst.plans.physical.SinglePartition +import org.apache.spark.sql.catalyst.plans.physical.{KeyGroupedPartitioning, SinglePartition} import org.apache.spark.sql.execution.{LeafExecNode, SparkPlan, SQLExecution} import org.apache.spark.sql.execution.datasources.{FilePartition, PartitionedFile} import org.apache.spark.sql.execution.datasources.v2.BatchScanExec @@ -278,6 +278,25 @@ case class NativeIcebergTableScanExec( } private def buildFilePartitions(): Array[FilePartition] = { + outputPartitioning match { + case partitioning: KeyGroupedPartitioning => + val taskGroups = plan.groupedScanTasks.getOrElse { + throw new IllegalStateException( + "Missing grouped Iceberg scan tasks for KeyGroupedPartitioning.") + } + require( + taskGroups.size == partitioning.numPartitions, + s"Key-grouped Iceberg task count ${taskGroups.size} did not match declared " + + s"partition count ${partitioning.numPartitions}.") + require( + taskGroups.flatten == scanTasks, + "Key-grouped Iceberg tasks did not flatten to the planned scan tasks.") + return taskGroups.zipWithIndex.map { case (tasks, index) => + FilePartition(index, tasks.map(partitionedFile).toArray) + }.toArray + case _ => + } + // Convert Iceberg scan tasks into Spark FilePartition groups for execution. if (scanTasks.isEmpty) { return Array.empty @@ -286,13 +305,7 @@ case class NativeIcebergTableScanExec( val sparkSession = Shims.get.getSqlContext(basedScan).sparkSession val maxSplitBytes = getMaxSplitBytes(sparkSession, scanTasks) val partitionedFiles = scanTasks - .map { task => - Shims.get.getPartitionedFile( - partitionValuesRow(task), - task.location, - task.start, - task.length) - } + .map(partitionedFile) .sortBy(_.length)(Ordering[Long].reverse) .toSeq @@ -304,6 +317,10 @@ case class NativeIcebergTableScanExec( } } + private def partitionedFile(task: IcebergNativeScanTask): PartitionedFile = { + Shims.get.getPartitionedFile(partitionValuesRow(task), task.location, task.start, task.length) + } + private def partitionValuesRow(task: IcebergNativeScanTask): InternalRow = { val values = partitionSchema.fields.zip(task.partitionValues).map { case (field, value) => Literal.create(value, field.dataType).eval() diff --git a/thirdparty/auron-iceberg/src/test/scala/org/apache/auron/iceberg/AuronIcebergIntegrationSuite.scala b/thirdparty/auron-iceberg/src/test/scala/org/apache/auron/iceberg/AuronIcebergIntegrationSuite.scala index 6d342bd48..fdd9bc6bb 100644 --- a/thirdparty/auron-iceberg/src/test/scala/org/apache/auron/iceberg/AuronIcebergIntegrationSuite.scala +++ b/thirdparty/auron-iceberg/src/test/scala/org/apache/auron/iceberg/AuronIcebergIntegrationSuite.scala @@ -30,6 +30,8 @@ import org.apache.iceberg.spark.Spark3Util import org.apache.spark.scheduler.{SparkListener, SparkListenerEvent} import org.apache.spark.sql.{DataFrame, Row} import org.apache.spark.sql.auron.iceberg.IcebergScanSupport +import org.apache.spark.sql.catalyst.expressions.{DynamicPruningExpression, Literal} +import org.apache.spark.sql.catalyst.plans.physical.KeyGroupedPartitioning import org.apache.spark.sql.catalyst.trees.TreeNodeTag import org.apache.spark.sql.execution.ExplainUtils.collectFirst import org.apache.spark.sql.execution.FormattedMode @@ -37,6 +39,7 @@ import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, QueryStageExec} import org.apache.spark.sql.execution.auron.plan.NativeIcebergTableScanExec import org.apache.spark.sql.execution.datasources.v2.BatchScanExec +import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike import org.apache.spark.sql.execution.ui.SparkListenerDriverAccumUpdates class AuronIcebergIntegrationSuite @@ -178,6 +181,199 @@ class AuronIcebergIntegrationSuite } } + test("native iceberg scan preserves key grouped partitioning") { + withTable("local.db.t_partitioned_join_left", "local.db.t_partitioned_join_right") { + sql(""" + |create table local.db.t_partitioned_join_left (id int, p int) + |using iceberg + |partitioned by (p) + |""".stripMargin) + sql("insert into local.db.t_partitioned_join_left values (0, 0)") + sql("insert into local.db.t_partitioned_join_left values (1, 1)") + + sql(""" + |create table local.db.t_partitioned_join_right (value int, p int) + |using iceberg + |partitioned by (p) + |""".stripMargin) + sql("insert into local.db.t_partitioned_join_right values (10, 0)") + sql("insert into local.db.t_partitioned_join_right values (11, 0)") + sql("insert into local.db.t_partitioned_join_right values (12, 1)") + + withSQLConf( + "spark.sql.adaptive.enabled" -> "false", + "spark.sql.autoBroadcastJoinThreshold" -> "-1", + "spark.sql.sources.v2.bucketing.enabled" -> "true", + "spark.sql.sources.v2.bucketing.pushPartValues.enabled" -> "true", + "spark.sql.iceberg.planning.preserve-data-grouping" -> "true", + "spark.sql.files.maxPartitionBytes" -> "1") { + val df = sql(""" + |select l.id, l.p, r.value + |from local.db.t_partitioned_join_left l + |join local.db.t_partitioned_join_right r on l.p = r.p + |""".stripMargin) + + val sourceScans = df.queryExecution.sparkPlan.collect { case scan: BatchScanExec => scan } + assert(sourceScans.size === 2) + assert(sourceScans.forall(_.outputPartitioning.isInstanceOf[KeyGroupedPartitioning])) + + val nativeScans = df.queryExecution.executedPlan.collect { + case scan: NativeIcebergTableScanExec => scan + } + assert(nativeScans.size === 2) + assert(df.queryExecution.executedPlan.collect { case e: ShuffleExchangeLike => + e + }.isEmpty) + + checkAnswer(df, Seq(Row(0, 0, 10), Row(0, 0, 11), Row(1, 1, 12))) + } + } + } + + test("native iceberg scan preserves empty key grouped partitions") { + withTable("local.db.t_grouped_left", "local.db.t_grouped_right") { + sql(""" + |create table local.db.t_grouped_left (id int, p int) + |using iceberg + |partitioned by (p) + |""".stripMargin) + sql("insert into local.db.t_grouped_left values (0, 0), (1, 1), (2, 2)") + + sql(""" + |create table local.db.t_grouped_right (value int, p int) + |using iceberg + |partitioned by (p) + |""".stripMargin) + sql("insert into local.db.t_grouped_right values (11, 1), (12, 2)") + + withSQLConf( + "spark.sql.adaptive.enabled" -> "false", + "spark.sql.autoBroadcastJoinThreshold" -> "-1", + "spark.sql.sources.v2.bucketing.enabled" -> "true", + "spark.sql.sources.v2.bucketing.pushPartValues.enabled" -> "true", + "spark.sql.iceberg.planning.preserve-data-grouping" -> "true") { + val df = sql(""" + |select l.id, l.p, r.value + |from local.db.t_grouped_left l + |join local.db.t_grouped_right r on l.p = r.p + |""".stripMargin) + + val sourceScans = df.queryExecution.sparkPlan.collect { case scan: BatchScanExec => scan } + assert(sourceScans.size === 2) + assert(sourceScans.forall(_.outputPartitioning.isInstanceOf[KeyGroupedPartitioning])) + + val nativeScans = df.queryExecution.executedPlan.collect { + case scan: NativeIcebergTableScanExec => scan + } + assert(nativeScans.size === 2) + assert(df.queryExecution.executedPlan.collect { case e: ShuffleExchangeLike => + e + }.isEmpty) + + checkAnswer(df, Seq(Row(1, 1, 11), Row(2, 2, 12))) + nativeScans.foreach { scan => + val partitioning = scan.outputPartitioning.asInstanceOf[KeyGroupedPartitioning] + assert(scan.metrics("numPartitions").value === partitioning.numPartitions) + } + } + } + } + + test("native iceberg scan ignores ineffective runtime filters for key grouped partitions") { + withTable("local.db.t_ineffective_dpp_fact", "local.db.t_ineffective_dpp_dim") { + sql(""" + |create table local.db.t_ineffective_dpp_fact (id int, p int) + |using iceberg + |partitioned by (p) + |""".stripMargin) + sql("insert into local.db.t_ineffective_dpp_fact values (1, 1), (2, 2), (3, 3)") + + sql(""" + |create table local.db.t_ineffective_dpp_dim (value int, p int) + |using iceberg + |partitioned by (p) + |""".stripMargin) + sql("insert into local.db.t_ineffective_dpp_dim values (10, 1), (20, 2), (30, 3)") + + withSQLConf( + "spark.sql.adaptive.enabled" -> "false", + "spark.sql.autoBroadcastJoinThreshold" -> "-1", + "spark.sql.optimizer.dynamicPartitionPruning.enabled" -> "true", + "spark.sql.optimizer.dynamicPartitionPruning.reuseBroadcastOnly" -> "true", + "spark.sql.sources.v2.bucketing.enabled" -> "true", + "spark.sql.sources.v2.bucketing.pushPartValues.enabled" -> "true", + "spark.sql.iceberg.planning.preserve-data-grouping" -> "true") { + val df = sql(""" + |select f.id, f.p, d.value + |from local.db.t_ineffective_dpp_fact f + |join local.db.t_ineffective_dpp_dim d on f.p = d.p + |where d.value = 20 + |""".stripMargin) + + checkAnswer(df, Seq(Row(2, 2, 20))) + val factScan = executedNativeIcebergTableScanExec(df, "t_ineffective_dpp_fact") + val explain = df.queryExecution.explainString(FormattedMode) + assert(factScan.runtimeFilters.nonEmpty, explain) + assert( + factScan.runtimeFilters.forall(_ == DynamicPruningExpression(Literal.TrueLiteral)), + explain) + assert(factScan.outputPartitioning.isInstanceOf[KeyGroupedPartitioning], explain) + } + } + } + + test("iceberg scan falls back for effective runtime filters on key grouped partitions") { + withTable("local.db.t_effective_dpp_fact", "local.db.t_effective_dpp_dim") { + sql(""" + |create table local.db.t_effective_dpp_fact (id int, p int) + |using iceberg + |partitioned by (p) + |""".stripMargin) + sql("insert into local.db.t_effective_dpp_fact values (1, 1), (2, 2), (3, 3)") + + sql(""" + |create table local.db.t_effective_dpp_dim (value int, p int) + |using iceberg + |partitioned by (p) + |""".stripMargin) + sql("insert into local.db.t_effective_dpp_dim values (10, 1), (20, 2), (30, 3)") + + withSQLConf( + "spark.sql.adaptive.enabled" -> "false", + "spark.sql.autoBroadcastJoinThreshold" -> "-1", + "spark.sql.optimizer.dynamicPartitionPruning.enabled" -> "true", + "spark.sql.optimizer.dynamicPartitionPruning.reuseBroadcastOnly" -> "false", + "spark.sql.sources.v2.bucketing.enabled" -> "true", + "spark.sql.sources.v2.bucketing.pushPartValues.enabled" -> "true", + "spark.sql.iceberg.planning.preserve-data-grouping" -> "true") { + val df = sql(""" + |select f.id, f.p, d.value + |from local.db.t_effective_dpp_fact f + |join local.db.t_effective_dpp_dim d on f.p = d.p + |where d.value = 20 + |""".stripMargin) + + checkAnswer(df, Seq(Row(2, 2, 20))) + val materializedPlan = collectMaterializedPlans(df.queryExecution.executedPlan) + val factScan = materializedPlan + .collectFirst { + case scan: BatchScanExec + if scan.scan.description().contains("t_effective_dpp_fact") => + scan + } + .getOrElse { + fail(df.queryExecution.explainString(FormattedMode)) + } + val explain = df.queryExecution.explainString(FormattedMode) + assert( + factScan.runtimeFilters + .exists(_ != DynamicPruningExpression(Literal.TrueLiteral)), + explain) + assert(factScan.outputPartitioning.isInstanceOf[KeyGroupedPartitioning], explain) + } + } + } + test("iceberg native scan preserves dynamic pruning runtime filters") { withTable("local.db.t_dpp_fact", "local.db.t_dpp_dim") { sql("""