rebase

Author: andygrove

native/core/src/execution/datafusion/planner.rs

@@ -1056,8 +1056,13 @@ impl PhysicalPlanner {
                     };
 
                 // The `ScanExec` operator will take actual arrays from Spark during execution
-                let scan =
-                    ScanExec::new(self.exec_context_id, input_source, &scan.source, data_types)?;
+                let scan = ScanExec::new(
+                    self.exec_context_id,
+                    input_source,
+                    &scan.source,
+                    data_types,
+                    scan.reuses_buffers,
+                )?;
                 Ok((
                     vec![scan.clone()],
                     Arc::new(SparkPlan::new(spark_plan.plan_id, Arc::new(scan), vec![])),
@@ -2138,8 +2143,13 @@ impl From<ExpressionError> for DataFusionError {
 fn can_reuse_input_batch(op: &Arc<dyn ExecutionPlan>) -> bool {
     if op.as_any().is::<ProjectionExec>() || op.as_any().is::<LocalLimitExec>() {
         can_reuse_input_batch(op.children()[0])
+    } else if op.as_any().is::<ScanExec>() {
+        op.as_any()
+            .downcast_ref::<ScanExec>()
+            .unwrap()
+            .reuses_buffers
     } else {
-        op.as_any().is::<ScanExec>()
+        false
     }
 }
 
@@ -2312,6 +2322,7 @@ mod tests {
                     type_info: None,
                 }],
                 source: "".to_string(),
+                reuses_buffers: false,
             })),
         };
 
@@ -2385,6 +2396,7 @@ mod tests {
                     type_info: None,
                 }],
                 source: "".to_string(),
+                reuses_buffers: false,
             })),
         };
 
@@ -2603,6 +2615,7 @@ mod tests {
             op_struct: Some(OpStruct::Scan(spark_operator::Scan {
                 fields: vec![create_proto_datatype()],
                 source: "".to_string(),
+                reuses_buffers: false,
             })),
         }
     }

native/core/src/execution/operators/scan.rs

@@ -64,6 +64,8 @@ pub struct ScanExec {
     pub input_source: Option<Arc<GlobalRef>>,
     /// A description of the input source for informational purposes
     pub input_source_description: String,
+    /// Some sources (currently only the native Parquet scan) re-use mutable buffers
+    pub reuses_buffers: bool,
     /// The data types of columns of the input batch. Converted from Spark schema.
     pub data_types: Vec<DataType>,
     /// Schema of first batch
@@ -85,6 +87,7 @@ impl ScanExec {
         input_source: Option<Arc<GlobalRef>>,
         input_source_description: &str,
         data_types: Vec<DataType>,
+        reuses_buffers: bool,
     ) -> Result<Self, CometError> {
         let metrics_set = ExecutionPlanMetricsSet::default();
         let baseline_metrics = BaselineMetrics::new(&metrics_set, 0);
@@ -119,6 +122,7 @@ impl ScanExec {
             exec_context_id,
             input_source,
             input_source_description: input_source_description.to_string(),
+            reuses_buffers,
             data_types,
             batch: Arc::new(Mutex::new(Some(first_batch))),
             cache,

native/proto/src/proto/operator.proto

@@ -55,6 +55,7 @@ message Scan {
   // is purely for informational purposes when viewing native query plans in
   // debug mode.
   string source = 2;
+  bool reuses_buffers = 3;
 }
 
 message Projection {

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -35,6 +35,9 @@ import org.apache.spark.sql.execution
 import org.apache.spark.sql.execution._
 import org.apache.spark.sql.execution.adaptive.{BroadcastQueryStageExec, ShuffleQueryStageExec}
 import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, HashAggregateExec, ObjectHashAggregateExec}
+import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat
+import org.apache.spark.sql.execution.datasources.v2.BatchScanExec
+import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan
 import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, ReusedExchangeExec, ShuffleExchangeExec}
 import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, HashJoin, ShuffledHashJoinExec, SortMergeJoinExec}
 import org.apache.spark.sql.execution.window.WindowExec
@@ -2910,6 +2913,7 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
         // These operators are source of Comet native execution chain
         val scanBuilder = OperatorOuterClass.Scan.newBuilder()
         scanBuilder.setSource(op.simpleStringWithNodeId())
+        scanBuilder.setReusesBuffers(isParquetScan(op))
 
         val scanTypes = op.output.flatten { attr =>
           serializeDataType(attr.dataType)
@@ -2944,6 +2948,17 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
     }
   }
 
+  def isParquetScan(op: SparkPlan): Boolean = {
+    op match {
+      case scan: FileSourceScanExec =>
+        scan.relation.fileFormat.isInstanceOf[ParquetFileFormat]
+      case scan: BatchScanExec =>
+        scan.scan.isInstanceOf[ParquetScan]
+      case _ =>
+        false
+    }
+  }
+
   /**
    * Whether the input Spark operator `op` can be considered as a Comet sink, i.e., the start of
    * native execution. If it is true, we'll wrap `op` with `CometScanWrapper` or