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 0b0d6e86b110c428a2b743baedff93635f4a02a8 Author: Andy Grove <[email protected]> AuthorDate: Sat Nov 9 08:09:02 2024 -0700 add partial support for multiple parquet files --- native/core/src/execution/datafusion/planner.rs | 14 ++++++++++---- native/core/src/execution/jni_api.rs | 18 +++++++++++------- native/proto/src/proto/operator.proto | 2 +- .../scala/org/apache/comet/serde/QueryPlanSerde.scala | 9 +++++++-- 4 files changed, 29 insertions(+), 14 deletions(-) diff --git a/native/core/src/execution/datafusion/planner.rs b/native/core/src/execution/datafusion/planner.rs index aa858e44..7d8bbbb7 100644 --- a/native/core/src/execution/datafusion/planner.rs +++ b/native/core/src/execution/datafusion/planner.rs @@ -84,7 +84,6 @@ use datafusion::{ }; use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr}; -use datafusion::datasource::file_format::FileFormat; use datafusion::datasource::listing::PartitionedFile; use datafusion::datasource::physical_plan::parquet::ParquetExecBuilder; use datafusion::datasource::physical_plan::FileScanConfig; @@ -1012,7 +1011,11 @@ impl PhysicalPlanner { println!("test_data_filters: {:?}", test_data_filters); let object_store_url = ObjectStoreUrl::local_filesystem(); - let path = Url::parse(&*scan.path).unwrap(); + let paths: Vec<Url> = scan + .path + .iter() + .map(|path| Url::parse(path).unwrap()) + .collect(); let object_store = object_store::local::LocalFileSystem::new(); // register the object store with the runtime environment @@ -1021,11 +1024,14 @@ impl PhysicalPlanner { .runtime_env() .register_object_store(&url, Arc::new(object_store)); - let file = PartitionedFile::from_path(path.path().to_string())?; + let files: Vec<PartitionedFile> = paths + .iter() + .map(|path| PartitionedFile::from_path(path.path().to_string()).unwrap()) + .collect(); let file_scan_config = FileScanConfig::new(object_store_url, data_schema_arrow.clone()) - .with_file(file) + .with_file_groups(vec![files]) .with_projection(Some(projection_vector)); let builder = diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index d7e8ccab..f1ddba44 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -17,8 +17,10 @@ //! Define JNI APIs which can be called from Java/Scala. +use super::{serde, utils::SparkArrowConvert, CometMemoryPool}; use arrow::datatypes::DataType as ArrowDataType; use arrow_array::RecordBatch; +use datafusion::physical_plan::ExecutionPlanProperties; use datafusion::{ execution::{ disk_manager::DiskManagerConfig, @@ -39,8 +41,6 @@ use jni::{ }; use std::{collections::HashMap, sync::Arc, task::Poll}; -use super::{serde, utils::SparkArrowConvert, CometMemoryPool}; - use crate::{ errors::{try_unwrap_or_throw, CometError, CometResult}, execution::{ @@ -370,11 +370,15 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( } let task_ctx = exec_context.session_ctx.task_ctx(); - let stream = exec_context - .root_op - .as_ref() - .unwrap() - .execute(0, task_ctx)?; + + let plan = exec_context.root_op.as_ref().unwrap(); + + println!( + "Executing partition 0 of {}", + plan.output_partitioning().partition_count() + ); + + let stream = plan.execute(0, task_ctx)?; exec_context.stream = Some(stream); } else { // Pull input batches diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 67070ee6..afd5d195 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -52,7 +52,7 @@ message Scan { // is purely for informational purposes when viewing native query plans in // debug mode. string source = 2; - string path = 3; + repeated string path = 3; string required_schema = 4; string data_schema = 5; repeated spark.spark_expression.Expr data_filters = 6; 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 e5929261..819a7c91 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -3188,9 +3188,14 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim scanBuilder.setRequiredSchema(requiredSchemaParquet.toString) scanBuilder.setDataSchema(dataSchemaParquet.toString) - scanBuilder.setPath(cometScan.relation.location.inputFiles(0)) + + // 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("") + // scanBuilder.setPath("") } Some(result.setScan(scanBuilder).build()) --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
