Skip to content

[AURON #2390] Preserve key-grouped partitions in native Iceberg scans #2391

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lyne7-sc wants to merge 5 commits into
base: master
Choose a base branch
from
Open

[AURON #2390] Preserve key-grouped partitions in native Iceberg scans #2391

Show file tree
Hide file tree
Changes from 4 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
23b286a
test: cover key grouped native Iceberg scans
lyne7-sc Jul 15, 2026
0a846f5
preserve key grouped partitioning for native Iceberg scans
lyne7-sc Jul 15, 2026
36c2db3
simplify
lyne7-sc Jul 16, 2026
f9a73a7
test: fix Iceberg test for spark 3.4
lyne7-sc Jul 16, 2026
776d798
ignore no-op DPP filters in key-grouped Iceberg scans
lyne7-sc Jul 19, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
109 changes: 97 additions & 12 deletions .../auron-iceberg/src/main/scala/org/apache/spark/sql/auron/iceberg/IcebergScanSupport.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,9 +28,10 @@ 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.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}

Expand All @@ -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(
Expand Down Expand Up @@ -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.")
Expand All @@ -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
Expand Down Expand Up @@ -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,
Expand All @@ -212,7 +229,8 @@ object IcebergScanSupport extends Logging {
fileSchema,
partitionSchema,
pruningPredicates,
fieldIdsByName))
fieldIdsByName,
groupedNativeTasks))
}

private def planChangelogScan(
Expand All @@ -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(
Expand All @@ -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
}
Expand Down Expand Up @@ -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,
Expand All @@ -300,7 +333,8 @@ object IcebergScanSupport extends Logging {
fileSchema,
partitionSchema,
pruningPredicates,
fieldIdsByName))
fieldIdsByName,
groupedNativeTasks))
}

private def inspectFieldIdSupport(
Expand Down Expand Up @@ -391,6 +425,57 @@ 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.
if (exec.runtimeFilters.nonEmpty) {

@weiqingy weiqingy Jul 17, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

runtimeFilters.nonEmpty and "runtime filtering actually happens" come apart in one case. With spark.sql.optimizer.dynamicPartitionPruning.reuseBroadcastOnly on (the default) and no reusable broadcast exchange to match, PlanDynamicPruningFilters substitutes DynamicPruningExpression(Literal.TrueLiteral) (v3.5.8, PlanDynamicPruningFilters.scala:78-80) — a filter that filters nothing. The AQE path does the same at PlanAdaptiveDynamicPruningFilters.scala:66-67, so this isn't limited to non-adaptive execution. Spark treats that value as absent in two places: doCanonicalize strips it explicitly (BatchScanExec.scala:237-239), and filteredPartitions translates runtime filters through DataSourceV2Strategy.translateRuntimeFilterV2, which only matches InSubqueryExec (DataSourceV2Strategy.scala:644-655) — so TrueLiteral translates to nothing, filteredPartitions returns partitions unfiltered, and the groups can't change.

Where that lands: select ... from fact f join dim d on f.p = d.p where d.x = 1, both tables Iceberg-partitioned on p, gets SPJ precisely because it's a sort-merge join (dim too large to broadcast). No BHJ ⇒ no reusable exchange ⇒ the DPP filter degrades to TrueLiteral ⇒ runtimeFilters.nonEmpty is true ⇒ the whole scan drops back to Spark, with zero runtime filtering having been performed.

That costs native coverage rather than correctness — on master this query went native and returned wrong rows, so this is better either way. Is the wider guard deliberate, or would mirroring Spark's own doCanonicalize predicate be closer to the intent? One option, in case it's useful:

val hasEffectiveRuntimeFilters = exec.runtimeFilters
  .exists(_ != DynamicPruningExpression(Literal.TrueLiteral))
if (hasEffectiveRuntimeFilters) None else keyGroupedInputPartitions(exec, partitioning)…

(Literal is already imported at :30; only DynamicPruningExpression would need adding.)

Related: would a test pinning this branch help? Nothing in the suite reaches it today — spark.sql.sources.v2.bucketing.enabled appears only in the two new tests, and every DPP test carries an explicit /*+ BROADCAST(d) */ hint (:281-340, and also :357 / :675), which forces the reusable exchange and a real InSubqueryExec. An SPJ query carrying runtime filters, asserting a clean fallback (right answer, no NativeIcebergTableScanExec, no exception), would pin the intended behavior — and if it turns out awkward to write because TrueLiteral makes the fallback fire where you'd expect native execution, that's the answer to the question above.

@lyne7-sc lyne7-sc Jul 19, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the detailed explanation! That makes sense to me.

I updated the check to ignore DynamicPruningExpression(Literal.TrueLiteral) while continuing to fall back to Spark for effective runtime filters.

I also added coverage for both sides:

  • A no-op TrueLiteral runtime filter keeps the scan native and preserves KeyGroupedPartitioning.
  • An effective runtime filter keeps the fact scan as BatchScanExec.

@weiqingy weiqingy Jul 19, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the fix and the tests on both sides. Nothing further from me.

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] = {
Expand Down
33 changes: 25 additions & 8 deletions ...src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeIcebergTableScanExec.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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

Expand All @@ -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()
Expand Down
100 changes: 100 additions & 0 deletions .../auron-iceberg/src/test/scala/org/apache/auron/iceberg/AuronIcebergIntegrationSuite.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,13 +30,15 @@ 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.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
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
Expand Down Expand Up @@ -178,6 +180,104 @@ 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("iceberg native scan preserves dynamic pruning runtime filters") {
withTable("local.db.t_dpp_fact", "local.db.t_dpp_dim") {
sql("""
Expand Down
Loading