This is an automated email from the ASF dual-hosted git repository. agrove pushed a commit to branch comet-parquet-exec in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
commit 08308407b19861fedf40ba5175ecef7d6d0dcedf Author: Andy Grove <[email protected]> AuthorDate: Sat Nov 9 09:28:56 2024 -0700 partitioning --- native/core/src/execution/datafusion/planner.rs | 34 +++++++++++++++------- native/core/src/execution/jni_api.rs | 11 +++++-- .../scala/org/apache/comet/CometExecIterator.scala | 4 ++- spark/src/main/scala/org/apache/comet/Native.scala | 7 ++++- .../org/apache/comet/serde/QueryPlanSerde.scala | 3 -- 5 files changed, 41 insertions(+), 18 deletions(-) diff --git a/native/core/src/execution/datafusion/planner.rs b/native/core/src/execution/datafusion/planner.rs index 7d8bbbb7..2c7b36ef 100644 --- a/native/core/src/execution/datafusion/planner.rs +++ b/native/core/src/execution/datafusion/planner.rs @@ -825,12 +825,13 @@ impl PhysicalPlanner { &'a self, spark_plan: &'a Operator, inputs: &mut Vec<Arc<GlobalRef>>, + partition_count: usize, ) -> Result<(Vec<ScanExec>, Arc<dyn ExecutionPlan>), ExecutionError> { let children = &spark_plan.children; match spark_plan.op_struct.as_ref().unwrap() { OpStruct::Projection(project) => { assert!(children.len() == 1); - let (scans, child) = self.create_plan(&children[0], inputs)?; + let (scans, child) = self.create_plan(&children[0], inputs, partition_count)?; let exprs: PhyExprResult = project .project_list .iter() @@ -844,7 +845,7 @@ impl PhysicalPlanner { } OpStruct::Filter(filter) => { assert!(children.len() == 1); - let (scans, child) = self.create_plan(&children[0], inputs)?; + let (scans, child) = self.create_plan(&children[0], inputs, partition_count)?; let predicate = self.create_expr(filter.predicate.as_ref().unwrap(), child.schema())?; @@ -852,7 +853,7 @@ impl PhysicalPlanner { } OpStruct::HashAgg(agg) => { assert!(children.len() == 1); - let (scans, child) = self.create_plan(&children[0], inputs)?; + let (scans, child) = self.create_plan(&children[0], inputs, partition_count)?; let group_exprs: PhyExprResult = agg .grouping_exprs @@ -917,13 +918,13 @@ impl PhysicalPlanner { OpStruct::Limit(limit) => { assert!(children.len() == 1); let num = limit.limit; - let (scans, child) = self.create_plan(&children[0], inputs)?; + let (scans, child) = self.create_plan(&children[0], inputs, partition_count)?; Ok((scans, Arc::new(LocalLimitExec::new(child, num as usize)))) } OpStruct::Sort(sort) => { assert!(children.len() == 1); - let (scans, child) = self.create_plan(&children[0], inputs)?; + let (scans, child) = self.create_plan(&children[0], inputs, partition_count)?; let exprs: Result<Vec<PhysicalSortExpr>, ExecutionError> = sort .sort_orders @@ -1029,9 +1030,17 @@ impl PhysicalPlanner { .map(|path| PartitionedFile::from_path(path.path().to_string()).unwrap()) .collect(); + // partition the files + // TODO really should partition the row groups + + let mut file_groups = vec![vec![]; partition_count]; + files.iter().enumerate().for_each(|(idx, file)| { + file_groups[partition_count % idx].push(file.clone()); + }); + let file_scan_config = FileScanConfig::new(object_store_url, data_schema_arrow.clone()) - .with_file_groups(vec![files]) + .with_file_groups(file_groups) .with_projection(Some(projection_vector)); let builder = @@ -1064,7 +1073,7 @@ impl PhysicalPlanner { } OpStruct::ShuffleWriter(writer) => { assert!(children.len() == 1); - let (scans, child) = self.create_plan(&children[0], inputs)?; + let (scans, child) = self.create_plan(&children[0], inputs, partition_count)?; let partitioning = self .create_partitioning(writer.partitioning.as_ref().unwrap(), child.schema())?; @@ -1081,7 +1090,7 @@ impl PhysicalPlanner { } OpStruct::Expand(expand) => { assert!(children.len() == 1); - let (scans, child) = self.create_plan(&children[0], inputs)?; + let (scans, child) = self.create_plan(&children[0], inputs, partition_count)?; let mut projections = vec![]; let mut projection = vec![]; @@ -1139,6 +1148,7 @@ impl PhysicalPlanner { &join.right_join_keys, join.join_type, &join.condition, + partition_count, )?; let sort_options = join @@ -1177,6 +1187,7 @@ impl PhysicalPlanner { &join.right_join_keys, join.join_type, &join.condition, + partition_count, )?; // HashJoinExec may cache the input batch internally. We need @@ -1209,7 +1220,7 @@ impl PhysicalPlanner { Ok((scans, hash_join)) } OpStruct::Window(wnd) => { - let (scans, child) = self.create_plan(&children[0], inputs)?; + let (scans, child) = self.create_plan(&children[0], inputs, partition_count)?; let input_schema = child.schema(); let sort_exprs: Result<Vec<PhysicalSortExpr>, ExecutionError> = wnd .order_by_list @@ -1261,10 +1272,11 @@ impl PhysicalPlanner { right_join_keys: &[Expr], join_type: i32, condition: &Option<Expr>, + partition_count: usize, ) -> Result<(JoinParameters, Vec<ScanExec>), ExecutionError> { assert!(children.len() == 2); - let (mut left_scans, left) = self.create_plan(&children[0], inputs)?; - let (mut right_scans, right) = self.create_plan(&children[1], inputs)?; + let (mut left_scans, left) = self.create_plan(&children[0], inputs, partition_count)?; + let (mut right_scans, right) = self.create_plan(&children[1], inputs, partition_count)?; left_scans.append(&mut right_scans); diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index f1ddba44..83f9570e 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -67,6 +67,8 @@ struct ExecutionContext { pub id: i64, /// The deserialized Spark plan pub spark_plan: Operator, + /// The number of partitions + pub partition_count: usize, /// The DataFusion root operator converted from the `spark_plan` pub root_op: Option<Arc<dyn ExecutionPlan>>, /// The input sources for the DataFusion plan @@ -102,6 +104,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( config_object: JObject, iterators: jobjectArray, serialized_query: jbyteArray, + partition_count: jint, metrics_node: JObject, comet_task_memory_manager_obj: JObject, ) -> jlong { @@ -170,6 +173,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( let exec_context = Box::new(ExecutionContext { id, spark_plan, + partition_count: partition_count as usize, root_op: None, scans: vec![], input_sources, @@ -340,6 +344,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( e: JNIEnv, _class: JClass, exec_context: jlong, + partition_id: jint, array_addrs: jlongArray, schema_addrs: jlongArray, ) -> jlong { @@ -358,6 +363,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( let (scans, root_op) = planner.create_plan( &exec_context.spark_plan, &mut exec_context.input_sources.clone(), + exec_context.partition_count, )?; exec_context.root_op = Some(Arc::clone(&root_op)); @@ -374,11 +380,12 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( let plan = exec_context.root_op.as_ref().unwrap(); println!( - "Executing partition 0 of {}", + "Executing partition {} of {}", + partition_id, plan.output_partitioning().partition_count() ); - let stream = plan.execute(0, task_ctx)?; + let stream = plan.execute(partition_id as usize, task_ctx)?; exec_context.stream = Some(stream); } else { // Pull input batches diff --git a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala index 07dd80c3..2b6bccdb 100644 --- a/spark/src/main/scala/org/apache/comet/CometExecIterator.scala +++ b/spark/src/main/scala/org/apache/comet/CometExecIterator.scala @@ -55,11 +55,13 @@ class CometExecIterator( }.toArray private val plan = { val configs = createNativeConf + TaskContext.get().numPartitions() nativeLib.createPlan( id, configs, cometBatchIterators, protobufQueryPlan, + TaskContext.get().numPartitions(), nativeMetrics, new CometTaskMemoryManager(id)) } @@ -104,7 +106,7 @@ class CometExecIterator( nativeUtil.getNextBatch( numOutputCols, (arrayAddrs, schemaAddrs) => { - nativeLib.executePlan(plan, arrayAddrs, schemaAddrs) + nativeLib.executePlan(plan, TaskContext.get().partitionId(), arrayAddrs, schemaAddrs) }) } diff --git a/spark/src/main/scala/org/apache/comet/Native.scala b/spark/src/main/scala/org/apache/comet/Native.scala index 03a9dea0..ce0e2612 100644 --- a/spark/src/main/scala/org/apache/comet/Native.scala +++ b/spark/src/main/scala/org/apache/comet/Native.scala @@ -50,6 +50,7 @@ class Native extends NativeBase { configMap: Map[String, String], iterators: Array[CometBatchIterator], plan: Array[Byte], + partitionCount: Int, metrics: CometMetricNode, taskMemoryManager: CometTaskMemoryManager): Long @@ -65,7 +66,11 @@ class Native extends NativeBase { * @return * the number of rows, if -1, it means end of the output. */ - @native def executePlan(plan: Long, arrayAddrs: Array[Long], schemaAddrs: Array[Long]): Long + @native def executePlan( + plan: Long, + partitionId: Int, + arrayAddrs: Array[Long], + schemaAddrs: Array[Long]): Long /** * Release and drop the native query plan object and context object. diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala index 819a7c91..ec5302d3 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -3189,13 +3189,10 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim scanBuilder.setRequiredSchema(requiredSchemaParquet.toString) scanBuilder.setDataSchema(dataSchemaParquet.toString) - // TODO this is not correct yet .. we really need to know how many partitions Spark is expecting - // and then split the files / row groups up across that number of partitions cometScan.relation.location.inputFiles.foreach { f => scanBuilder.addPath(f) } case _ => - // scanBuilder.setPath("") } Some(result.setScan(scanBuilder).build()) --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
