akshaytayal commented on code in PR #13126: URL: https://github.com/apache/gluten/pull/13126#discussion_r4108288811
########## shims/spark42/src/main/scala/org/apache/spark/sql/execution/datasources/v2/BatchScanExecShim.scala: ########## @@ -0,0 +1,155 @@ +/* + * 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.spark.sql.execution.datasources.v2 + +import org.apache.spark.SparkException +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.plans.physical.KeyedPartitioning +import org.apache.spark.sql.catalyst.util.InternalRowComparableWrapper +import org.apache.spark.sql.connector.catalog.Table +import org.apache.spark.sql.connector.catalog.functions.Reducer +import org.apache.spark.sql.connector.expressions.aggregate.Aggregation +import org.apache.spark.sql.connector.read.{HasPartitionKey, InputPartition, Scan} +import org.apache.spark.sql.execution.datasources.v2.orc.OrcScan +import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan +import org.apache.spark.sql.execution.metric.SQLMetric +import org.apache.spark.sql.vectorized.ColumnarBatch + +// Spark 4.2 removed `StoragePartitionJoinParams` and no longer accepts the SPJ parameters +// (`joinKeyPositions`, `commonPartitionValues`, `reducers`, `applyPartialClustering`, +// `replicatePartitions`) on the scan node -- that grouping/replication now happens in +// `GroupPartitionsExec`. To keep the public constructor identical to the other Spark shims +// (Gluten's own planner reads these vals), they are kept here as shim-local fields and are simply +// not forwarded into the Spark superclass, which now only takes `keyGroupedPartitioning`. +abstract class BatchScanExecShim( + output: Seq[AttributeReference], + @transient scan: Scan, + runtimeFilters: Seq[Expression], + keyGroupedPartitioning: Option[Seq[Expression]] = None, + ordering: Option[Seq[SortOrder]] = None, + @transient val table: Table, + val joinKeyPositions: Option[Seq[Int]] = None, + val commonPartitionValues: Option[Seq[(InternalRow, Int)]] = None, + val reducers: Option[Seq[Option[Reducer[_, _]]]] = None, + val applyPartialClustering: Boolean = false, + val replicatePartitions: Boolean = false) + extends AbstractBatchScanExec( + output, + scan, + runtimeFilters, + ordering, + table, + keyGroupedPartitioning + ) { + + // Note: "metrics" is made transient to avoid sending driver-side metrics to tasks. + @transient override lazy val metrics: Map[String, SQLMetric] = Map() + + lazy val metadataColumns: Seq[AttributeReference] = output.collect { + case FileSourceConstantMetadataAttribute(attr) => attr + case FileSourceGeneratedMetadataAttribute(attr, _) => attr + } + + def hasUnsupportedColumns: Boolean = { + // TODO, fallback if user define same name column due to we can't right now + // detect which column is metadata column which is user defined column. + val metadataColumnsNames = metadataColumns.map(_.name) + output + .filterNot(metadataColumns.toSet) + .exists(v => metadataColumnsNames.contains(v.name)) + } + + // Spark 4.2 moved `postDriverMetrics` to SupportsCustomDriverMetrics and made the reported + // task metrics an explicit argument (see BatchScanExec in Spark 4.2). + def doPostDriverMetrics(): Unit = { + postDriverMetrics(scan.reportDriverMetrics()) + } + + override def doExecuteColumnar(): RDD[ColumnarBatch] = { + throw new UnsupportedOperationException("Need to implement this method") + } + + @transient protected lazy val filteredPartitions: Seq[Seq[InputPartition]] = { + val originalPartitioning = outputPartitioning + + val filtered = PushDownUtils.pushRuntimeFilters(scan, runtimeFilters, table, output) + // call toBatch again to get filtered partitions if any runtime filter was pushed + val newPartitions = + if (filtered) scan.toBatch.planInputPartitions().toSeq else inputPartitions + + originalPartitioning match { + case k: KeyedPartitioning => + if (newPartitions.exists(!_.isInstanceOf[HasPartitionKey])) { + throw new SparkException( + "Data source must have preserved the original partitioning " + + "during runtime filtering: not all partitions implement HasPartitionKey after " + + "filtering") + } + + if (filtered) { + // Validate that runtime filtering only removed partition keys, never introduced new ones. + val newPartitionKeys = newPartitions + .map( + partition => + InternalRowComparableWrapper( + partition.asInstanceOf[HasPartitionKey].partitionKey(), + k.expressions)) + .toSet + val oldPartitionKeys = k.partitionKeys.toSet + // We require the new number of partition keys to be equal or less than the old number. + if (oldPartitionKeys.size < newPartitionKeys.size) { + throw new SparkException( + "During runtime filtering, data source must either report " + + "the same number of partition values, or a subset of partition values from the " + + s"original. Before: ${oldPartitionKeys.size} partition values. " + + s"After: ${newPartitionKeys.size} partition values") + } + if (!newPartitionKeys.forall(oldPartitionKeys.contains)) { + throw new SparkException( + "During runtime filtering, data source must not report new " + + "partition values that are not present in the original partitioning.") + } + } + + // Group the splits that share the same partition key into a single group and sort the + // groups by partition key in ascending order. This reproduces the key-grouped layout that + // Spark 4.1's `BatchScanExec`/`KeyGroupedPartitionedScan` used to produce and that Gluten's + // planner (`SparkShims.orderPartitions`) still expects. In Spark 4.2 this grouping is + // otherwise deferred to `GroupPartitionsExec`. + newPartitions + .map(part => (part.asInstanceOf[HasPartitionKey].partitionKey(), part)) + .groupBy { case (key, _) => InternalRowComparableWrapper(key, k.expressions) } Review Comment: Thanks — I verified this independently and it's a real Spark 4.2 change: GroupPartitionsExec is new in 4.2.0 (absent in 4.1.1), so the scan now emits ungrouped per-split partitions and grouping is deferred to that operator, whereas the shim still rebuilds the 4.1 grouped layout. As you note, it's currently dormant: FileScan doesn't report KeyGroupedPartitioning (confirmed), Delta doesn't use DSv2 SPJ, and Iceberg 4.2 isn't enabled here — so the path isn't reachable or behaviorally testable in this profile-only PR. I've opened #[13139 ]( https://github.com/apache/gluten/issues/13139)to track the proper fix (preserve per-split slots + padding, let GroupPartitionsExec group) and the interim explicit fallback, together with SPJ tests over an actually-offloaded scan — to land alongside enabling native keyed DSv2 / the Spark 4.2 UT module. Deferring here to avoid shipping an unverifiable partitioning change. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
