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 9535117d0d094ce0d7b381e778b0a0195bb9bacd Author: Matt Butrovich <[email protected]> AuthorDate: Fri Nov 8 13:54:27 2024 -0500 Stash changes. Maybe runs TPC-H q1? --- native/Cargo.lock | 68 ++++++++++++++++++++-- native/Cargo.toml | 4 +- native/core/Cargo.toml | 4 +- native/core/src/execution/datafusion/planner.rs | 66 +++++++++++++++++++++ native/proto/src/proto/operator.proto | 3 + .../org/apache/comet/serde/QueryPlanSerde.scala | 41 ++++++++++++- 6 files changed, 179 insertions(+), 7 deletions(-) diff --git a/native/Cargo.lock b/native/Cargo.lock index 39db7f9d..95595848 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -436,7 +436,18 @@ checksum = "d640d25bc63c50fb1f0b545ffd80207d2e10a4c965530809b40ba3386825c391" dependencies = [ "alloc-no-stdlib", "alloc-stdlib", - "brotli-decompressor", + "brotli-decompressor 2.5.1", +] + +[[package]] +name = "brotli" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index"; +checksum = "74f7971dbd9326d58187408ab83117d8ac1bb9c17b085fdacd1cf2f598719b6b" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", + "brotli-decompressor 4.0.1", ] [[package]] @@ -449,6 +460,16 @@ dependencies = [ "alloc-stdlib", ] +[[package]] +name = "brotli-decompressor" +version = "4.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index"; +checksum = "9a45bd2e4095a8b518033b128020dd4a55aab1c0a381ba4404a472630f4bc362" +dependencies = [ + "alloc-no-stdlib", + "alloc-stdlib", +] + [[package]] name = "bumpalo" version = "3.16.0" @@ -841,6 +862,7 @@ dependencies = [ "num_cpus", "object_store", "parking_lot", + "parquet", "paste", "pin-project-lite", "rand", @@ -877,7 +899,7 @@ dependencies = [ "arrow-schema", "assertables", "async-trait", - "brotli", + "brotli 3.5.0", "bytes", "crc32fast", "criterion", @@ -901,6 +923,7 @@ dependencies = [ "lz4", "mimalloc", "num", + "object_store", "once_cell", "parquet", "paste", @@ -914,7 +937,8 @@ dependencies = [ "tempfile", "thiserror", "tokio", - "zstd", + "url", + "zstd 0.11.2+zstd.1.5.2", ] [[package]] @@ -964,6 +988,7 @@ dependencies = [ "libc", "num_cpus", "object_store", + "parquet", "paste", "sqlparser", "tokio", @@ -2175,16 +2200,33 @@ source = "registry+https://github.com/rust-lang/crates.io-index"; checksum = "310c46a70a3ba90d98fec39fa2da6d9d731e544191da6fb56c9d199484d0dd3e" dependencies = [ "ahash", + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-data", + "arrow-ipc", + "arrow-schema", + "arrow-select", + "base64", + "brotli 6.0.0", "bytes", "chrono", + "flate2", + "futures", "half", "hashbrown", + "lz4_flex", "num", "num-bigint", + "object_store", "paste", "seq-macro", + "snap", "thrift", + "tokio", "twox-hash 1.6.3", + "zstd 0.13.2", + "zstd-sys", ] [[package]] @@ -3395,7 +3437,16 @@ version = "0.11.2+zstd.1.5.2" source = "registry+https://github.com/rust-lang/crates.io-index"; checksum = "20cc960326ece64f010d2d2107537f26dc589a6573a316bd5b1dba685fa5fde4" dependencies = [ - "zstd-safe", + "zstd-safe 5.0.2+zstd.1.5.2", +] + +[[package]] +name = "zstd" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index"; +checksum = "fcf2b778a664581e31e389454a7072dab1647606d44f7feea22cd5abb9c9f3f9" +dependencies = [ + "zstd-safe 7.2.1", ] [[package]] @@ -3408,6 +3459,15 @@ dependencies = [ "zstd-sys", ] +[[package]] +name = "zstd-safe" +version = "7.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index"; +checksum = "54a3ab4db68cea366acc5c897c7b4d4d1b8994a9cd6e6f841f8964566a419059" +dependencies = [ + "zstd-sys", +] + [[package]] name = "zstd-sys" version = "2.0.13+zstd.1.5.6" diff --git a/native/Cargo.toml b/native/Cargo.toml index 943d591d..f9af172e 100644 --- a/native/Cargo.toml +++ b/native/Cargo.toml @@ -40,7 +40,7 @@ arrow-data = { version = "53.1.0" } arrow-schema = { version = "53.1.0" } parquet = { version = "53.1.0", default-features = false, features = ["experimental"] } datafusion-common = { git = "https://github.com/apache/datafusion.git";, rev = "3b6aac2" } -datafusion = { git = "https://github.com/apache/datafusion.git";, rev = "3b6aac2", default-features = false, features = ["unicode_expressions", "crypto_expressions"] } +datafusion = { git = "https://github.com/apache/datafusion.git";, rev = "3b6aac2", default-features = false, features = ["unicode_expressions", "crypto_expressions", "parquet"] } datafusion-functions = { git = "https://github.com/apache/datafusion.git";, rev = "3b6aac2", features = ["crypto_expressions"] } datafusion-functions-nested = { git = "https://github.com/apache/datafusion.git";, rev = "3b6aac2", default-features = false } datafusion-expr = { git = "https://github.com/apache/datafusion.git";, rev = "3b6aac2", default-features = false } @@ -55,6 +55,8 @@ num = "0.4" rand = "0.8" regex = "1.9.6" thiserror = "1" +object_store = "0.11.0" +url = "2.2" [profile.release] debug = true diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index daa0837c..8d30b38c 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -65,7 +65,7 @@ tempfile = "3.8.0" ahash = { version = "0.8", default-features = false } itertools = "0.11.0" paste = "1.0.14" -datafusion-common = { workspace = true } +datafusion-common = { workspace = true, features= ["object_store"] } datafusion = { workspace = true } datafusion-functions-nested = { workspace = true } datafusion-expr = { workspace = true } @@ -77,6 +77,8 @@ crc32fast = "1.3.2" simd-adler32 = "0.3.7" datafusion-comet-spark-expr = { workspace = true } datafusion-comet-proto = { workspace = true } +object_store = { workspace = true } +url = { workspace = true } [dev-dependencies] pprof = { version = "0.13.0", features = ["flamegraph"] } diff --git a/native/core/src/execution/datafusion/planner.rs b/native/core/src/execution/datafusion/planner.rs index 5b53cb39..65ad73f8 100644 --- a/native/core/src/execution/datafusion/planner.rs +++ b/native/core/src/execution/datafusion/planner.rs @@ -84,6 +84,10 @@ 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; use datafusion_comet_proto::{ spark_expression::{ self, agg_expr::ExprStruct as AggExprStruct, expr::ExprStruct, literal::Value, AggExpr, @@ -105,6 +109,7 @@ use datafusion_common::{ tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRecursion, TreeNodeRewriter}, JoinType as DFJoinType, ScalarValue, }; +use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_expr::expr::find_df_window_func; use datafusion_expr::{ AggregateUDF, WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, @@ -114,8 +119,11 @@ use datafusion_physical_expr::window::WindowExpr; use itertools::Itertools; use jni::objects::GlobalRef; use num::{BigInt, ToPrimitive}; +use object_store::path::Path; +use parquet::schema::parser::parse_message_type; use std::cmp::max; use std::{collections::HashMap, sync::Arc}; +use url::Url; // For clippy error on type_complexity. type ExecResult<T> = Result<T, ExecutionError>; @@ -941,6 +949,64 @@ impl PhysicalPlanner { OpStruct::Scan(scan) => { let data_types = scan.fields.iter().map(to_arrow_datatype).collect_vec(); + if scan.source == "CometScan parquet (unknown)" { + let data_schema = parse_message_type(&*scan.data_schema).unwrap(); + let required_schema = parse_message_type(&*scan.required_schema).unwrap(); + println!("data_schema: {:?}", data_schema); + println!("required_schema: {:?}", required_schema); + + let data_schema_descriptor = + parquet::schema::types::SchemaDescriptor::new(Arc::new(data_schema)); + let data_schema_arrow = parquet::arrow::schema::parquet_to_arrow_schema( + &data_schema_descriptor, + None, + ) + .unwrap(); + println!("data_schema_arrow: {:?}", data_schema_arrow); + + let required_schema_descriptor = + parquet::schema::types::SchemaDescriptor::new(Arc::new(required_schema)); + let required_schema_arrow = parquet::arrow::schema::parquet_to_arrow_schema( + &required_schema_descriptor, + None, + ) + .unwrap(); + println!("required_schema_arrow: {:?}", required_schema_arrow); + + assert!(!required_schema_arrow.fields.is_empty()); + + let mut projection_vector: Vec<usize> = + Vec::with_capacity(required_schema_arrow.fields.len()); + // TODO: could be faster with a hashmap rather than iterating over data_schema_arrow with index_of. + required_schema_arrow.fields.iter().for_each(|field| { + projection_vector.push(data_schema_arrow.index_of(field.name()).unwrap()); + }); + println!("projection_vector: {:?}", projection_vector); + + assert_eq!(projection_vector.len(), required_schema_arrow.fields.len()); + + let object_store_url = ObjectStoreUrl::local_filesystem(); + let path = Url::parse(&*scan.path).unwrap(); + + let object_store = object_store::local::LocalFileSystem::new(); + // register the object store with the runtime environment + let url = Url::try_from("file://").unwrap(); + self.session_ctx + .runtime_env() + .register_object_store(&url, Arc::new(object_store)); + + let file = PartitionedFile::from_path(path.path().to_string())?; + + let file_scan_config = + FileScanConfig::new(object_store_url, Arc::new(data_schema_arrow)) + .with_file(file) + .with_projection(Some(projection_vector)); + + let builder = ParquetExecBuilder::new(file_scan_config); + let scan = builder.build(); + return Ok((vec![], Arc::new(scan))); + } + // If it is not test execution context for unit test, we should have at least one // input source if self.exec_context_id != TEST_EXEC_CONTEXT_ID && inputs.is_empty() { diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 533d504c..950b35ae 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -52,6 +52,9 @@ message Scan { // is purely for informational purposes when viewing native query plans in // debug mode. string source = 2; + string path = 3; + string required_schema = 4; + string data_schema = 5; } message Projection { 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 abb138b0..3775c7b2 100644 --- a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala +++ b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala @@ -29,12 +29,13 @@ import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, Normalize import org.apache.spark.sql.catalyst.plans._ import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, Partitioning, RangePartitioning, RoundRobinPartitioning, SinglePartition} import org.apache.spark.sql.catalyst.util.CharVarcharCodegenUtils -import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometSinkPlaceHolder, CometSparkToColumnarExec, DecimalPrecision} +import org.apache.spark.sql.comet.{CometBroadcastExchangeExec, CometScanExec, CometSinkPlaceHolder, CometSparkToColumnarExec, DecimalPrecision} import org.apache.spark.sql.comet.execution.shuffle.CometShuffleExchangeExec 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.SparkToParquetSchemaConverter 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 @@ -3151,6 +3152,44 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim // Sink operators don't have children result.clearChildren() + op match { + case cometScan: CometScanExec => + // scalastyle:off println +// System.out.println(op.simpleStringWithNodeId()) +// System.out.println(scanTypes.asJava) // Spark types for output. +// System.out.println(cometScan.output) // This is the names of the output columns. +// System.out.println(cometScan.requiredSchema); // This is the projected columns. +// System.out.println( +// cometScan.dataFilters +// ); // This is the filter expressions that have been pushed down. +// System.out.println(cometScan.relation.location.inputFiles(0)) +// System.out.println(cometScan.partitionFilters); +// System.out.println(cometScan.relation.partitionSchema) +// System.out.println(cometScan.metadata); + +// System.out.println("requiredSchema:") +// cometScan.requiredSchema.fields.foreach(field => { +// System.out.println(field.dataType) +// }) + +// System.out.println("relation.dataSchema:") +// cometScan.relation.dataSchema.fields.foreach(field => { +// System.out.println(field.dataType) +// }) + // scalastyle:on println + + val requiredSchemaParquet = + new SparkToParquetSchemaConverter(conf).convert(cometScan.requiredSchema) + val dataSchemaParquet = + new SparkToParquetSchemaConverter(conf).convert(cometScan.relation.dataSchema) + + scanBuilder.setRequiredSchema(requiredSchemaParquet.toString) + scanBuilder.setDataSchema(dataSchemaParquet.toString) + scanBuilder.setPath(cometScan.relation.location.inputFiles(0)) + case _ => + scanBuilder.setPath("") + } + Some(result.setScan(scanBuilder).build()) } else { // There are unsupported scan type --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
