Adopt PR apache#3349's per-partition scan logic to CometNativeScan. Add DPP.

Author: mbutrovich

common/src/main/scala/org/apache/comet/CometConf.scala

@@ -530,13 +530,6 @@ object CometConf extends ShimCometConf {
     .doubleConf
     .createWithDefault(1.0)
 
-  val COMET_DPP_FALLBACK_ENABLED: ConfigEntry[Boolean] =
-    conf("spark.comet.dppFallback.enabled")
-      .category(CATEGORY_EXEC)
-      .doc("Whether to fall back to Spark for queries that use DPP.")
-      .booleanConf
-      .createWithDefault(true)
-
   val COMET_DEBUG_ENABLED: ConfigEntry[Boolean] =
     conf("spark.comet.debug.enabled")
       .category(CATEGORY_EXEC)

native/core/src/execution/planner.rs

@@ -965,20 +965,29 @@ impl PhysicalPlanner {
                 ))
             }
             OpStruct::NativeScan(scan) => {
-                let data_schema = convert_spark_types_to_arrow_schema(scan.data_schema.as_slice());
+                let common = scan
+                    .common
+                    .as_ref()
+                    .ok_or_else(|| GeneralError("NativeScan missing common data".to_string()))?;
+
+                let data_schema =
+                    convert_spark_types_to_arrow_schema(common.data_schema.as_slice());
                 let required_schema: SchemaRef =
-                    convert_spark_types_to_arrow_schema(scan.required_schema.as_slice());
+                    convert_spark_types_to_arrow_schema(common.required_schema.as_slice());
                 let partition_schema: SchemaRef =
-                    convert_spark_types_to_arrow_schema(scan.partition_schema.as_slice());
-                let projection_vector: Vec<usize> = scan
+                    convert_spark_types_to_arrow_schema(common.partition_schema.as_slice());
+                let projection_vector: Vec<usize> = common
                     .projection_vector
                     .iter()
                     .map(|offset| *offset as usize)
                     .collect();
 
-                // Check if this partition has any files (bucketed scan with bucket pruning may have empty partitions)
-                let partition_files = &scan.file_partitions[self.partition as usize];
+                // Get this partition's files (injected at execution time)
+                let partition_files = scan.file_partition.as_ref().ok_or_else(|| {
+                    GeneralError("NativeScan missing file_partition data".to_string())
+                })?;
 
+                // Check if this partition has any files (bucketed scan with bucket pruning may have empty partitions)
                 if partition_files.partitioned_file.is_empty() {
                     let empty_exec = Arc::new(EmptyExec::new(required_schema));
                     return Ok((
@@ -988,19 +997,19 @@ impl PhysicalPlanner {
                 }
 
                 // Convert the Spark expressions to Physical expressions
-                let data_filters: Result<Vec<Arc<dyn PhysicalExpr>>, ExecutionError> = scan
+                let data_filters: Result<Vec<Arc<dyn PhysicalExpr>>, ExecutionError> = common
                     .data_filters
                     .iter()
                     .map(|expr| self.create_expr(expr, Arc::clone(&required_schema)))
                     .collect();
 
-                let default_values: Option<HashMap<usize, ScalarValue>> = if !scan
+                let default_values: Option<HashMap<usize, ScalarValue>> = if !common
                     .default_values
                     .is_empty()
                 {
                     // We have default values. Extract the two lists (same length) of values and
                     // indexes in the schema, and then create a HashMap to use in the SchemaMapper.
-                    let default_values: Result<Vec<ScalarValue>, DataFusionError> = scan
+                    let default_values: Result<Vec<ScalarValue>, DataFusionError> = common
                         .default_values
                         .iter()
                         .map(|expr| {
@@ -1015,7 +1024,7 @@ impl PhysicalPlanner {
                         })
                         .collect();
                     let default_values = default_values?;
-                    let default_values_indexes: Vec<usize> = scan
+                    let default_values_indexes: Vec<usize> = common
                         .default_values_indexes
                         .iter()
                         .map(|offset| *offset as usize)
@@ -1037,7 +1046,7 @@ impl PhysicalPlanner {
                     .map(|f| f.file_path.clone())
                     .expect("partition should have files after empty check");
 
-                let object_store_options: HashMap<String, String> = scan
+                let object_store_options: HashMap<String, String> = common
                     .object_store_options
                     .iter()
                     .map(|(k, v)| (k.clone(), v.clone()))
@@ -1048,10 +1057,7 @@ impl PhysicalPlanner {
                     &object_store_options,
                 )?;
 
-                // Comet serializes all partitions' PartitionedFiles, but we only want to read this
-                // Spark partition's PartitionedFiles
-                let files =
-                    self.get_partitioned_files(&scan.file_partitions[self.partition as usize])?;
+                let files = self.get_partitioned_files(partition_files)?;
                 let file_groups: Vec<Vec<PartitionedFile>> = vec![files];
                 let partition_fields: Vec<Field> = partition_schema
                     .fields()
@@ -1060,7 +1066,7 @@ impl PhysicalPlanner {
                         Field::new(field.name(), field.data_type().clone(), field.is_nullable())
                     })
                     .collect_vec();
-                let scan = init_datasource_exec(
+                let datasource_exec = init_datasource_exec(
                     required_schema,
                     Some(data_schema),
                     Some(partition_schema),
@@ -1070,14 +1076,14 @@ impl PhysicalPlanner {
                     Some(projection_vector),
                     Some(data_filters?),
                     default_values,
-                    scan.session_timezone.as_str(),
-                    scan.case_sensitive,
+                    common.session_timezone.as_str(),
+                    common.case_sensitive,
                     self.session_ctx(),
-                    scan.encryption_enabled,
+                    common.encryption_enabled,
                 )?;
                 Ok((
                     vec![],
-                    Arc::new(SparkPlan::new(spark_plan.plan_id, scan, vec![])),
+                    Arc::new(SparkPlan::new(spark_plan.plan_id, datasource_exec, vec![])),
                 ))
             }
             OpStruct::CsvScan(scan) => {

native/proto/src/proto/operator.proto

@@ -85,7 +85,8 @@ message Scan {
   bool arrow_ffi_safe = 3;
 }
 
-message NativeScan {
+// Common data shared across all partitions for NativeScan (sent once via commonBytes)
+message NativeScanCommon {
   repeated spark.spark_expression.DataType fields = 1;
   // The source of the scan (e.g. file scan, broadcast exchange, shuffle, etc). This
   // is purely for informational purposes when viewing native query plans in
@@ -95,20 +96,30 @@ message NativeScan {
   repeated SparkStructField data_schema = 4;
   repeated SparkStructField partition_schema = 5;
   repeated spark.spark_expression.Expr data_filters = 6;
-  repeated SparkFilePartition file_partitions = 7;
-  repeated int64 projection_vector = 8;
-  string session_timezone = 9;
-  repeated spark.spark_expression.Expr default_values = 10;
-  repeated int64 default_values_indexes = 11;
-  bool case_sensitive = 12;
+  repeated int64 projection_vector = 7;
+  string session_timezone = 8;
+  repeated spark.spark_expression.Expr default_values = 9;
+  repeated int64 default_values_indexes = 10;
+  bool case_sensitive = 11;
   // Options for configuring object stores such as AWS S3, GCS, etc. The key-value pairs are taken
   // from Hadoop configuration for compatibility with Hadoop FileSystem implementations of object
   // stores.
   // The configuration values have hadoop. or spark.hadoop. prefix trimmed. For instance, the
   // configuration value "spark.hadoop.fs.s3a.access.key" will be stored as "fs.s3a.access.key" in
   // the map.
-  map<string, string> object_store_options = 13;
-  bool encryption_enabled = 14;
+  map<string, string> object_store_options = 12;
+  bool encryption_enabled = 13;
+
+  // Unique identifier for this scan, used to match planning data at execution time
+  string scan_id = 14;
+}
+
+message NativeScan {
+  // Common data shared across partitions
+  NativeScanCommon common = 1;
+
+  // This partition's files only (injected at execution time by NativePlanDataInjector)
+  SparkFilePartition file_partition = 2;
 }
 
 message CsvScan {

spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala

@@ -180,9 +180,8 @@ case class CometExecRule(session: SparkSession) extends Rule[SparkPlan] {
   // spotless:on
   private def transform(plan: SparkPlan): SparkPlan = {
     def convertNode(op: SparkPlan): SparkPlan = op match {
-      // Fully native scan for V1
-      case scan: CometScanExec if scan.scanImpl == CometConf.SCAN_NATIVE_DATAFUSION =>
-        convertToComet(scan, CometNativeScan).getOrElse(scan)
+      // CometNativeScanExec is created directly by CometScanRule and handles its own execution
+      // No conversion needed here - it passes through unchanged
 
       // Fully native Iceberg scan for V2 (iceberg-rust path)
       // Only handle scans with native metadata; SupportsComet scans fall through to isCometScan

spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala

@@ -32,7 +32,7 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, DynamicPruningExpre
 import org.apache.spark.sql.catalyst.rules.Rule
 import org.apache.spark.sql.catalyst.util.{sideBySide, ArrayBasedMapData, GenericArrayData, MetadataColumnHelper}
 import org.apache.spark.sql.catalyst.util.ResolveDefaultColumns.getExistenceDefaultValues
-import org.apache.spark.sql.comet.{CometBatchScanExec, CometScanExec}
+import org.apache.spark.sql.comet.{CometBatchScanExec, CometNativeScanExec, CometScanExec}
 import org.apache.spark.sql.execution.{FileSourceScanExec, InSubqueryExec, SparkPlan, SubqueryAdaptiveBroadcastExec}
 import org.apache.spark.sql.execution.datasources.HadoopFsRelation
 import org.apache.spark.sql.execution.datasources.parquet.ParquetUtils
@@ -50,6 +50,7 @@ import org.apache.comet.iceberg.{CometIcebergNativeScanMetadata, IcebergReflecti
 import org.apache.comet.objectstore.NativeConfig
 import org.apache.comet.parquet.{CometParquetScan, Native, SupportsComet}
 import org.apache.comet.parquet.CometParquetUtils.{encryptionEnabled, isEncryptionConfigSupported}
+import org.apache.comet.serde.OperatorOuterClass
 import org.apache.comet.serde.operator.CometNativeScan
 import org.apache.comet.shims.{CometTypeShim, ShimFileFormat, ShimSubqueryBroadcast}
 
@@ -140,10 +141,9 @@ case class CometScanRule(session: SparkSession)
 
   private def transformV1Scan(scanExec: FileSourceScanExec): SparkPlan = {
 
-    if (COMET_DPP_FALLBACK_ENABLED.get() &&
-      scanExec.partitionFilters.exists(isDynamicPruningFilter)) {
-      return withInfo(scanExec, "Dynamic Partition Pruning is not supported")
-    }
+    // Check for DPP - only some scan implementations support it
+    val dppFilters = scanExec.partitionFilters.filter(isDynamicPruningFilter)
+    val hasDPP = dppFilters.nonEmpty
 
     scanExec.relation match {
       case r: HadoopFsRelation =>
@@ -170,13 +170,23 @@ case class CometScanRule(session: SparkSession)
         COMET_NATIVE_SCAN_IMPL.get() match {
           case SCAN_AUTO =>
             // TODO add support for native_datafusion in the future
+            if (hasDPP) {
+              return withInfo(scanExec, "Dynamic Partition Pruning is not supported")
+            }
             nativeIcebergCompatScan(session, scanExec, r, hadoopConf)
               .getOrElse(scanExec)
           case SCAN_NATIVE_DATAFUSION =>
+            // native_datafusion supports DPP
             nativeDataFusionScan(session, scanExec, r, hadoopConf).getOrElse(scanExec)
           case SCAN_NATIVE_ICEBERG_COMPAT =>
+            if (hasDPP) {
+              return withInfo(scanExec, "Dynamic Partition Pruning is not supported")
+            }
             nativeIcebergCompatScan(session, scanExec, r, hadoopConf).getOrElse(scanExec)
           case SCAN_NATIVE_COMET =>
+            if (hasDPP) {
+              return withInfo(scanExec, "Dynamic Partition Pruning is not supported")
+            }
             nativeCometScan(session, scanExec, r, hadoopConf).getOrElse(scanExec)
         }
 
@@ -213,7 +223,12 @@ case class CometScanRule(session: SparkSession)
     if (!isSchemaSupported(scanExec, SCAN_NATIVE_DATAFUSION, r)) {
       return None
     }
-    Some(CometScanExec(scanExec, session, SCAN_NATIVE_DATAFUSION))
+
+    // Create placeholder NativeScan operator
+    val builder = OperatorOuterClass.Operator.newBuilder()
+    CometNativeScan.convert(scanExec, builder).map { nativeOp =>
+      CometNativeScanExec(nativeOp, scanExec)
+    }
   }
 
   private def nativeIcebergCompatScan(