diff --git a/dev/ci/check-suites.py b/dev/ci/check-suites.py index 68bf972c5cf..5ced2ba6971 100644 --- a/dev/ci/check-suites.py +++ b/dev/ci/check-suites.py @@ -42,6 +42,7 @@ def file_to_class_name(path: Path) -> str | None: "org.apache.comet.CometIcebergSchedulerFailureSuite", # manual, multiple hosts + SIGKILL "org.apache.comet.cloud.s3.CometS3CredentialBridgeSuite", # manual test suite "org.apache.comet.cloud.s3.HadoopS3ACredentialProviderAdapterBridgeSuite", # manual test suite + "org.apache.comet.exec.CometNativeExecutorsSuite", # manual, needs comet-scheduler and comet-executor binaries "org.apache.comet.shuffle.CelebornReflectionCompatibilitySuite", # dedicated version matrix "org.apache.spark.sql.comet.CometPlanStabilitySuite", # abstract "org.apache.spark.sql.comet.ParquetDatetimeRebaseSuite", # abstract diff --git a/docs/source/contributor-guide/index.md b/docs/source/contributor-guide/index.md index 53acfa1be69..350d61f7931 100644 --- a/docs/source/contributor-guide/index.md +++ b/docs/source/contributor-guide/index.md @@ -51,6 +51,7 @@ Floating-Point Semantics JVM Shuffle Native Shuffle Memory Management +Comet Native Executors Iceberg Writes ANSI Error Propagation S3 Credential Provider Design diff --git a/docs/source/contributor-guide/native_executors.md b/docs/source/contributor-guide/native_executors.md new file mode 100644 index 00000000000..d81be8fb0fd --- /dev/null +++ b/docs/source/contributor-guide/native_executors.md @@ -0,0 +1,165 @@ + + +# Comet Native Executors + +**Experimental.** Comet Native Executors run whole Spark queries without a JVM in the data plane. A +Spark driver plans each query with Catalyst and Comet, then submits the plan as a stage graph of +Comet native blocks to a cluster of one `comet-scheduler` and any number of `comet-executor` +processes, built on [Apache DataFusion Ballista](https://github.com/apache/datafusion-ballista). +The design is in [issue #7288](https://github.com/apache/datafusion-comet/issues/7288), and +`native/native-executors/README.md` describes the cluster, how to build it and how to test it. + +## Current limitation: whole queries only + +The driver offloads a query only if every operator and expression in it can run natively +(`CometNativeExecutors` in `CometNativeExecutorsExec.scala`). Otherwise it refuses the query, +naming the operator that couldn't run and Comet's fallback reasons for it. That includes plans that +call back into the JVM, such as JVM scalar UDFs and Arrow Python UDFs, because the cluster has no +JVM. + +Data already crosses between Spark and the cluster in two places: + +- **Scalar subqueries.** Spark runs each scalar subquery, and the driver writes its result into the + plan as a literal before submitting the query. Data flows from Spark to the cluster, but only one + value. +- **Query results.** The driver reads the final stage's output over Arrow Flight, one endpoint per + result partition (`ResultEndpoint`). Data flows from the cluster to Spark, but only at the root + of the plan. + +## Possible future research: hybrid Spark and native execution + +A query could run partly on the cluster and partly in Spark, so that a query with one unsupported +operator or UDF can still run most of its work natively. The approaches below generalize the two +crossings above. They differ mainly in which way data crosses the boundary. Moving data from the +cluster into Spark is much easier than moving data from Spark into the cluster. + +### A native subtree as a leaf of the Spark plan + +The driver could offload the largest native subtrees under an operator that must run in Spark, +rather than requiring the whole plan to be native. Spark would see each offloaded subtree as a leaf +with one partition per result endpoint, and each Spark task would read its partition from the +cluster over Flight. The batches arrive as Arrow, so Comet's columnar operators above the leaf can +consume them without converting them to rows. + +This reuses the result path that already works. What remains: + +- Serving partitioned results in pipelined mode. Today the scheduler coalesces a pipelined job's + results to one partition. +- Reporting the cluster's output statistics to Spark's adaptive query execution. +- Choosing where the native region should end. + +It suits the common shape where the scans, joins and partial aggregates are native, and an +unsupported operator or UDF sits near the top of the plan. + +### Small Spark outputs as inputs to native stages + +Spark could compute a whole relation that fits in a broadcast, such as the build side of a +broadcast join, rather than a single value. The driver would encode the relation as Arrow IPC and +send it with the plan or upload it to the scheduler, where a native leaf reads it. This only +works for small data, but it costs little, because it uses the same round trip through the driver +that scalar subqueries already make. + +### A shuffle both engines can read and write + +Mixing stages freely needs a shuffle that Spark tasks can write and cluster tasks can read, and the +other way round. Some options: + +- **A remote shuffle service.** Comet already integrates with Apache Celeborn. A Celeborn shuffle + reader and writer for Ballista would let either engine run any stage. +- **An object store or shared filesystem.** Comet's native shuffle already writes Arrow IPC. Spark + stages would write it to a location the cluster can read, and the scheduler would register those + files as a finished stage's output. +- **Pushing batches over Flight.** Spark tasks would send batches to executors with `DoPut`. This is + simpler to deploy, but it ties the lifetimes of Spark tasks and cluster tasks together and makes + retrying a failed task difficult. + +The hard part is scheduling rather than the data format. Spark's DAG scheduler and Ballista's +scheduler would have to agree on stage dependencies, failure handling and adaptive replanning. +Realistically, Spark stays the coordinator, and the cluster runs individual stages on its behalf. + +### Running JVM expressions on native executors + +Instead of splitting the plan, `comet-executor` could embed a JVM to evaluate the expressions and +UDFs that have no native implementation, reusing Comet's existing bridges for calling back into +the JVM. The plan would stay fully native and scheduling would not change. The costs are that every +executor needs Spark's classpath and enough session state to evaluate the expressions, and that the +plan must carry those expressions in serialized Spark form. + +### Speaking Spark's protocols + +Another approach is for `comet-executor` to speak Spark's own protocols, so that Spark tasks and +native tasks exchange data the way Spark executors do. A Spark executor uses two protocols, and +they differ in how practical they are to implement. + +The control plane runs over Spark's Netty RPC. Its messages, such as `RegisterExecutor`, +`LaunchTask`, `StatusUpdate` and heartbeats, are Scala case classes in Java serialization. A +`LaunchTask` carries a serialized `ShuffleMapTask` or `ResultTask`, whose RDD chain and closures +are JVM code, and the formats change between Spark versions. + +The data plane is Spark's block transfer protocol in `network-common` (`OpenBlocks`, +`FetchShuffleBlocks`, `StreamHandle`, `ChunkFetchRequest` and so on). Its messages have a compact +binary encoding rather than Java serialization, the protocol has been stable for many releases, and +the external shuffle service speaks it too. Shuffle data sits in Spark's sort-based layout of a data +file and an index file per map task. Comet's native shuffle already writes that layout: the data +comes from Rust and `CometNativeShuffleWriter` commits the index through Spark's +`IndexShuffleBlockResolver`. + +That suggests three options: + +- **A fully native Spark executor.** Not practical. Running a Spark task means deserializing Java + objects and running JVM closures. Even a task that only wraps a Comet plan would need its plan + extracted from a Java-serialized object graph, separately for each Spark version. +- **A native data plane, with Spark coordinating.** `comet-executor` writes shuffle output in + Spark's layout, serves those blocks over the block transfer protocol, and fetches blocks from + Spark executors and the external shuffle service. When a native stage finishes, the driver + registers its map outputs with Spark's `MapOutputTrackerMaster`, and downstream Spark stages read + them as they would any other shuffle. In the other direction, the driver gives a native stage the + block locations of the Spark shuffle it reads. This provides a shuffle both engines can read and + write without a separate shuffle service. The work includes Spark's authentication and encryption + (SASL, or `spark.network.crypto`), which a first version could require to be off, mapping fetch + failures onto Spark's `FetchFailed` handling and stage retries, and possibly push-based shuffle + later. A native block fetch client would also benefit Comet on Spark, by moving shuffle reads out + of the JVM. +- **An embedded JVM for the control plane.** `comet-executor` hosts a small JVM that runs a cut-down + `CoarseGrainedExecutorBackend`. It registers with the driver, deserializes tasks, passes each + task's Comet plan to native code, and reports map statuses and accumulators. This works, but the + result is a Spark executor with a small heap whose tasks run natively: Comet on Spark with a more + native data plane. Spark schedules every task one stage at a time, so the cluster's pipelined + execution and lower per-task overhead are lost. An embedded JVM is more useful as an evaluator for + expressions that have no native implementation, as described above, where Spark's protocols are + not involved. + +The data-plane option combines with either way of placing native work. Native executors could join +a Spark application as executors, which needs the embedded JVM for their control plane, with +stage-level scheduling sending fully native stages to them through a `ResourceProfile` that asks for +a resource only native executors offer. Or the driver could keep offloading whole native subgraphs +to the cluster, which keeps its own scheduling and pipelining, and use Spark-compatible shuffle as +the boundary between Spark stages and native stages. The second keeps more of what the cluster +provides. + +A first experiment ([#7384](https://github.com/apache/datafusion-comet/issues/7384)) is a Rust client and server for `OpenBlocks` and `FetchShuffleBlocks`, tested +against a Spark executor and the external shuffle service by reading a shuffle that Comet's native +shuffle wrote. That shows whether the data plane works independently of any scheduling decision. + +### When splitting pays off + +Whichever approach is used, the driver should decide whether to split a query based on cost. If +most of a query's work has to run in Spark, running all of it with Comet on Spark is likely to beat +sending part of it to the cluster and reading the results back over the network. diff --git a/native/Cargo.lock b/native/Cargo.lock index 673325aedfe..670961ad86e 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -347,6 +347,27 @@ dependencies = [ "num-traits", ] +[[package]] +name = "arrow-flight" +version = "59.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b543210778ac15d5f7a10153dc39e099f1eee756363589a842029e6c5c76e423" +dependencies = [ + "arrow-array", + "arrow-buffer", + "arrow-cast", + "arrow-data", + "arrow-ipc", + "arrow-schema", + "base64 0.23.1", + "bytes", + "futures", + "prost", + "prost-types", + "tonic", + "tonic-prost", +] + [[package]] name = "arrow-ipc" version = "59.3.1" @@ -1033,6 +1054,49 @@ dependencies = [ "tracing", ] +[[package]] +name = "axum" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90" +dependencies = [ + "axum-core", + "bytes", + "futures-util", + "http", + "http-body", + "http-body-util", + "itoa", + "matchit", + "memchr", + "mime", + "percent-encoding", + "pin-project-lite", + "serde_core", + "sync_wrapper", + "tower", + "tower-layer", + "tower-service", +] + +[[package]] +name = "axum-core" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "http-body-util", + "mime", + "pin-project-lite", + "sync_wrapper", + "tower-layer", + "tower-service", +] + [[package]] name = "backon" version = "1.6.0" @@ -1979,6 +2043,7 @@ name = "datafusion-comet" version = "1.3.0" dependencies = [ "arrow", + "arrow-flight", "assertables", "async-trait", "aws-config", @@ -2044,6 +2109,7 @@ dependencies = [ "tikv-jemalloc-ctl", "tikv-jemallocator", "tokio", + "tonic", "typetag", "url", "uuid", @@ -3465,6 +3531,12 @@ version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" +[[package]] +name = "httpdate" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" + [[package]] name = "humantime" version = "2.4.0" @@ -3494,6 +3566,7 @@ dependencies = [ "http", "http-body", "httparse", + "httpdate", "itoa", "pin-project-lite", "smallvec", @@ -3517,6 +3590,19 @@ dependencies = [ "tower-service", ] +[[package]] +name = "hyper-timeout" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" +dependencies = [ + "hyper", + "hyper-util", + "pin-project-lite", + "tokio", + "tower-service", +] + [[package]] name = "hyper-util" version = "0.1.21" @@ -4241,6 +4327,12 @@ dependencies = [ "twox-hash", ] +[[package]] +name = "matchit" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" + [[package]] name = "md-5" version = "0.10.6" @@ -4295,6 +4387,12 @@ dependencies = [ "libmimalloc-sys", ] +[[package]] +name = "mime" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" + [[package]] name = "miniz_oxide" version = "0.8.9" @@ -6883,6 +6981,47 @@ dependencies = [ "tokio", ] +[[package]] +name = "tonic" +version = "0.14.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd1fa9b1a51188f3266a823208078315811269d06d272c3a61be0b9a194194ff" +dependencies = [ + "async-trait", + "axum", + "base64 0.22.1", + "bytes", + "h2", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-timeout", + "hyper-util", + "percent-encoding", + "pin-project", + "socket2", + "sync_wrapper", + "tokio", + "tokio-stream", + "tokio-util", + "tower", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tonic-prost" +version = "0.14.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f63c616ffaa648f7c5dd5657e97d9a285fcb44ff3b70da761d436f5a13ff7be6" +dependencies = [ + "bytes", + "prost", + "tonic", +] + [[package]] name = "tower" version = "0.5.3" @@ -6891,11 +7030,15 @@ checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" dependencies = [ "futures-core", "futures-util", + "indexmap 2.14.2", "pin-project-lite", + "slab", "sync_wrapper", "tokio", + "tokio-util", "tower-layer", "tower-service", + "tracing", ] [[package]] diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index 254e0740e97..c84c09492a2 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -54,6 +54,9 @@ async-trait = { workspace = true } log = "0.4" log4rs = "1.4.0" prost = "0.14.3" +# The driver's client of a Comet Native Executor cluster; see the native-executors feature. +arrow-flight = { version = "59.2.0", optional = true } +tonic = { version = "0.14", optional = true } jni = { workspace = true } rand = { workspace = true } num = { workspace = true } @@ -132,6 +135,9 @@ python-udf = ["dep:pyo3", "dep:libc"] contrib-lance = ["dep:comet-contrib-lance"] hdfs-opendal = ["opendal/services-hdfs", "object_store_opendal", "hdfs-sys"] jemalloc = ["tikv-jemallocator", "tikv-jemalloc-ctl"] +# The JNI client a Spark driver uses to run queries on Comet Native Executors (issue #7288). Off +# by default while the feature is experimental, so the standard build links no gRPC client. +native-executors = ["dep:arrow-flight", "dep:tonic"] # Native Delta Lake scan support for the JVM-planned path (contrib/delta-spark). # In the default set so trying the contrib needs only the jar and the config, # not a custom native build. It is inert at runtime unless the contrib jar is diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index d915abfed0b..de5e539625d 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -1034,6 +1034,8 @@ fn configure_skip_partial_aggregation(config: &mut SessionConfig, plan: &Operato && agg.agg_exprs.iter().all(|expr| { matches!(&expr.expr_struct, Some(AggExprStruct::Count(count)) if count.children.len() == 1) + || matches!(&expr.expr_struct, Some(AggExprStruct::Avg(avg)) + if avg.child.is_some()) }) } // PartialMerge consumes states, not rows. @@ -1057,6 +1059,9 @@ fn configure_skip_partial_aggregation(config: &mut SessionConfig, plan: &Operato } } +/// Set in the configuration of a Comet Native Executors fragment's session. +pub(crate) const NATIVE_EXECUTORS_FRAGMENT: &str = "spark.comet.nativeExecutors.fragment"; + /// Configure DataFusion session context. fn prepare_datafusion_session_context( batch_size: usize, @@ -1151,9 +1156,25 @@ pub(crate) fn comet_session_context( crate::execution::c_udf::policy::NativeUdfPolicy::from_spark_config(spark_config), )); + // A Comet Native Executors fragment has no ShuffleWriter root, but its output goes to the + // cluster's shuffle all the same, so it is checked as if it had one. + let fragment_root; + let gate_plan = if spark_config + .get(NATIVE_EXECUTORS_FRAGMENT) + .is_some_and(|value| value == "true") + { + fragment_root = Operator { + op_struct: Some(OpStruct::ShuffleWriter(Default::default())), + children: vec![spark_plan.clone()], + ..Default::default() + }; + &fragment_root + } else { + spark_plan + }; configure_skip_partial_aggregation( &mut session_config, - spark_plan, + gate_plan, spark_config.get_bool(COMET_EXEC_AGGREGATE_SKIP_PARTIAL_ENABLED), ); @@ -1233,7 +1254,7 @@ fn prepare_output( /// # Safety /// Each address must point to a writable `FFI_ArrowArray` / `FFI_ArrowSchema` the caller owns. /// Whatever the structs hold is overwritten without being released. -unsafe fn export_batch( +pub(crate) unsafe fn export_batch( array_addrs: &[i64], schema_addrs: &[i64], output_batch: RecordBatch, @@ -2712,7 +2733,7 @@ mod tests { use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::ExecutionPlan; use datafusion_comet_proto::spark_expression; - use datafusion_comet_proto::spark_expression::{AggExpr, Count, Expr, Sum}; + use datafusion_comet_proto::spark_expression::{AggExpr, Avg, Count, Expr, Sum}; use datafusion_comet_proto::spark_operator::{HashAggregate, ShuffleWriter}; use std::cell::Cell; use std::future::Future; @@ -2729,6 +2750,13 @@ mod tests { expr_struct: Some(AggExprStruct::Sum(Sum::default())), ..Default::default() }; + let avg = |child: Option| AggExpr { + expr_struct: Some(AggExprStruct::Avg(Avg { + child, + ..Default::default() + })), + ..Default::default() + }; let partial = HashAggregate { grouping_exprs: vec![Expr::default()], agg_exprs: vec![count.clone()], @@ -2773,6 +2801,10 @@ mod tests { agg_exprs: vec![], ..partial.clone() }, + HashAggregate { + agg_exprs: vec![count.clone(), avg(Some(Expr::default()))], + ..partial.clone() + }, HashAggregate { agg_exprs: vec![count.clone(), count], ..partial.clone() @@ -2797,6 +2829,10 @@ mod tests { agg_exprs: vec![AggExpr::default()], ..partial.clone() }, + HashAggregate { + agg_exprs: vec![avg(None)], + ..partial.clone() + }, HashAggregate { agg_exprs: vec![AggExpr { expr_struct: Some(AggExprStruct::Count(Count { diff --git a/native/core/src/execution/native_executors/client.rs b/native/core/src/execution/native_executors/client.rs new file mode 100644 index 00000000000..0eea7fb60ae --- /dev/null +++ b/native/core/src/execution/native_executors/client.rs @@ -0,0 +1,185 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The JNI client a Spark driver uses to run queries on a Comet Native Executor cluster: the +//! native methods of `org.apache.comet.NativeExecutorsClient`. +//! +//! The driver submits a query to the scheduler's Arrow Flight service, which answers with one +//! endpoint per result partition, and then reads each partition from its executor with `DoGet`. + +use arrow_flight::decode::FlightRecordBatchStream; +use arrow_flight::flight_service_client::FlightServiceClient; +use arrow_flight::{FlightClient, FlightDescriptor, Ticket}; +use datafusion_comet_proto::spark_native_executors::{ + QueryRequest, ResultEndpoint, ResultEndpoints, +}; +use datafusion_comet_proto::PROTO_FINGERPRINT; +use futures::StreamExt; +use jni::objects::{JByteArray, JClass, JLongArray, JString, ReleaseMode}; +use jni::sys::{jbyteArray, jlong}; +use jni::EnvUnowned; +use prost::Message; +use tonic::transport::{Channel, Endpoint}; + +use super::COMET_VERSION; +use crate::errors::{try_unwrap_or_throw, CometError, CometResult}; +use crate::execution::jni_api::{export_batch, get_runtime}; + +fn client_error(context: &str, error: impl std::fmt::Display) -> CometError { + CometError::Internal(format!("Comet Native Executors: {context}: {error}")) +} + +/// Connects to the Flight service at `location`, a Flight location such as +/// `grpc+tcp://host:port` or an `http://` URL. +async fn connect(location: &str) -> CometResult { + let uri = match location.strip_prefix("grpc+tcp://") { + Some(address) => format!("http://{address}"), + None => location.to_string(), + }; + let channel: Channel = Endpoint::from_shared(uri) + .map_err(|e| client_error(&format!("invalid location {location}"), e))? + .connect() + .await + .map_err(|e| client_error(&format!("cannot connect to {location}"), e))?; + // Query plans and result batches can be larger than gRPC's default 4 MiB message limit. + let inner = FlightServiceClient::new(channel) + .max_decoding_message_size(usize::MAX) + .max_encoding_message_size(usize::MAX); + Ok(FlightClient::new_from_inner(inner)) +} + +/// Submits `query`, a serialized `QueryRequest`, to the scheduler at `scheduler_url` and waits +/// for it to finish. Returns where its result partitions are. +async fn submit(scheduler_url: &str, mut query: Vec) -> CometResult { + // Stamp the query with this build's version and plan format. A field that appears again + // replaces the first, so appending them spares decoding and encoding the whole query. + QueryRequest { + comet_version: COMET_VERSION.to_string(), + proto_fingerprint: PROTO_FINGERPRINT.to_string(), + ..Default::default() + } + .encode(&mut query) + .map_err(|e| client_error("bad query", e))?; + let mut client = connect(scheduler_url).await?; + let info = client + .get_flight_info(FlightDescriptor::new_cmd(query)) + .await + .map_err(|e| client_error("the query failed", e))?; + let endpoints = info + .endpoint + .into_iter() + .map(|endpoint| { + // A partition without a location is served by the scheduler itself. + let location = endpoint.location.first().map_or_else( + || scheduler_url.to_string(), + |location| location.uri.clone(), + ); + let ticket = endpoint + .ticket + .map(|ticket| ticket.ticket.to_vec()) + .ok_or_else(|| client_error("bad result", "a partition has no ticket"))?; + Ok(ResultEndpoint { location, ticket }) + }) + .collect::>>()?; + Ok(ResultEndpoints { endpoints }) +} + +/// Runs a query: `query` is a serialized `QueryRequest`, which this fills in with the driver's +/// Comet version and plan format. Returns a serialized `ResultEndpoints`. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_NativeExecutorsClient_submit( + e: EnvUnowned, + _class: JClass, + scheduler_url: JString, + query: JByteArray, +) -> jbyteArray { + try_unwrap_or_throw(&e, |env| { + let scheduler_url = scheduler_url + .try_to_string(env) + .map_err(|e| client_error("bad scheduler URL", e))?; + let query = env.convert_byte_array(query)?; + let endpoints = get_runtime().block_on(submit(&scheduler_url, query))?; + Ok(env + .byte_array_from_slice(&endpoints.encode_to_vec())? + .into_raw()) + }) +} + +/// Starts reading a result partition from the executor at `location`. Returns a handle for +/// `nextBatch` and `closePartition`. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_NativeExecutorsClient_openPartition( + e: EnvUnowned, + _class: JClass, + location: JString, + ticket: JByteArray, +) -> jlong { + try_unwrap_or_throw(&e, |env| { + let location = location + .try_to_string(env) + .map_err(|e| client_error("bad location", e))?; + let ticket = env.convert_byte_array(ticket)?; + let stream = get_runtime().block_on(async { + connect(&location) + .await? + .do_get(Ticket::new(ticket)) + .await + .map_err(|e| client_error(&format!("cannot read a partition from {location}"), e)) + })?; + Ok(Box::into_raw(Box::new(stream)) as jlong) + }) +} + +/// Moves the partition's next batch into the Arrow C Data Interface structs at `array_addrs` +/// and `schema_addrs`, one pair per column. Returns its row count, or -1 at the end. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_NativeExecutorsClient_nextBatch( + e: EnvUnowned, + _class: JClass, + handle: jlong, + array_addrs: JLongArray, + schema_addrs: JLongArray, +) -> jlong { + try_unwrap_or_throw(&e, |env| { + // SAFETY: `handle` comes from `openPartition` and hasn't been closed. + let stream = unsafe { &mut *(handle as *mut FlightRecordBatchStream) }; + let Some(batch) = get_runtime().block_on(stream.next()) else { + return Ok(-1); + }; + let batch = batch.map_err(|e| client_error("cannot read a partition", e))?; + let array_addrs = unsafe { array_addrs.get_elements(env, ReleaseMode::NoCopyBack)? }; + let schema_addrs = unsafe { schema_addrs.get_elements(env, ReleaseMode::NoCopyBack)? }; + // SAFETY: the JVM allocated one array and one schema struct per column. + unsafe { export_batch(&array_addrs, &schema_addrs, batch, false) } + }) +} + +/// Stops reading a partition and releases its handle. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_NativeExecutorsClient_closePartition( + e: EnvUnowned, + _class: JClass, + handle: jlong, +) { + try_unwrap_or_throw(&e, |_env| { + if handle != 0 { + // SAFETY: `handle` comes from `openPartition`, and the JVM closes it once. + drop(unsafe { Box::from_raw(handle as *mut FlightRecordBatchStream) }); + } + Ok(()) + }) +} diff --git a/native/core/src/execution/native_executors/fragment.rs b/native/core/src/execution/native_executors/fragment.rs index 4e4aec790e0..fdd67ab9280 100644 --- a/native/core/src/execution/native_executors/fragment.rs +++ b/native/core/src/execution/native_executors/fragment.rs @@ -24,6 +24,7 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef, UInt32Type}; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{internal_datafusion_err, DataFusionError, Result}; use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; use datafusion::execution::disk_manager::DiskManager; use datafusion::execution::memory_pool::{MemoryPool, UnboundedMemoryPool}; use datafusion::execution::TaskContext; @@ -41,7 +42,9 @@ use datafusion_comet_proto::spark_operator::{Operator, SparkStructField}; use futures::stream::BoxStream; use futures::{StreamExt, TryStreamExt}; -use crate::execution::jni_api::{comet_runtime_env, comet_session_context}; +use crate::execution::jni_api::{ + comet_runtime_env, comet_session_context, NATIVE_EXECUTORS_FRAGMENT, +}; use crate::execution::planner::{ align_shuffle_writer_input, convert_spark_types_to_arrow_schema, PhysicalPlanner, }; @@ -88,6 +91,29 @@ pub fn partition_descriptor( Ok(MemorySourceConfig::try_new_exec(&partitions, schema, None)?) } +/// The partitions' own copies of the block in a fragment's partition `descriptor`, one per +/// partition, or none when its partitions all run the fragment's block. +pub fn partition_plans(descriptor: &Arc) -> Result>> { + let memory = descriptor + .downcast_ref::() + .and_then(|source| source.data_source().downcast_ref::()) + .ok_or_else(|| internal_datafusion_err!("Not a fragment's partition descriptor"))?; + let mut plans = Vec::new(); + for partition in memory.partitions() { + let [batch] = partition.as_slice() else { + return Err(internal_datafusion_err!( + "A fragment's partition descriptor must have one row per partition" + )); + }; + let column = batch.column(1).as_binary::(); + if column.is_null(0) { + return Ok(vec![]); + } + plans.push(column.value(0).to_vec()); + } + Ok(plans) +} + /// Runs a Comet native block for each partition of a stage. /// /// The first child is the partition descriptor (see [`partition_descriptor`]), and each other @@ -126,15 +152,6 @@ impl CometFragmentExec { )); } let partitions = descriptor.output_partitioning().partition_count(); - for (index, (input, broadcast)) in inputs.iter().zip(&fragment.broadcast).enumerate() { - let expected = if *broadcast { 1 } else { partitions }; - let actual = input.output_partitioning().partition_count(); - if actual != expected { - return Err(internal_datafusion_err!( - "CometFragmentExec input {index} has {actual} partitions, expected {expected}" - )); - } - } let props = Arc::new(PlanProperties::new( EquivalenceProperties::new(fragment_schema(&fragment.output_schema)), Partitioning::UnknownPartitioning(partitions), @@ -151,13 +168,34 @@ impl CometFragmentExec { pub fn fragment(&self) -> &Arc { &self.fragment } + + /// Checks that each input's partitions line up with the descriptor's. This is checked when the + /// plan runs rather than when it is built, since an optimizer can rebuild the node with an + /// input whose exchange it has taken away and is about to put back. + fn check_inputs(&self) -> Result<()> { + let partitions = self.props.output_partitioning().partition_count(); + for (index, (input, broadcast)) in self.children[1..] + .iter() + .zip(&self.fragment.broadcast) + .enumerate() + { + let expected = if *broadcast { 1 } else { partitions }; + let actual = input.output_partitioning().partition_count(); + if actual != expected { + return Err(internal_datafusion_err!( + "CometFragmentExec input {index} has {actual} partitions, expected {expected}" + )); + } + } + Ok(()) + } } /// The schema of a fragment's output: the types Spark declared for it, every column nullable. /// The cast to those types keeps a column nullable wherever the block's output is, even where /// Spark knows it has no nulls, such as a join key a filter has checked, so a fragment can't /// promise more. Nothing between stages relies on a column having no nulls. -fn fragment_schema(output_schema: &[SparkStructField]) -> SchemaRef { +pub fn fragment_schema(output_schema: &[SparkStructField]) -> SchemaRef { let declared = convert_spark_types_to_arrow_schema(output_schema); let fields: Vec = declared .fields() @@ -193,10 +231,13 @@ impl ExecutionPlan for CometFragmentExec { /// The descriptor's and the partitioned inputs' partitions line up with this plan's, while a /// broadcast input is read whole by every partition, so a distributed engine must not split - /// it between tasks. + /// it between tasks. A stage with one partition reads every input as one partition, which + /// keeps an optimizer such as `EnsureRequirements` from removing the coalesce that gathers + /// it. fn input_distribution_requirements(&self) -> InputDistributionRequirements { + let single = self.props.output_partitioning().partition_count() == 1; let inputs = self.fragment.broadcast.iter().map(|broadcast| { - if *broadcast { + if *broadcast || single { Distribution::SinglePartition } else { Distribution::UnspecifiedDistribution @@ -237,6 +278,7 @@ impl ExecutionPlan for CometFragmentExec { partition: usize, context: Arc, ) -> Result { + self.check_inputs()?; let descriptor = Arc::clone(&self.children[0]); let inputs = self.children[1..] .iter() @@ -350,11 +392,13 @@ fn fragment_session( ) -> Result> { let mut runtime = comet_runtime_env(memory_pool); runtime.disk_manager = disk_manager; + let mut conf = conf.clone(); + conf.insert(NATIVE_EXECUTORS_FRAGMENT.to_string(), "true".to_string()); let session = comet_session_context( runtime, conf.get_usize(COMET_BATCH_SIZE, DEFAULT_BATCH_SIZE), 1, - conf, + &conf, spark_plan, ) .map_err(|e| DataFusionError::External(Box::new(e)))?; diff --git a/native/core/src/execution/native_executors/mod.rs b/native/core/src/execution/native_executors/mod.rs index fc0b9d6b2f4..9201fd652b4 100644 --- a/native/core/src/execution/native_executors/mod.rs +++ b/native/core/src/execution/native_executors/mod.rs @@ -24,9 +24,13 @@ //! partition of a stage, reading its inputs from other DataFusion plans, such as a distributed //! engine's shuffle readers, instead of from the JVM. +#[cfg(feature = "native-executors")] +mod client; mod fragment; -pub use fragment::{block_schema, partition_descriptor, CometFragmentExec}; +pub use fragment::{ + block_schema, fragment_schema, partition_descriptor, partition_plans, CometFragmentExec, +}; /// The version of this Comet build. A Comet Native Executor cluster requires a driver's to match /// its own, together with [`datafusion_comet_proto::PROTO_FINGERPRINT`]. diff --git a/native/native-executors/Cargo.lock b/native/native-executors/Cargo.lock index 4d21ac8a3e5..b70bba3af43 100644 --- a/native/native-executors/Cargo.lock +++ b/native/native-executors/Cargo.lock @@ -996,11 +996,19 @@ dependencies = [ "tokio", ] +[[package]] +name = "ballista-api-types" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion-ballista?branch=comet-native#0f092ee494d47d611f6e83f3660e60a4e4038d8c" +dependencies = [ + "serde", + "utoipa", +] + [[package]] name = "ballista-core" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a090511fa9b3661dbf17444ecbbf5efcd11793e140ddd7fe2fb85f2de88f561" +source = "git+https://github.com/apache/datafusion-ballista?branch=comet-native#0f092ee494d47d611f6e83f3660e60a4e4038d8c" dependencies = [ "arrow-flight", "async-trait", @@ -1023,14 +1031,14 @@ dependencies = [ "tonic-prost", "tonic-prost-build", "url", + "utoipa", "uuid", ] [[package]] name = "ballista-executor" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "403cdac6a4e822e6025348f4f7e43765ba0ca176ca1f225888a9eeb90b9c66ce" +source = "git+https://github.com/apache/datafusion-ballista?branch=comet-native#0f092ee494d47d611f6e83f3660e60a4e4038d8c" dependencies = [ "arrow", "arrow-flight", @@ -1056,16 +1064,29 @@ dependencies = [ "uuid", ] +[[package]] +name = "ballista-history" +version = "55.0.0" +source = "git+https://github.com/apache/datafusion-ballista?branch=comet-native#0f092ee494d47d611f6e83f3660e60a4e4038d8c" +dependencies = [ + "ballista-api-types", + "log", + "serde", + "serde_json", + "tokio", +] + [[package]] name = "ballista-scheduler" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "845bb1c2c763ae2b50ecfd9debdfe8e092d13a4fefe84c20e9262bb191a1f1e7" +source = "git+https://github.com/apache/datafusion-ballista?branch=comet-native#0f092ee494d47d611f6e83f3660e60a4e4038d8c" dependencies = [ "arrow-flight", "async-trait", "axum", + "ballista-api-types", "ballista-core", + "ballista-history", "dashmap", "datafusion", "datafusion-proto", @@ -1079,11 +1100,14 @@ dependencies = [ "prost", "rand 0.10.3", "serde", + "serde_json", "tokio", "tokio-stream", "tonic", "tower-http 0.7.1", + "tracing", "url-escape", + "utoipa", "uuid", ] @@ -3471,7 +3495,7 @@ dependencies = [ [[package]] name = "iceberg" version = "0.10.1" -source = "git+https://github.com/apache/iceberg-rust?rev=1f3bc34058e3f3e45336bc5428f37b92649a608b#1f3bc34058e3f3e45336bc5428f37b92649a608b" +source = "git+https://github.com/apache/iceberg-rust?rev=b154603a710d730d2f43f8af4db257cc137237ef#b154603a710d730d2f43f8af4db257cc137237ef" dependencies = [ "aes-gcm", "anyhow", @@ -3529,7 +3553,7 @@ dependencies = [ [[package]] name = "iceberg-property-macro" version = "0.10.1" -source = "git+https://github.com/apache/iceberg-rust?rev=1f3bc34058e3f3e45336bc5428f37b92649a608b#1f3bc34058e3f3e45336bc5428f37b92649a608b" +source = "git+https://github.com/apache/iceberg-rust?rev=b154603a710d730d2f43f8af4db257cc137237ef#b154603a710d730d2f43f8af4db257cc137237ef" dependencies = [ "proc-macro2", "quote", @@ -3539,7 +3563,7 @@ dependencies = [ [[package]] name = "iceberg-storage-opendal" version = "0.10.1" -source = "git+https://github.com/apache/iceberg-rust?rev=1f3bc34058e3f3e45336bc5428f37b92649a608b#1f3bc34058e3f3e45336bc5428f37b92649a608b" +source = "git+https://github.com/apache/iceberg-rust?rev=b154603a710d730d2f43f8af4db257cc137237ef#b154603a710d730d2f43f8af4db257cc137237ef" dependencies = [ "anyhow", "async-trait", @@ -4081,9 +4105,9 @@ dependencies = [ [[package]] name = "lru" -version = "0.18.5" +version = "0.19.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef9ac18847474e638e3702b76c65d4eb93428471a74778ef0f1be711717f89b5" +checksum = "4a9f71e3591c7ec44d3e563e87bf14ddc2bada8519dca046e5b0e560cd004e2d" dependencies = [ "hashbrown 0.17.1", ] @@ -5076,8 +5100,8 @@ version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "03da047801ff44bb6a4d407d4860c05fd70bb81714e6b2f3812603d5b145b042" dependencies = [ - "heck 0.5.0", - "itertools 0.14.0", + "heck 0.4.1", + "itertools 0.13.0", "log", "multimap", "petgraph", @@ -5098,7 +5122,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" dependencies = [ "anyhow", - "itertools 0.14.0", + "itertools 0.13.0", "proc-macro2", "quote", "syn 2.0.119", @@ -6880,6 +6904,29 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" +[[package]] +name = "utoipa" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8765fe27aeff71012a3f90fa474cf1df863a7d0dd81ed25de4e63fff2544994f" +dependencies = [ + "indexmap 2.14.2", + "serde", + "serde_json", + "utoipa-gen", +] + +[[package]] +name = "utoipa-gen" +version = "6.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d935f1c83fdf8b88f09bbe050d0cea78ea4afee6d7b0317403228c1200b2c8ab" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.9", +] + [[package]] name = "uuid" version = "1.29.0" diff --git a/native/native-executors/Cargo.toml b/native/native-executors/Cargo.toml index 117cd7e3df4..c298f5c9f08 100644 --- a/native/native-executors/Cargo.toml +++ b/native/native-executors/Cargo.toml @@ -19,8 +19,10 @@ # DataFusion Ballista. See https://github.com/apache/datafusion-comet/issues/7288. # # This is its own workspace, with its own lockfile, so that the main Comet build never depends -# on Ballista. It builds Comet's native crates by path together with the Ballista release on the -# same DataFusion major version. +# on Ballista. It builds Comet's native crates by path together with Ballista's `comet-native` +# branch, which is Ballista's main branch with the enhancements Comet Native Executors need that +# haven't merged yet (apache/datafusion-ballista#2682, #2721, #2722, #2723 and #2743), on the same +# DataFusion major version. [workspace] [package] @@ -35,9 +37,11 @@ publish = false [dependencies] arrow-flight = "59.2.0" -ballista-core = "55.0.0" -ballista-executor = { version = "55.0.0", default-features = false, features = ["arrow-ipc-optimizations"] } -ballista-scheduler = { version = "55.0.0", default-features = false } +ballista-core = { git = "https://github.com/apache/datafusion-ballista", branch = "comet-native" } +ballista-executor = { git = "https://github.com/apache/datafusion-ballista", branch = "comet-native", default-features = false, features = ["arrow-ipc-optimizations"] } +# `rest-api` compiles the scheduler's event-log writer (COMET_SCHEDULER_EVENT_LOG_DIR) and its REST API. +# Its code calls `tracing`, which Ballista links only for its binaries (`build-binary`), so it is added here. +ballista-scheduler = { git = "https://github.com/apache/datafusion-ballista", branch = "comet-native", default-features = false, features = ["rest-api", "tracing"] } datafusion = { version = "55.2.0", default-features = false } # Comet's core crate, under the name of its library. Without the default hdfs-opendal feature, # which links libjvm through libhdfs. diff --git a/native/native-executors/README.md b/native/native-executors/README.md index a271ba863f9..8d4f5235ff0 100644 --- a/native/native-executors/README.md +++ b/native/native-executors/README.md @@ -48,15 +48,22 @@ A client runs a query through the scheduler's Arrow Flight service: A stage that several other stages read runs once for each of them, since Ballista's planner doesn't share a stage between consumers yet. +Spark plans a global sort as a range exchange followed by a sort of each partition, and its +driver samples the exchange's input for the ranges' bounds. A cluster doesn't range-partition: +the stage reads the producer's partitions as they are, sorts each, and merges the sorted +partitions in the order the driver sends with the stage, as DataFusion plans a global sort. A +stage whose order isn't over columns of types Arrow orders as Spark does, such as floats, reads +the exchange as one partition and sorts it in a single task. + Each stage runs on the executors as a `CometFragmentExec` (`native/core/src/execution/native_executors`), which plans the stage's block with Comet's planner for each partition. Ballista moves the data between stages: -| Exchange | Ballista plan | -| --------------------------- | --------------------------------------------------------------------------------------------------------------------------------- | -| Hash | `RepartitionExec(Hash)`: Ballista's hash shuffle, with DataFusion's hash | -| Single partition, broadcast | `CoalescePartitionsExec` | -| Range | `CoalescePartitionsExec`, so a sort runs in one task ([ballista#2680](https://github.com/apache/datafusion-ballista/issues/2680)) | +| Exchange | Ballista plan | +| --------------------------- | -------------------------------------------------------------------------------------------- | +| Hash | `RepartitionExec(Hash)`: Ballista's hash shuffle, with DataFusion's hash | +| Single partition, broadcast | `CoalescePartitionsExec` | +| Range | None: the stage sorts the producer's partitions, and a `SortPreservingMergeExec` merges them | A stage that fails to plan is refused, not run wrongly or loaded with code it names: @@ -67,6 +74,24 @@ A stage that fails to plan is refused, not run wrongly or loaded with code it na - A stage that reads files can't also read a hash exchange. Spark places the rows of a bucketed scan with Murmur3, so a join with a side shuffled by DataFusion's hash would return wrong rows. +## Joins whose build side turns out small + +Spark plans a join's sides from its estimates, and its estimate of a filtered relation ignores the +filter, so it shuffles both sides of a join whose build side is far smaller than it thinks. The +scheduler moves such a join into the stage that produces its probe side, with the build side +broadcast to it, so the probe side's exchange carries the join's output rather than all of its +input (`src/adaptive.rs`). It picks the joins Spark's estimates say are worth measuring: an inner +hash join built on the right, with column keys and no condition, whose build side Spark estimates +within Ballista's maximum estimate multiple of its broadcast threshold and whose probe side many +times larger. + +- In blocking mode, a rule in Ballista's adaptive planner holds the probe side's stage back until + the build side's stage finishes, then moves the join if the build side measured under + `ballista.optimizer.broadcast_join_threshold_bytes`, or releases the probe side's stage as it + was. +- In pipelined mode every task is placed before any runs, so a build side cannot be measured first + and the join stays as Spark planned it. + ## Building and running ```shell @@ -90,8 +115,16 @@ binaries and runs queries through the scheduler's Flight service, as a driver wo ## Keeping it building -This workspace builds Comet's native crates by path against the Ballista release on the same -DataFusion major version, and CI builds and tests it in the `linux-test-native-executors` job. When -Comet moves to a DataFusion major that Ballista hasn't released yet, the upgrade can turn that job -off until Ballista catches up. Dependabot updates this workspace's lockfile, which is separate from -`native/Cargo.lock`, so the two can resolve a crate to different versions between updates. +This workspace builds Comet's native crates by path against Ballista's +[`comet-native`](https://github.com/apache/datafusion-ballista/tree/comet-native) branch: Ballista's +main branch with the enhancements Comet Native Executors need that haven't merged yet +([#2682](https://github.com/apache/datafusion-ballista/pull/2682), +[#2721](https://github.com/apache/datafusion-ballista/pull/2721), +[#2722](https://github.com/apache/datafusion-ballista/pull/2722), +[#2723](https://github.com/apache/datafusion-ballista/pull/2723) and +[#2743](https://github.com/apache/datafusion-ballista/pull/2743)), on the same DataFusion major +version. `Cargo.lock` pins the branch's commit. CI builds and tests the workspace in the +`linux-test-native-executors` job. When Comet moves to a DataFusion major that Ballista hasn't +caught up with, the upgrade can turn that job off until it does. Dependabot updates this +workspace's lockfile, which is separate from `native/Cargo.lock`, so the two can resolve a crate to +different versions between updates. diff --git a/native/native-executors/src/adaptive.rs b/native/native-executors/src/adaptive.rs new file mode 100644 index 00000000000..b93b606664c --- /dev/null +++ b/native/native-executors/src/adaptive.rs @@ -0,0 +1,499 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Adaptive join planning for Comet fragments. +//! +//! A stage's joins are inside its Comet block, so Ballista's adaptive join rules see the stage +//! as a leaf, and Spark has already placed an exchange on each side of every shuffled join. When +//! a join's build side turns out small, broadcasting it is not enough: the probe side would still +//! be shuffled whole to reach the join. These rules move the join into the stage that produces +//! the probe side instead, with the build side broadcast to it, so the probe side's exchange +//! carries the join's output rather than all of its input. +//! +//! That only pays off if the probe side's stage hasn't run yet when the build side's size is +//! known. [`HoldProbeStages`] holds it back before anything runs, where Spark's estimates say the +//! measurement could change the plan, and [`BroadcastIntoProbe`] decides once the build side's +//! stage finishes: it moves the join where the build side measured small, and otherwise releases +//! the probe side's stage as it was. Holding serialises the two stages, so like Ballista's own +//! build-side staging it uses Ballista's broadcast threshold, minimum probe ratio and maximum +//! estimate multiple. + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; + +use ballista_core::config::BallistaConfig; +use ballista_scheduler::state::adaptive::{AdaptiveRule, AdaptiveRules, ExchangeExec}; +use comet::execution::native_executors::{ + partition_descriptor, partition_plans, CometFragmentExec, +}; +use datafusion::common::tree_node::{Transformed, TreeNode}; +use datafusion::common::{internal_datafusion_err, Result}; +use datafusion::config::ConfigOptions; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_plan::statistics::StatisticsContext; +use datafusion::physical_plan::{ + ChildrenPropertiesMode, ExecutionPlan, Partitioning, ReplaceChildrenOptions, StatisticsArgs, +}; +use datafusion_comet_proto::spark_expression::expr::ExprStruct; +use datafusion_comet_proto::spark_native_executors::Fragment; +use datafusion_comet_proto::spark_operator::operator::OpStruct; +use datafusion_comet_proto::spark_operator::{BuildSide, JoinType, Operator, ShuffleScan}; +use log::info; +use prost::Message; + +/// The rules Comet adds to Ballista's adaptive planning. +#[derive(Debug, Default)] +pub struct CometAdaptiveRules; + +impl AdaptiveRules for CometAdaptiveRules { + fn preparation_rules(&self, _plan_ids: Arc) -> Vec { + vec![Arc::new(HoldProbeStages)] + } + + fn replan_rules(&self, plan_ids: Arc) -> Vec { + vec![Arc::new(BroadcastIntoProbe { plan_ids })] + } +} + +/// A join in a stage's block that can run in its probe side's stage instead. +/// +/// It is an inner hash join built on the right, with no condition and with keys that are +/// columns, whose inputs are both input leaves of the block. Moved below the probe side's +/// exchange its output keeps the probe side's columns first, so the exchange's keys still name +/// the same columns. +struct MovableJoin { + /// The child indexes from the block's root to the join. + path: Vec, + /// The stage inputs that the join's probe (left) and build (right) leaves read. + probe_input: usize, + build_input: usize, +} + +/// The joins of `block` that [`MovableJoin`] describes. +fn movable_joins(block: &Operator) -> Vec { + fn walk( + op: &Operator, + path: &mut Vec, + leaves: &mut usize, + found: &mut Vec, + ) { + if is_input_leaf(op) { + *leaves += 1; + return; + } + if let Some(OpStruct::HashJoin(join)) = &op.op_struct { + let movable = join.join_type == JoinType::Inner as i32 + && join.build_side == BuildSide::BuildRight as i32 + && join.condition.is_none() + && !join.null_aware_anti_join + && join + .left_join_keys + .iter() + .chain(&join.right_join_keys) + .all(|key| matches!(key.expr_struct, Some(ExprStruct::Bound(_)))) + && op.children.len() == 2 + && op.children.iter().all(is_input_leaf); + if movable { + found.push(MovableJoin { + path: path.clone(), + probe_input: *leaves, + build_input: *leaves + 1, + }); + } + } + for (index, child) in op.children.iter().enumerate() { + path.push(index); + walk(child, path, leaves, found); + path.pop(); + } + } + let mut found = vec![]; + walk(block, &mut vec![], &mut 0, &mut found); + found +} + +/// Whether `op` reads one of its stage's inputs. Comet's planner feeds them to these leaves in +/// the order a depth-first walk meets them. +fn is_input_leaf(op: &Operator) -> bool { + matches!( + op.op_struct, + Some(OpStruct::Scan(_)) | Some(OpStruct::ShuffleScan(_)) + ) +} + +fn leaf_fields(op: &Operator) -> &[datafusion_comet_proto::spark_expression::DataType] { + match &op.op_struct { + Some(OpStruct::Scan(scan)) => &scan.fields, + Some(OpStruct::ShuffleScan(scan)) => &scan.fields, + _ => &[], + } +} + +fn node_at<'a>(block: &'a Operator, path: &[usize]) -> &'a Operator { + path.iter().fold(block, |op, index| &op.children[*index]) +} + +/// A stage's block, or `None` for a stage that reads files, whose partitions each run their own +/// copy of it. +fn block_of(fragment: &CometFragmentExec) -> Result> { + let block = &fragment.fragment().block; + if block.is_empty() { + return Ok(None); + } + Ok(Some(Operator::decode(block.as_slice()).map_err(|e| { + internal_datafusion_err!("Cannot decode a stage's block: {e}") + })?)) +} + +/// A stage input's exchange, when it is a hash shuffle of another stage's output. +fn hash_exchange(input: &Arc) -> Option<(&ExchangeExec, &CometFragmentExec)> { + let exchange = input.downcast_ref::()?; + if exchange.is_broadcast() || !matches!(exchange.partitioning(), Some(Partitioning::Hash(..))) { + return None; + } + let producer = exchange.input().downcast_ref::()?; + Some((exchange, producer)) +} + +fn ballista_config(config: &ConfigOptions) -> BallistaConfig { + config + .extensions + .get::() + .cloned() + .unwrap_or_default() +} + +/// Holds back the probe side's stage of each [`MovableJoin`] whose build side Spark estimates +/// close enough to the broadcast threshold that measuring it could make it a broadcast, and +/// whose probe side Spark estimates large enough that not shuffling it would pay for the wait. +/// +/// Spark's estimate of a filtered relation ignores the filter, so a build side under a selective +/// filter is often far smaller than its estimate. +#[derive(Debug)] +pub struct HoldProbeStages; + +impl PhysicalOptimizerRule for HoldProbeStages { + fn optimize( + &self, + plan: Arc, + config: &ConfigOptions, + ) -> Result> { + let config = ballista_config(config); + let threshold = config.broadcast_join_threshold_bytes() as u64; + if threshold == 0 || !config.adaptive_join_enabled() { + return Ok(plan); + } + let max_build = + threshold.saturating_mul(config.stage_build_side_max_estimate_multiple() as u64); + let min_probe_ratio = config.stage_build_side_min_probe_ratio() as u64; + Ok(plan + .transform_up(|node| { + let Some(fragment) = node.downcast_ref::() else { + return Ok(Transformed::no(node)); + }; + let Some(block) = block_of(fragment)? else { + return Ok(Transformed::no(node)); + }; + let inputs = &fragment.children()[1..]; + let mut held = inputs + .iter() + .map(|input| Arc::clone(*input)) + .collect::>(); + let mut changed = false; + for join in movable_joins(&block) { + let (Some((probe, probe_stage)), Some((build, build_stage))) = ( + hash_exchange(inputs[join.probe_input]), + hash_exchange(inputs[join.build_input]), + ) else { + continue; + }; + let (Some(probe_bytes), Some(build_bytes)) = ( + probe_stage.fragment().estimated_output_bytes, + build_stage.fragment().estimated_output_bytes, + ) else { + continue; + }; + if probe.shuffle_created() + || build.shuffle_created() + || probe.is_inactive_stage() + || build_bytes > max_build + || probe_bytes < build_bytes.saturating_mul(min_probe_ratio) + { + continue; + } + info!( + "Holding back a join's probe side (estimated {probe_bytes} bytes) until \ + its build side (estimated {build_bytes} bytes) is measured" + ); + held[join.probe_input] = Arc::new(probe.with_inactive_stage(true)); + changed = true; + } + if !changed { + return Ok(Transformed::no(node)); + } + let children = std::iter::once(Arc::clone(fragment.children()[0])) + .chain(held) + .collect(); + Ok(Transformed::yes(with_children(&node, children)?)) + })? + .data) + } + + fn name(&self) -> &str { + "HoldProbeStages" + } + + fn schema_check(&self) -> bool { + true + } +} + +/// Decides each join whose probe side [`HoldProbeStages`] held back, once its build side's stage +/// has finished: a build side under the broadcast threshold is broadcast to the probe side's +/// stage, which runs the join, and otherwise the probe side's stage is released unchanged. +#[derive(Debug)] +pub struct BroadcastIntoProbe { + plan_ids: Arc, +} + +impl PhysicalOptimizerRule for BroadcastIntoProbe { + fn optimize( + &self, + plan: Arc, + config: &ConfigOptions, + ) -> Result> { + let config = ballista_config(config); + let threshold_bytes = config.broadcast_join_threshold_bytes(); + let threshold_rows = config.broadcast_join_threshold_rows(); + Ok(plan + .transform_up(|node| { + let mut node = node; + let mut changed = false; + // Each decision changes the block, so find the joins again after each. + while let Some(decided) = self.decide_one(&node, threshold_bytes, threshold_rows)? { + node = decided; + changed = true; + } + Ok(if changed { + Transformed::yes(node) + } else { + Transformed::no(node) + }) + })? + .data) + } + + fn name(&self) -> &str { + "BroadcastIntoProbe" + } + + fn schema_check(&self) -> bool { + true + } +} + +impl BroadcastIntoProbe { + /// `node` with one held join decided, or `None` when it has none whose build side is ready. + fn decide_one( + &self, + node: &Arc, + threshold_bytes: usize, + threshold_rows: usize, + ) -> Result>> { + let Some(fragment) = node.downcast_ref::() else { + return Ok(None); + }; + let Some(block) = block_of(fragment)? else { + return Ok(None); + }; + let inputs = &fragment.children()[1..]; + for join in movable_joins(&block) { + let (Some((probe, _)), Some((build, _))) = ( + hash_exchange(inputs[join.probe_input]), + hash_exchange(inputs[join.build_input]), + ) else { + continue; + }; + if !probe.is_inactive_stage() || !build.shuffle_created() { + continue; + } + let stats = StatisticsContext::new().compute(build, &StatisticsArgs::new())?; + let small = match ( + stats.total_byte_size.get_value(), + stats.num_rows.get_value(), + ) { + (Some(bytes), _) => *bytes <= threshold_bytes, + (None, Some(rows)) => *rows <= threshold_rows, + (None, None) => false, + }; + info!( + "A held join's build side measured {:?} bytes and {:?} rows: {}", + stats.total_byte_size, + stats.num_rows, + if small { + "broadcasting it to the probe side's stage" + } else { + "shuffling the probe side" + } + ); + let decided = if small { + self.broadcast_into_probe(fragment, &block, &join)? + } else { + let mut children = fragment + .children() + .into_iter() + .map(Arc::clone) + .collect::>(); + children[1 + join.probe_input] = Arc::new(probe.with_inactive_stage(false)); + with_children(node, children)? + }; + return Ok(Some(decided)); + } + Ok(None) + } + + /// `consumer` with `join` moved into the stage that produces its probe side, and its build + /// side broadcast to that stage. + fn broadcast_into_probe( + &self, + consumer: &CometFragmentExec, + block: &Operator, + join: &MovableJoin, + ) -> Result> { + let inputs = &consumer.children()[1..]; + let (_, probe_stage) = hash_exchange(inputs[join.probe_input]) + .ok_or_else(|| internal_datafusion_err!("The join's probe side is not an exchange"))?; + let (build, build_stage) = hash_exchange(inputs[join.build_input]) + .ok_or_else(|| internal_datafusion_err!("The join's build side is not an exchange"))?; + let join_op = node_at(block, &join.path); + let build_leaf = &join_op.children[1]; + + // The probe side's stage runs the join on its output, reading the build side whole as + // its last input. The join becomes the root of the stage's block, so it takes over the + // block's pool of SQL texts. + let into_probe = |root: Operator| -> Operator { + let mut root = root; + let mut moved = join_op.clone(); + moved.sql_text_pool = std::mem::take(&mut root.sql_text_pool); + moved.children = vec![root, build_leaf.clone()]; + moved + }; + let probe_fragment = probe_stage.fragment(); + let block_of_probe = match block_of(probe_stage)? { + Some(root) => into_probe(root).encode_to_vec(), + None => vec![], + }; + let descriptor = Arc::clone(probe_stage.children()[0]); + let own_plans = partition_plans(&descriptor)?; + let descriptor = if own_plans.is_empty() { + descriptor + } else { + let plans = own_plans + .iter() + .map(|plan| { + let root = Operator::decode(plan.as_slice()).map_err(|e| { + internal_datafusion_err!("Cannot decode a partition's block: {e}") + })?; + Ok(into_probe(root).encode_to_vec()) + }) + .collect::>>()?; + partition_descriptor(plans.len(), &plans)? + }; + let mut broadcast = probe_fragment.broadcast.clone(); + broadcast.push(true); + let mut output_schema = probe_fragment.output_schema.clone(); + output_schema.extend(build_stage.fragment().output_schema.iter().cloned()); + let joined = Fragment { + block: block_of_probe, + output_schema, + broadcast, + conf: probe_fragment.conf.clone(), + estimated_output_bytes: None, + }; + let mut children = vec![descriptor]; + children.extend( + probe_stage.children()[1..] + .iter() + .map(|child| Arc::clone(*child)), + ); + children.push(Arc::new( + build.to_broadcast(self.plan_ids.fetch_add(1, Ordering::Relaxed)), + )); + let joined: Arc = + Arc::new(CometFragmentExec::try_new(Arc::new(joined), children)?); + let exchange = with_children(inputs[join.probe_input], vec![joined])?; + let exchange = exchange + .downcast_ref::() + .ok_or_else(|| internal_datafusion_err!("An exchange rebuilt as something else"))? + .with_inactive_stage(false); + + // The consumer reads the join's output where it read the join's probe side, and no + // longer reads the build side. + let mut fields = leaf_fields(&join_op.children[0]).to_vec(); + fields.extend(leaf_fields(build_leaf).iter().cloned()); + let leaf = Operator { + children: vec![], + plan_id: join_op.plan_id, + sql_text_pool: if join.path.is_empty() { + join_op.sql_text_pool.clone() + } else { + vec![] + }, + op_struct: Some(OpStruct::ShuffleScan(ShuffleScan { + fields, + source: "a join moved into the probe side's stage".to_string(), + })), + }; + let mut consumer_block = block.clone(); + replace_at(&mut consumer_block, &join.path, leaf); + let mut consumer_broadcast = consumer.fragment().broadcast.clone(); + consumer_broadcast.remove(join.build_input); + let consumer_fragment = Fragment { + block: consumer_block.encode_to_vec(), + broadcast: consumer_broadcast, + ..consumer.fragment().as_ref().clone() + }; + let mut children = consumer + .children() + .into_iter() + .map(Arc::clone) + .collect::>(); + children[1 + join.probe_input] = Arc::new(exchange); + children.remove(1 + join.build_input); + Ok(Arc::new(CometFragmentExec::try_new( + Arc::new(consumer_fragment), + children, + )?)) + } +} + +fn replace_at(block: &mut Operator, path: &[usize], replacement: Operator) { + let target = path + .iter() + .fold(block, |op, index| &mut op.children[*index]); + *target = replacement; +} + +fn with_children( + plan: &Arc, + children: Vec>, +) -> Result> { + Arc::clone(plan).replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) +} diff --git a/native/native-executors/src/assembly.rs b/native/native-executors/src/assembly.rs index 19cca64beae..89f79d2d7d6 100644 --- a/native/native-executors/src/assembly.rs +++ b/native/native-executors/src/assembly.rs @@ -21,19 +21,24 @@ use std::collections::HashMap; use std::sync::Arc; use comet::execution::native_executors::{ - block_schema, partition_descriptor, CometFragmentExec, COMET_VERSION, + block_schema, fragment_schema, partition_descriptor, CometFragmentExec, COMET_VERSION, }; +use datafusion::arrow::compute::SortOptions; +use datafusion::arrow::datatypes::{DataType, SchemaRef}; use datafusion::common::{plan_datafusion_err, plan_err, Result}; use datafusion::physical_expr::expressions::Column; -use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_expr::{LexOrdering, PhysicalExpr, PhysicalSortExpr}; use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion::physical_plan::repartition::RepartitionExec; +use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion::physical_plan::ExecutionPlanProperties; use datafusion::physical_plan::{ExecutionPlan, Partitioning}; use datafusion_comet_proto::spark_expression::expr::ExprStruct; +use datafusion_comet_proto::spark_expression::{Expr, NullOrdering, SortDirection}; use datafusion_comet_proto::spark_native_executors::{Fragment, QueryRequest, Stage, StageInput}; use datafusion_comet_proto::spark_partitioning::partitioning::PartitioningStruct; use datafusion_comet_proto::PROTO_FINGERPRINT; +use log::warn; use prost::Message; /// Checks that `request` comes from a driver running this Comet build. The plan format is not @@ -74,11 +79,10 @@ fn build_stage( producers: &[Arc], conf: &HashMap, ) -> Result> { - let block = stage + let plan = stage .plan .as_ref() - .ok_or_else(|| plan_datafusion_err!("The stage has no plan"))? - .encode_to_vec(); + .ok_or_else(|| plan_datafusion_err!("The stage has no plan"))?; // Spark places the rows of a bucketed scan by Murmur3 and shuffles only the other side of a // join to match. The hash exchange below would place the other side by DataFusion's hash, so // partition `p` of one side would no longer meet the rows of bucket `p` of the other, and the @@ -100,6 +104,12 @@ fn build_stage( differently from DataFusion's, so the two sides of a bucketed join wouldn't meet" ); } + // A stage that sorts what a range exchange delivers sorts the producer's partitions as they + // are, and the sorted partitions are merged, as DataFusion plans a global sort. + let merge = merge_ordering( + &stage.output_ordering, + &fragment_schema(&stage.output_schema), + ); let mut inputs = Vec::with_capacity(stage.inputs.len()); for input in &stage.inputs { let producer = producers.get(input.producer as usize).ok_or_else(|| { @@ -108,17 +118,21 @@ fn build_stage( input.producer ) })?; - inputs.push(exchange(input, Arc::clone(producer))?); + inputs.push(exchange(input, Arc::clone(producer), merge.is_some())?); } - // A stage that reads exchanges has as many partitions as they deliver, which - // `CometFragmentExec` checks they agree on. - let num_partitions = stage + // A stage that reads exchanges has as many partitions as they deliver, and they must agree. + let mut partitioned = stage .inputs .iter() .zip(&inputs) - .find(|(input, _)| !input.broadcast) - .map(|(_, exchange)| exchange.output_partitioning().partition_count()) - .unwrap_or(stage.num_partitions as usize); + .filter(|(input, _)| !input.broadcast) + .map(|(_, exchange)| exchange.output_partitioning().partition_count()); + let num_partitions = partitioned.next().unwrap_or(stage.num_partitions as usize); + if let Some(other) = partitioned.find(|count| *count != num_partitions) { + return plan_err!( + "The stage reads exchanges of {num_partitions} and {other} partitions, which must agree" + ); + } if num_partitions == 0 { return plan_err!("The stage has no partitions"); } @@ -128,8 +142,14 @@ fn build_stage( stage.partition_plans.len() ); } - // A stage that reads files carries its scans' schemas and options only in each partition's - // copy of the block, so check one of those. Checking reads none of its files. + // A stage that reads files runs each partition's own copy of the block, so the block itself + // is only needed for a stage that doesn't. Checking one of those copies reads none of its + // files. + let block = if stage.partition_plans.is_empty() { + plan.encode_to_vec() + } else { + vec![] + }; let planned = block_schema(stage.partition_plans.first().unwrap_or(&block), conf)?; if planned.fields().len() != stage.output_schema.len() { return plan_err!( @@ -145,20 +165,20 @@ fn build_stage( )?]; children.extend(inputs); let fragment = Fragment { - // A stage that reads files runs each partition's own copy of the block instead. - block: if stage.partition_plans.is_empty() { - block - } else { - vec![] - }, + block, broadcast: stage.inputs.iter().map(|input| input.broadcast).collect(), output_schema: stage.output_schema, conf: conf.clone(), + estimated_output_bytes: stage.estimated_output_bytes, }; - Ok(Arc::new(CometFragmentExec::try_new( - Arc::new(fragment), - children, - )?)) + let fragment: Arc = + Arc::new(CometFragmentExec::try_new(Arc::new(fragment), children)?); + Ok(match merge { + Some(ordering) if num_partitions > 1 => { + Arc::new(SortPreservingMergeExec::new(ordering, fragment)) + } + _ => fragment, + }) } /// The exchange that carries `producer`'s output to a stage, as a node Ballista's planner turns @@ -167,13 +187,18 @@ fn build_stage( /// A hash exchange partitions with DataFusion's hash function, not Spark's Murmur3. That is /// consistent within a query, which is all joins, aggregates and windows need, but a plan that /// relies on Spark's own placement of rows, such as a join with a bucketed table, would be wrong, -/// which `build_stage` refuses. +/// which `build_stage` refuses (https://github.com/apache/datafusion-comet/issues/7368). +/// +/// A range exchange whose stage merges its sorted partitions is no exchange at all: the stage +/// reads the producer's partitions as they are. Otherwise it is read as one partition, and the +/// stage sorts it in a single task. fn exchange( input: &StageInput, producer: Arc, + merged: bool, ) -> Result> { if input.broadcast { - return Ok(Arc::new(CoalescePartitionsExec::new(producer))); + return Ok(one_partition(producer)); } let partitioning = input .partitioning @@ -185,15 +210,13 @@ fn exchange( let keys = hash .hash_expression .iter() - .map(|expr| match expr.expr_struct.as_ref() { - Some(ExprStruct::Bound(bound)) - if (bound.index as usize) < schema.fields().len() => - { - let index = bound.index as usize; - Ok(Arc::new(Column::new(schema.field(index).name(), index)) - as Arc) - } - _ => plan_err!("Hash partitioning keys must be columns of the stage's input"), + .map(|expr| { + let index = bound_column(expr, schema.fields().len()).ok_or_else(|| { + plan_datafusion_err!( + "Hash partitioning keys must be columns of the stage's input" + ) + })?; + Ok(column(&schema, index)) }) .collect::>>()?; Ok(Arc::new(RepartitionExec::try_new( @@ -201,12 +224,83 @@ fn exchange( Partitioning::Hash(keys, hash.num_partitions as usize), )?)) } - // Ballista can't plan a range shuffle for a physical plan yet - // (https://github.com/apache/datafusion-ballista/issues/2680), so the consuming stage - // reads a range exchange as one partition too. That sorts correctly, in a single task. - Some(PartitioningStruct::SinglePartition(_) | PartitioningStruct::RangePartition(_)) => { - Ok(Arc::new(CoalescePartitionsExec::new(producer))) - } + Some(PartitioningStruct::RangePartition(_)) if merged => Ok(producer), + Some(PartitioningStruct::RangePartition(_)) => Ok(one_partition(producer)), + Some(PartitioningStruct::SinglePartition(_)) => Ok(one_partition(producer)), other => plan_err!("Comet Native Executors don't support this exchange yet: {other:?}"), } } + +/// `producer`'s output gathered into one partition. +fn one_partition(producer: Arc) -> Arc { + Arc::new(CoalescePartitionsExec::new(producer)) +} + +fn column(schema: &SchemaRef, index: usize) -> Arc { + Arc::new(Column::new(schema.field(index).name(), index)) +} + +/// The index of the column `expr` reads, if it reads a column of an input with `num_fields`. +fn bound_column(expr: &Expr, num_fields: usize) -> Option { + match expr.expr_struct.as_ref() { + Some(ExprStruct::Bound(bound)) if (bound.index as usize) < num_fields => { + Some(bound.index as usize) + } + _ => None, + } +} + +/// Whether Arrow orders values of `data_type` as Spark does, so that a merge, which compares raw +/// values, keeps the order Comet's sort put each partition in. Comet's sort normalizes floats +/// first (a NaN of either sign sorts last), so a float key, or a nested one that may hold +/// floats, can't be merged this way (https://github.com/apache/datafusion-comet/issues/7366). +fn orders_like_spark(data_type: &DataType) -> bool { + data_type.is_integer() + || matches!( + data_type, + DataType::Boolean + | DataType::Decimal32(_, _) + | DataType::Decimal64(_, _) + | DataType::Decimal128(_, _) + | DataType::Decimal256(_, _) + | DataType::Date32 + | DataType::Date64 + | DataType::Timestamp(_, _) + | DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Utf8View + | DataType::Binary + | DataType::LargeBinary + | DataType::BinaryView + ) +} + +/// The order to merge a stage's sorted partitions in: its declared output ordering, when every +/// key is a column of a type whose order Arrow and Spark agree on. `None` when there's nothing +/// to merge by, and the stage then reads its range exchange as one partition. +fn merge_ordering(orders: &[Expr], schema: &SchemaRef) -> Option { + let sort_exprs = orders + .iter() + .map(|order| { + let Some(ExprStruct::SortOrder(order)) = order.expr_struct.as_ref() else { + return None; + }; + let index = bound_column(order.child.as_deref()?, schema.fields().len())?; + if !orders_like_spark(schema.field(index).data_type()) { + warn!( + "A global sort by a {} column runs in one task", + schema.field(index).data_type() + ); + return None; + } + Some(PhysicalSortExpr::new( + column(schema, index), + SortOptions { + descending: order.direction == SortDirection::Descending as i32, + nulls_first: order.null_ordering == NullOrdering::NullsFirst as i32, + }, + )) + }) + .collect::>>()?; + LexOrdering::new(sort_exprs) +} diff --git a/native/native-executors/src/bin/comet-executor.rs b/native/native-executors/src/bin/comet-executor.rs index f1b7cec6a99..9a08a7c62df 100644 --- a/native/native-executors/src/bin/comet-executor.rs +++ b/native/native-executors/src/bin/comet-executor.rs @@ -26,6 +26,11 @@ //! - `COMET_EXECUTOR_GRPC_PORT` (default `50052`) //! - `COMET_SCHEDULER_HOST` (default `localhost`) and `COMET_SCHEDULER_PORT` (default `50050`) //! - `COMET_EXECUTOR_TASKS`: concurrent tasks (default: the number of CPUs) +//! - `COMET_GRPC_MAX_MESSAGE_SIZE`: the largest gRPC message, in bytes, the executor accepts or +//! sends, such as a task with many partitions of a scan (default: Ballista's 16 MiB) +//! - `COMET_EXECUTOR_MEMORY_POOL_SIZE`: the bytes of memory the executor's tasks share (default: +//! Ballista's, 70% of the container's memory limit). Memory outside the pool, such as batches +//! in flight between stages, comes on top of it. //! - `COMET_EXECUTOR_WORK_DIR`: where shuffle and spill files go (default: a temporary directory) use std::sync::Arc; @@ -61,6 +66,13 @@ async fn main() -> Result<(), Box> { config.vcores = tasks; } config.work_dir = env_var("COMET_EXECUTOR_WORK_DIR")?; + if let Some(size) = env_var("COMET_EXECUTOR_MEMORY_POOL_SIZE")? { + config.memory_pool_size = Some(size); + } + if let Some(size) = env_var("COMET_GRPC_MAX_MESSAGE_SIZE")? { + config.grpc_max_decoding_message_size = size; + config.grpc_max_encoding_message_size = size; + } start_executor_process(Arc::new(config)).await?; Ok(()) diff --git a/native/native-executors/src/bin/comet-scheduler.rs b/native/native-executors/src/bin/comet-scheduler.rs index d1506f335ac..f5c53a7489f 100644 --- a/native/native-executors/src/bin/comet-scheduler.rs +++ b/native/native-executors/src/bin/comet-scheduler.rs @@ -20,10 +20,19 @@ //! //! Configured through environment variables, with Ballista's defaults otherwise: //! - `COMET_SCHEDULER_BIND_HOST` (default `127.0.0.1`) +//! - `COMET_SCHEDULER_EXTERNAL_HOST`: the host executors reach this scheduler at, which they +//! report task status to (default `localhost`) //! - `COMET_SCHEDULER_PORT`: Ballista's gRPC port, for executors (default `50050`) //! - `COMET_SCHEDULER_FLIGHT_PORT`: the Flight port, for drivers (default `50060`) //! - `COMET_SCHEDULER_RESULT_RETENTION_SECONDS`: how long a query's results stay readable after //! the query ends, whether or not a driver has read them (default `300`). `0` keeps them. +//! - `COMET_GRPC_MAX_MESSAGE_SIZE`: the largest gRPC message, in bytes, the scheduler sends or +//! accepts, such as a task with many partitions of a scan (default: Ballista's 16 MiB) +//! - `COMET_SCHEDULER_BALLISTA_CONFIG`: Ballista settings for every query, as comma-separated +//! `key=value` pairs, such as `ballista.execution.mode=pipelined` +//! - `COMET_SCHEDULER_EVENT_LOG_DIR`: a directory to write one event log per query to, as +//! Ballista's scheduler does: the query's plan, a record for each task, and each stage's rows +//! and compute time (default: no event logs) use std::net::SocketAddr; use std::sync::Arc; @@ -32,6 +41,7 @@ use arrow_flight::flight_service_server::FlightServiceServer; use ballista_scheduler::cluster::BallistaCluster; use ballista_scheduler::config::SchedulerConfig; use ballista_scheduler::scheduler_process::{create_scheduler, start_grpc_service}; +use comet_native_executors::adaptive::CometAdaptiveRules; use comet_native_executors::codec::CometPhysicalCodec; use comet_native_executors::frontend::CometFlightService; use comet_native_executors::{env_var, init_logging}; @@ -43,18 +53,32 @@ async fn main() -> Result<(), Box> { init_logging(); let mut config = SchedulerConfig { override_physical_codec: Some(Arc::new(CometPhysicalCodec::default())), + adaptive_rules: Some(Arc::new(CometAdaptiveRules)), ..Default::default() }; if let Some(host) = env_var("COMET_SCHEDULER_BIND_HOST")? { config.bind_host = host; } + if let Some(host) = env_var("COMET_SCHEDULER_EXTERNAL_HOST")? { + config.external_host = host; + } + if let Some(size) = env_var("COMET_GRPC_MAX_MESSAGE_SIZE")? { + config.grpc_server_max_decoding_message_size = size; + config.grpc_server_max_encoding_message_size = size; + config.grpc_client_max_message_size = size; + } if let Some(port) = env_var("COMET_SCHEDULER_PORT")? { config.bind_port = port; } if let Some(seconds) = env_var("COMET_SCHEDULER_RESULT_RETENTION_SECONDS")? { config.finished_job_data_clean_up_interval_seconds = seconds; } + if let Some(dir) = env_var::("COMET_SCHEDULER_EVENT_LOG_DIR")? { + config.event_log_dir = Some(dir); + } let flight_port: u16 = env_var("COMET_SCHEDULER_FLIGHT_PORT")?.unwrap_or(50060); + let ballista_settings = + parse_settings(&env_var::("COMET_SCHEDULER_BALLISTA_CONFIG")?.unwrap_or_default())?; let grpc_addr: SocketAddr = format!("{}:{}", config.bind_host, config.bind_port).parse()?; let flight_addr = SocketAddr::new(grpc_addr.ip(), flight_port); let config = Arc::new(config); @@ -68,7 +92,7 @@ async fn main() -> Result<(), Box> { let grpc = tokio::spawn(start_grpc_service(grpc_addr, scheduler.clone())); let flight = tonic::transport::Server::builder() .add_service( - FlightServiceServer::new(CometFlightService::new(scheduler)) + FlightServiceServer::new(CometFlightService::try_new(scheduler, &ballista_settings)?) .max_decoding_message_size(usize::MAX) .max_encoding_message_size(usize::MAX), ) @@ -79,3 +103,18 @@ async fn main() -> Result<(), Box> { } Ok(()) } + +/// `key=value` pairs separated by commas. +fn parse_settings(settings: &str) -> Result, String> { + settings + .split(',') + .map(str::trim) + .filter(|setting| !setting.is_empty()) + .map(|setting| { + setting + .split_once('=') + .map(|(key, value)| (key.trim().to_string(), value.trim().to_string())) + .ok_or_else(|| format!("Invalid Ballista setting {setting:?}: expected key=value")) + }) + .collect() +} diff --git a/native/native-executors/src/frontend.rs b/native/native-executors/src/frontend.rs index b6077669819..f457cd04d27 100644 --- a/native/native-executors/src/frontend.rs +++ b/native/native-executors/src/frontend.rs @@ -23,23 +23,34 @@ //! also says when the partition goes: Ballista deletes a query's results a set time after the //! query ends, whether or not they have been read (see the comet-scheduler binary). -use std::sync::Arc; -use std::time::{Duration, UNIX_EPOCH}; +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use arrow_flight::encode::FlightDataEncoderBuilder; use arrow_flight::flight_service_server::FlightService; use arrow_flight::{ Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightEndpoint, FlightInfo, HandshakeRequest, HandshakeResponse, PollInfo, PutResult, SchemaResult, Ticket, }; +use ballista_core::client::BallistaClient; +use ballista_core::config::{BallistaConfig, ExecutionMode}; +use ballista_core::error::BallistaError; use ballista_core::extension::SessionConfigExt; -use ballista_core::serde::protobuf::{job_status, PartitionLocation, SuccessfulJob}; +use ballista_core::serde::protobuf::{ + job_status, JobStatus, PartitionLocation, RunningJob, SuccessfulJob, +}; use ballista_core::serde::scheduler::{Action as BallistaAction, ShuffleFileKind}; use ballista_scheduler::scheduler_server::SchedulerServer; -use datafusion::physical_plan::ExecutionPlan; +use datafusion::arrow::array::RecordBatch; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use datafusion::prelude::SessionConfig; use datafusion_comet_proto::spark_native_executors::QueryRequest; use datafusion_proto::protobuf::{LogicalPlanNode, PhysicalPlanNode}; use futures::stream::BoxStream; +use futures::TryStreamExt; use log::info; use prost::Message; use prost_types::Timestamp; @@ -55,57 +66,68 @@ const SESSION_ID: &str = "comet"; /// wraps. pub struct CometFlightService { scheduler: SchedulerServer, + /// The results of pipelined queries, by job id, until they expire. + stored_results: Arc>>>, + /// The session every query runs in, with the cluster's Ballista settings. + session: SessionConfig, } impl CometFlightService { - pub fn new(scheduler: SchedulerServer) -> Self { - Self { scheduler } + /// A service that runs each query in a session with Ballista's defaults overridden by + /// `ballista_settings`, such as `ballista.execution.mode=pipelined`. + pub fn try_new( + scheduler: SchedulerServer, + ballista_settings: &[(String, String)], + ) -> Result { + let session = ballista_settings.iter().fold( + SessionConfig::new_with_ballista(), + |session, (key, value)| session.set_str(key, value), + ); + // Check the settings now rather than when the first query runs. + session + .ballista_config() + .execution_mode() + .map_err(|e| format!("Invalid Ballista settings: {e}"))?; + Ok(Self { + scheduler, + stored_results: Arc::default(), + session, + }) } - /// Runs `plan` to completion and returns its result: where its output partitions are, in - /// partition order, and when it ended. - async fn run(&self, plan: Arc) -> Result { - let ctx = self - .scheduler - .state - .session_manager - .create_or_update_session(SESSION_ID, &SessionConfig::new_with_ballista()) - .await - .map_err(|e| Status::internal(format!("Cannot create a session: {e}")))?; - let (sender, mut receiver) = tokio::sync::mpsc::channel(16); - let job_id = self - .scheduler - .submit_physical_plan("comet", ctx, plan, Some(sender)) - .await - .map_err(|e| Status::internal(format!("Cannot submit the query: {e}")))?; - info!("Submitted query as job {job_id}"); - while let Some(status) = receiver.recv().await { - match status.status { - Some(job_status::Status::Successful(mut job)) => { - job.partition_location.sort_by_key(|location| { - let partition = location.partition_id.as_ref().map(|p| p.partition_id); - (partition, location.map_partition_id) - }); - return Ok(job); - } - Some(job_status::Status::Failed(job)) => { - return Err(Status::internal(format!( - "Job {job_id} failed: {}", - job.error - ))) - } - _ => {} - } - } - Err(Status::internal(format!( - "Job {job_id} ended without a final status" - ))) + /// Whether queries run in Ballista's pipelined mode, which streams batches between stages + /// instead of writing shuffle files. + fn pipelined(&self) -> bool { + matches!( + self.session.ballista_config().execution_mode(), + Ok(ExecutionMode::Pipelined) + ) } /// When the executors delete the results of a query that ended at `ended_at`, in milliseconds /// since the epoch. Ballista deletes them that long after the query ends, unless the scheduler /// keeps them. Ballista starts that wait only after it records the end, so they are still /// there until this time. + /// Keeps a pipelined query's result until it expires, as long as a blocking job's. + fn store(&self, job_id: String, result: StoredResult) { + self.stored_results + .lock() + .unwrap() + .insert(job_id.clone(), Arc::new(result)); + let retention = self + .scheduler + .state + .config + .finished_job_data_clean_up_interval_seconds; + if retention > 0 { + let stored = Arc::clone(&self.stored_results); + tokio::spawn(async move { + tokio::time::sleep(Duration::from_secs(retention)).await; + stored.lock().unwrap().remove(&job_id); + }); + } + } + fn results_expire_at(&self, ended_at: u64) -> Option { let retention = self .scheduler @@ -122,6 +144,173 @@ impl CometFlightService { } } +/// Runs `plan` as a Ballista job on `scheduler` to completion and returns its result: where its +/// output partitions are, in partition order, and when it ended. +async fn run( + scheduler: &SchedulerServer, + session: &SessionConfig, + plan: Arc, +) -> Result { + let (job_id, mut receiver) = submit(scheduler, session, plan).await?; + while let Some(status) = receiver.recv().await { + match status.status { + Some(job_status::Status::Successful(mut job)) => { + job.partition_location.sort_by_key(|location| { + let partition = location.partition_id.as_ref().map(|p| p.partition_id); + (partition, location.map_partition_id) + }); + return Ok(job); + } + Some(job_status::Status::Failed(job)) => { + return Err(Status::internal(format!( + "Job {job_id} failed: {}", + job.error + ))) + } + _ => {} + } + } + Err(Status::internal(format!( + "Job {job_id} ended without a final status" + ))) +} + +/// Submits `plan` as a Ballista job on `scheduler`, and returns its id and the channel its status +/// changes arrive on. +async fn submit( + scheduler: &SchedulerServer, + session: &SessionConfig, + plan: Arc, +) -> Result<(String, tokio::sync::mpsc::Receiver), Status> { + let ctx = scheduler + .state + .session_manager + .create_or_update_session(SESSION_ID, session) + .await + .map_err(|e| Status::internal(format!("Cannot create a session: {e}")))?; + let (sender, receiver) = tokio::sync::mpsc::channel(16); + let job_id = scheduler + .submit_physical_plan("comet", ctx, plan, Some(sender)) + .await + .map_err(|e| Status::internal(format!("Cannot submit the query: {e}")))?; + info!("Submitted query as job {job_id}"); + Ok((job_id.to_string(), receiver)) +} + +/// Runs `plan` as a pipelined Ballista job and reads its results into the scheduler. +/// +/// A pipelined job's final stage streams each output partition to one reader, once, and can't +/// finish until it has been read. A Spark driver may read a result partition more than once, +/// since Spark can compute a partition of an RDD again, so the scheduler reads the streams itself +/// and serves them from its own Flight service for as long as a blocking job's results are kept. +/// +/// The result is gathered into one partition first. A sorted query's result is already one +/// partition, the merge of its sorted partitions. +async fn run_pipelined( + scheduler: &SchedulerServer, + session: &SessionConfig, + plan: Arc, +) -> Result<(String, StoredResult), Status> { + let plan: Arc = if plan.output_partitioning().partition_count() > 1 { + Arc::new(CoalescePartitionsExec::new(plan)) + } else { + plan + }; + let schema = plan.schema(); + let config = session.ballista_config(); + let (job_id, mut receiver) = submit(scheduler, session, plan).await?; + let failed = + |job_id: &str, error: &str| Status::internal(format!("Job {job_id} failed: {error}")); + let locations = loop { + let Some(status) = receiver.recv().await else { + return Err(Status::internal(format!( + "Job {job_id} ended without reporting where its results are" + ))); + }; + match status.status { + Some(job_status::Status::Running(RunningJob { + pipelined_output_locations, + .. + })) + | Some(job_status::Status::Successful(SuccessfulJob { + pipelined_output_locations, + .. + })) if !pipelined_output_locations.is_empty() => break pipelined_output_locations, + Some(job_status::Status::Failed(job)) => return Err(failed(&job_id, &job.error)), + _ => {} + } + }; + // Every partition at once: the final stage can't finish until all of them are read. + let partitions = futures::future::try_join_all( + locations + .into_iter() + .map(|location| fetch_pipelined(location, &config)), + ) + .await; + // The job's own failure explains a failed read better than the read does. + while let Some(status) = receiver.recv().await { + match status.status { + Some(job_status::Status::Successful(_)) => break, + Some(job_status::Status::Failed(job)) => return Err(failed(&job_id, &job.error)), + _ => {} + } + } + Ok(( + job_id, + StoredResult { + schema, + partitions: partitions?, + }, + )) +} + +/// Reads one result partition of a pipelined job from the executor that streams it. +async fn fetch_pipelined( + location: ballista_core::serde::protobuf::PipelinedPartitionLocation, + config: &BallistaConfig, +) -> Result, Status> { + let location: ballista_core::serde::scheduler::PipelinedPartitionLocation = location + .try_into() + .map_err(|e: BallistaError| Status::internal(format!("Bad result location: {e}")))?; + let executor = &location.executor_meta; + let failed = |e: BallistaError| { + Status::internal(format!( + "Cannot read a result from {}:{}: {e}", + executor.host, executor.port + )) + }; + let mut client = BallistaClient::try_new( + &executor.host, + executor.port, + config.grpc_client_max_message_size(), + false, + None, + config.io_retries_times() as u8, + config.io_retry_wait_time_ms() as u64, + 0, + 0, + ) + .await + .map_err(failed)?; + let stream = client + .fetch_pipelined_partition(&location, config.pipelined_registration_timeout()) + .await + .map_err(failed)?; + datafusion::physical_plan::common::collect(stream) + .await + .map_err(|e| failed(BallistaError::DataFusionError(Box::new(e)))) +} + +/// The results of a pipelined query, which the scheduler serves itself. +struct StoredResult { + schema: SchemaRef, + partitions: Vec>, +} + +/// Prefix of a ticket for a result partition the scheduler serves itself. The rest is the job id, +/// a NUL, and the partition's index. +const STORED_RESULT_MAGIC: &str = "COMET_RESULT\0"; + /// The Flight endpoint for a result partition: the executor's location, a ticket the executor's /// Flight service reads as a Ballista partition fetch, and when the partition is deleted. fn endpoint( @@ -181,15 +370,32 @@ impl FlightService for CometFlightService { let query = QueryRequest::decode(descriptor.cmd.as_ref()) .map_err(|e| Status::invalid_argument(format!("Cannot decode the query: {e}")))?; check_version(&query).map_err(|e| Status::failed_precondition(e.to_string()))?; + let pipelined = self.pipelined(); let plan = tokio::task::spawn_blocking(move || build_plan(query)) .await .map_err(|e| Status::internal(e.to_string()))? .map_err(|e| Status::invalid_argument(e.to_string()))?; - let job = self.run(plan).await?; - let expires_at = self.results_expire_at(job.ended_at); let mut info = FlightInfo::new(); - for location in &job.partition_location { - info = info.with_endpoint(endpoint(location, expires_at)?); + if pipelined { + let (job_id, result) = run_pipelined(&self.scheduler, &self.session, plan).await?; + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |now| now.as_millis() as u64); + let expires_at = self.results_expire_at(now); + // No location: the driver reads them from this service. + for index in 0..result.partitions.len() { + let ticket = format!("{STORED_RESULT_MAGIC}{job_id}\0{index}"); + let mut endpoint = FlightEndpoint::new().with_ticket(Ticket::new(ticket)); + endpoint.expiration_time = expires_at; + info = info.with_endpoint(endpoint); + } + self.store(job_id, result); + } else { + let job = run(&self.scheduler, &self.session, plan).await?; + let expires_at = self.results_expire_at(job.ended_at); + for location in &job.partition_location { + info = info.with_endpoint(endpoint(location, expires_at)?); + } } Ok(Response::new(info)) } @@ -222,11 +428,38 @@ impl FlightService for CometFlightService { unimplemented("GetSchema") } + /// Serves a result partition of a pipelined query. A blocking query's results are read from + /// the executors instead. async fn do_get( &self, - _request: Request, + request: Request, ) -> Result, Status> { - unimplemented("DoGet; fetch result partitions from the executors") + let ticket = request.into_inner().ticket; + let ticket = std::str::from_utf8(&ticket) + .ok() + .and_then(|ticket| ticket.strip_prefix(STORED_RESULT_MAGIC)) + .and_then(|ticket| ticket.split_once('\0')) + .and_then(|(job_id, index)| Some((job_id.to_string(), index.parse::().ok()?))); + let Some((job_id, index)) = ticket else { + return unimplemented("DoGet for this ticket; fetch it from an executor"); + }; + let result = self + .stored_results + .lock() + .unwrap() + .get(&job_id) + .cloned() + .ok_or_else(|| { + Status::not_found(format!("The results of job {job_id} have expired")) + })?; + let batches = result.partitions.get(index).cloned().ok_or_else(|| { + Status::not_found(format!("Job {job_id} has no result partition {index}")) + })?; + let stream = FlightDataEncoderBuilder::new() + .with_schema(Arc::clone(&result.schema)) + .build(futures::stream::iter(batches.into_iter().map(Ok))) + .map_err(Status::from); + Ok(Response::new(Box::pin(stream))) } async fn do_put( diff --git a/native/native-executors/src/lib.rs b/native/native-executors/src/lib.rs index 02251c83cb2..2457cf22cc7 100644 --- a/native/native-executors/src/lib.rs +++ b/native/native-executors/src/lib.rs @@ -28,6 +28,7 @@ use std::fmt::Display; use std::str::FromStr; +pub mod adaptive; pub mod assembly; pub mod codec; pub mod frontend; diff --git a/native/native-executors/tests/assembly.rs b/native/native-executors/tests/assembly.rs index 1776fc6a655..792bd203831 100644 --- a/native/native-executors/tests/assembly.rs +++ b/native/native-executors/tests/assembly.rs @@ -22,7 +22,10 @@ mod common; use comet_native_executors::assembly::build_plan; use common::*; +use datafusion::arrow::array::AsArray; use datafusion::arrow::datatypes::{DataType, Int32Type, Int64Type}; +use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion::physical_plan::{collect, ExecutionPlan, ExecutionPlanProperties}; use datafusion::prelude::SessionContext; use datafusion_comet_proto::spark_expression::{ @@ -76,17 +79,63 @@ async fn single_partition_exchange_collects_every_row() { } #[tokio::test(flavor = "multi_thread")] -async fn range_exchange_is_read_as_one_partition() { - // The driver asks for the exchange's partition count. +async fn range_exchange_without_an_ordering_is_read_as_one_partition() { + // A stage that declares no ordering has nothing to merge its partitions by. let request = query(vec![ range_stage(100, 4), - reading_stage(0, range_partitioning(4), 4), + reading_stage(0, range_partitioning(4, vec![]), 4), ]); let plan = assemble(request.clone()).await.unwrap(); assert_eq!(plan.output_partitioning().partition_count(), 1); assert_eq!(run(request).await, (0..100).collect::>()); } +#[tokio::test(flavor = "multi_thread")] +async fn range_exchange_is_sorted_in_each_partition_and_merged() { + let order = first_column_descending(long_type()); + let request = query(vec![ + range_stage(1000, 4), + sorting_stage(0, range_partitioning(4, vec![]), order, long_type()), + ]); + let plan = assemble(request).await.unwrap(); + let merge = plan + .downcast_ref::() + .expect("the stage's sorted partitions are merged"); + assert!(merge.expr()[0].options.descending); + // The stage reads the producer's four partitions without an exchange. + assert_eq!(merge.input().output_partitioning().partition_count(), 4); + let batches = collect(plan, SessionContext::new().task_ctx()) + .await + .unwrap(); + let values: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_primitive::() + .values() + .to_vec() + }) + .collect(); + assert_eq!(values, (0..1000).rev().collect::>()); +} + +/// A merge compares raw values, which order floats differently from Comet's sort. +#[tokio::test(flavor = "multi_thread")] +async fn range_exchange_sorted_by_a_float_is_read_as_one_partition() { + let mut producer = range_stage(1000, 4); + producer.output_schema = vec![field(double_type())]; + let order = first_column_descending(double_type()); + let request = query(vec![ + producer, + sorting_stage(0, range_partitioning(4, vec![]), order, double_type()), + ]); + let plan = assemble(request).await.unwrap(); + assert!(!plan.is::()); + assert!(plan.children()[1].is::()); + assert_eq!(plan.output_partitioning().partition_count(), 1); +} + #[tokio::test(flavor = "multi_thread")] async fn partition_plans_must_match_the_partitions() { let mut stage = range_stage(100, 4); diff --git a/native/native-executors/tests/cluster.rs b/native/native-executors/tests/cluster.rs index dee8dcaa7b5..5cd539a27cd 100644 --- a/native/native-executors/tests/cluster.rs +++ b/native/native-executors/tests/cluster.rs @@ -30,7 +30,7 @@ use std::time::{Duration, Instant, SystemTime}; use arrow_flight::error::FlightError; use arrow_flight::{FlightClient, FlightDescriptor, FlightEndpoint, FlightInfo}; use common::*; -use datafusion::arrow::array::RecordBatch; +use datafusion::arrow::array::{AsArray, RecordBatch}; use datafusion::arrow::datatypes::Int64Type; use datafusion_comet_proto::spark_native_executors::QueryRequest; use futures::TryStreamExt; @@ -43,6 +43,8 @@ struct Cluster { processes: Vec, work_dir: PathBuf, flight_url: String, + /// The scheduler's log, with this crate's messages at info level. + scheduler_log: PathBuf, } /// Held while a cluster starts, so that two tests don't pick the same free port in between. @@ -67,6 +69,10 @@ impl Cluster { std::process::id() )), flight_url: format!("http://127.0.0.1:{flight_port}"), + scheduler_log: std::env::temp_dir().join(format!( + "comet-scheduler-{}-{flight_port}.log", + std::process::id() + )), }; let mut env = vec![ ("COMET_SCHEDULER_BIND_HOST", "127.0.0.1".to_string()), @@ -74,7 +80,13 @@ impl Cluster { ("COMET_SCHEDULER_FLIGHT_PORT", flight_port.to_string()), ]; env.extend_from_slice(scheduler_env); - cluster.spawn(env!("CARGO_BIN_EXE_comet-scheduler"), &env); + let log = std::fs::File::create(&cluster.scheduler_log).unwrap(); + cluster.spawn_logging( + env!("CARGO_BIN_EXE_comet-scheduler"), + &env, + "warn,comet_native_executors=info", + Stdio::from(log), + ); wait_for(scheduler_port); wait_for(flight_port); let work_dir = cluster.work_dir.to_string_lossy().into_owned(); @@ -95,10 +107,15 @@ impl Cluster { } fn spawn(&mut self, binary: &str, env: &[(&str, String)]) { + self.spawn_logging(binary, env, "warn", Stdio::inherit()); + } + + fn spawn_logging(&mut self, binary: &str, env: &[(&str, String)], log: &str, stderr: Stdio) { let child = Command::new(binary) .envs(env.iter().map(|(key, value)| (*key, value.as_str()))) - .env("RUST_LOG", "warn") + .env("RUST_LOG", log) .stdout(Stdio::null()) + .stderr(stderr) .spawn() .unwrap_or_else(|e| panic!("Cannot start {binary}: {e}")); self.processes.push(child); @@ -138,6 +155,7 @@ impl Drop for Cluster { let _ = process.wait(); } let _ = std::fs::remove_dir_all(&self.work_dir); + let _ = std::fs::remove_file(&self.scheduler_log); } } @@ -197,7 +215,7 @@ async fn runs_a_query_across_a_range_exchange() { let cluster = Cluster::start(); let request = query(vec![ range_stage(1000, 4), - reading_stage(0, range_partitioning(4), 4), + reading_stage(0, range_partitioning(4, vec![]), 4), ]); let partitions = cluster.run(&request).await; assert_eq!(partitions.len(), 1); @@ -207,6 +225,87 @@ async fn runs_a_query_across_a_range_exchange() { ); } +/// A global sort runs as a sort of each partition and a merge, so its result comes back as one +/// partition, in order. +#[tokio::test(flavor = "multi_thread")] +async fn sorts_across_a_range_exchange_by_merging_sorted_partitions() { + let cluster = Cluster::start(); + let order = first_column_descending(long_type()); + let request = query(vec![ + range_stage(1000, 4), + sorting_stage(0, range_partitioning(4, vec![]), order, long_type()), + ]); + let partitions = cluster.run(&request).await; + assert_eq!(partitions.len(), 1); + let values: Vec = partitions[0] + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_primitive::() + .values() + .to_vec() + }) + .collect(); + assert_eq!(values, (0..1000).rev().collect::>()); +} + +#[tokio::test(flavor = "multi_thread")] +async fn broadcasts_a_build_side_that_measures_small_into_the_probe_side_stage() { + // Spark's estimates make the build side worth measuring before the probe side is shuffled: + // under Ballista's maximum estimate multiple, and the probe side many times larger. It + // measures far under the broadcast threshold, so the join runs in the probe side's stage. + let cluster = Cluster::start(); + let mut probe = range_stage(10_000, 4); + probe.estimated_output_bytes = Some(1 << 40); + let mut build = range_stage(100, 2); + build.estimated_output_bytes = Some(1 << 30); + let request = query(vec![probe, build, joining_stage(0, 1, 3)]); + + let partitions = cluster.run(&request).await; + + let rows = partitions.concat(); + assert_eq!( + sorted_values::(&rows), + (0..100).collect::>() + ); + for batch in &rows { + assert_eq!(batch.column(0), batch.column(1)); + } + let log = std::fs::read_to_string(&cluster.scheduler_log).unwrap(); + assert!(log.contains("Holding back a join's probe side"), "{log}"); + assert!( + log.contains("broadcasting it to the probe side's stage"), + "{log}" + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn shuffles_the_probe_side_when_the_build_side_measures_large() { + // With a broadcast threshold of one byte, a held join's build side always measures too + // large, so the probe side's stage is released as it was and the join is shuffled. + let cluster = Cluster::start_with(&[( + "COMET_SCHEDULER_BALLISTA_CONFIG", + "ballista.optimizer.broadcast_join_threshold_bytes=1".to_string(), + )]); + let mut probe = range_stage(10_000, 4); + probe.estimated_output_bytes = Some(1 << 40); + let mut build = range_stage(100, 2); + build.estimated_output_bytes = Some(16); + let request = query(vec![probe, build, joining_stage(0, 1, 3)]); + + let partitions = cluster.run(&request).await; + + let rows = partitions.concat(); + assert_eq!( + sorted_values::(&rows), + (0..100).collect::>() + ); + let log = std::fs::read_to_string(&cluster.scheduler_log).unwrap(); + assert!(log.contains("Holding back a join's probe side"), "{log}"); + assert!(log.contains("shuffling the probe side"), "{log}"); +} + #[tokio::test(flavor = "multi_thread")] async fn refuses_a_query_from_another_build() { let cluster = Cluster::start(); diff --git a/native/native-executors/tests/common/mod.rs b/native/native-executors/tests/common/mod.rs index 5dc6b530cf7..b7e1b3acaa4 100644 --- a/native/native-executors/tests/common/mod.rs +++ b/native/native-executors/tests/common/mod.rs @@ -25,10 +25,13 @@ use comet::execution::native_executors::COMET_VERSION; use datafusion::arrow::array::{AsArray, RecordBatch}; use datafusion::arrow::datatypes::ArrowPrimitiveType; use datafusion_comet_proto::spark_expression::data_type::DataTypeId; -use datafusion_comet_proto::spark_expression::{self, expr::ExprStruct, BoundReference, Expr}; +use datafusion_comet_proto::spark_expression::{ + self, expr::ExprStruct, BoundReference, Expr, NullOrdering, SortDirection, SortOrder, +}; use datafusion_comet_proto::spark_native_executors::{QueryRequest, Stage, StageInput}; use datafusion_comet_proto::spark_operator::{ - operator::OpStruct, Operator, RangeScan, Scan, SparkStructField, + operator::OpStruct, BuildSide, HashJoin, JoinType, Operator, RangeScan, Scan, ShuffleScan, + Sort, SparkStructField, }; use datafusion_comet_proto::spark_partitioning::{ partitioning::PartitioningStruct, HashPartition, Partitioning, RangePartition, SinglePartition, @@ -60,6 +63,10 @@ pub fn int_type() -> spark_expression::DataType { data_type(DataTypeId::Int32) } +pub fn double_type() -> spark_expression::DataType { + data_type(DataTypeId::Double) +} + pub fn bool_type() -> spark_expression::DataType { data_type(DataTypeId::Bool) } @@ -125,6 +132,43 @@ pub fn reading_stage(producer: u32, partitioning: Partitioning, num_partitions: } } +/// A stage that sorts what stage `producer` sends it through an exchange with `partitioning`, by +/// `order`, of a column of type `data_type`, and declares that order, as a driver does for a +/// stage that reads a range exchange. +pub fn sorting_stage( + producer: u32, + partitioning: Partitioning, + order: Expr, + data_type: spark_expression::DataType, +) -> Stage { + let scan = Operator { + op_struct: Some(OpStruct::Scan(Scan { + fields: vec![data_type.clone()], + source: "test".to_string(), + })), + ..Default::default() + }; + Stage { + plan: Some(Operator { + children: vec![scan], + op_struct: Some(OpStruct::Sort(Sort { + sort_orders: vec![order.clone()], + ..Default::default() + })), + ..Default::default() + }), + inputs: vec![StageInput { + producer, + partitioning: Some(partitioning), + broadcast: false, + }], + num_partitions: 1, + output_schema: vec![field(data_type)], + output_ordering: vec![order], + ..Default::default() + } +} + pub fn hash_on_first_column(num_partitions: i32) -> Partitioning { Partitioning { partitioning_struct: Some(PartitioningStruct::HashPartition(HashPartition { @@ -140,16 +184,29 @@ pub fn single_partition() -> Partitioning { } } -/// A range exchange as a driver sends it: its sort keys don't matter here, and it has no bounds. -pub fn range_partitioning(num_partitions: i32) -> Partitioning { +/// A range exchange as a driver sends it, with no bounds: sorting by `sort_orders`. +pub fn range_partitioning(num_partitions: i32, sort_orders: Vec) -> Partitioning { Partitioning { partitioning_struct: Some(PartitioningStruct::RangePartition(RangePartition { + sort_orders, num_partitions, ..Default::default() })), } } +/// The first column, of type `data_type`, descending, as a sort key. +pub fn first_column_descending(data_type: spark_expression::DataType) -> Expr { + Expr { + expr_struct: Some(ExprStruct::SortOrder(Box::new(SortOrder { + child: Some(Box::new(bound(0, data_type))), + direction: SortDirection::Descending as i32, + null_ordering: NullOrdering::NullsLast as i32, + }))), + ..Default::default() + } +} + /// The values of the batches' single column, of type `T`, sorted. pub fn sorted_values(batches: &[RecordBatch]) -> Vec where @@ -162,3 +219,38 @@ where values.sort_unstable(); values } + +/// A stage that joins what stages `probe` and `build` send it through hash exchanges of +/// `num_partitions` partitions on their first column, as an inner hash join built on the build +/// side, and returns both sides' columns. +pub fn joining_stage(probe: u32, build: u32, num_partitions: i32) -> Stage { + let leaf = || Operator { + op_struct: Some(OpStruct::ShuffleScan(ShuffleScan { + fields: vec![long_type()], + source: "test".to_string(), + })), + ..Default::default() + }; + let input = |producer| StageInput { + producer, + partitioning: Some(hash_on_first_column(num_partitions)), + broadcast: false, + }; + Stage { + plan: Some(Operator { + children: vec![leaf(), leaf()], + op_struct: Some(OpStruct::HashJoin(HashJoin { + left_join_keys: vec![bound(0, long_type())], + right_join_keys: vec![bound(0, long_type())], + join_type: JoinType::Inner as i32, + build_side: BuildSide::BuildRight as i32, + ..Default::default() + })), + ..Default::default() + }), + inputs: vec![input(probe), input(build)], + num_partitions: num_partitions as u32, + output_schema: vec![field(long_type()), field(long_type())], + ..Default::default() + } +} diff --git a/native/proto/src/proto/native_executors.proto b/native/proto/src/proto/native_executors.proto index 6c0ba50afb4..8fdb2261be4 100644 --- a/native/proto/src/proto/native_executors.proto +++ b/native/proto/src/proto/native_executors.proto @@ -23,6 +23,7 @@ syntax = "proto3"; package spark.spark_native_executors; +import "expr.proto"; import "operator.proto"; import "partitioning.proto"; @@ -58,6 +59,14 @@ message Stage { // leaves the stage, as a Spark task's shuffle writer does, because a hash exchange places // rows by the types of their values too. repeated spark.spark_operator.SparkStructField output_schema = 5; + // For a stage that reads a range exchange, the order its block sorts each partition by, as + // SortOrder expressions over the block's output. Rather than range-partition the exchange, + // which needs bounds Spark would sample for, the cluster reads the producer's partitions as + // they are, sorts each, and merges them in this order. + repeated spark.spark_expression.Expr output_ordering = 6; + // Spark's estimate of the stage's output size in bytes, when it has one. The scheduler reads + // it to decide whether measuring a join's build side before shuffling its probe side pays off. + optional uint64 estimated_output_bytes = 7; } // How a scheduler sends a stage to its executors, inside the task plans Ballista ships. The @@ -70,6 +79,8 @@ message Fragment { // Whether every partition reads all of each input, one entry per input. repeated bool broadcast = 3; map conf = 4; + // Spark's estimate of the stage's output size in bytes, from Stage.estimated_output_bytes. + optional uint64 estimated_output_bytes = 5; } message StageInput { @@ -80,3 +91,17 @@ message StageInput { // Every partition of this stage reads the producer's whole output. bool broadcast = 3; } + +// The result partitions of a query, in order. The driver's native library converts the +// endpoints of the FlightInfo the scheduler returns into this message for the JVM. +message ResultEndpoints { + repeated ResultEndpoint endpoints = 1; +} + +message ResultEndpoint { + // Arrow Flight location of the executor that holds the partition, such as + // grpc+tcp://host:port. + string location = 1; + // Opaque ticket to pass to DoGet. + bytes ticket = 2; +} diff --git a/native/spark-expr/src/agg_funcs/avg.rs b/native/spark-expr/src/agg_funcs/avg.rs index 2b2422ce775..c44d30eda2a 100644 --- a/native/spark-expr/src/agg_funcs/avg.rs +++ b/native/spark-expr/src/agg_funcs/avg.rs @@ -334,12 +334,33 @@ where ]) } + /// Each row's own partial state, as `state` emits it: the value and a count of one, or a + /// zero sum and count for a null or filtered-out row. fn convert_to_state( &self, - _values: &[ArrayRef], - _opt_filter: Option<&arrow::array::BooleanArray>, + values: &[ArrayRef], + opt_filter: Option<&arrow::array::BooleanArray>, ) -> Result> { - not_impl_err!("Input batch conversion to state not implemented") + let values = values[0].as_primitive::(); + let mut sums = Vec::with_capacity(values.len()); + let mut counts = Vec::with_capacity(values.len()); + for idx in 0..values.len() { + let included = values.is_valid(idx) + && opt_filter.is_none_or(|filter| filter.is_valid(idx) && filter.value(idx)); + if included { + sums.push(values.value(idx)); + counts.push(1); + } else { + sums.push(T::default_value()); + counts.push(0); + } + } + let sums = PrimitiveArray::::new(sums.into(), None) + .with_data_type(self.return_data_type.clone()); + Ok(vec![ + Arc::new(sums) as ArrayRef, + Arc::new(Int64Array::new(counts.into(), None)) as ArrayRef, + ]) } fn size(&self) -> usize { @@ -365,4 +386,32 @@ mod tests { assert_eq!(partial.evaluate()?, ScalarValue::Float64(None)); Ok(()) } + + /// Converting rows straight to partial states, as a partial aggregate that stops aggregating + /// does, gives each row the state that aggregating it alone into a group would. + #[test] + fn convert_to_state_matches_one_row_groups() -> Result<()> { + use arrow::array::BooleanArray; + let values: ArrayRef = Arc::new(Float64Array::from(vec![ + Some(1.5), + None, + Some(-2.0), + Some(4.0), + ])); + let filter = BooleanArray::from(vec![Some(true), Some(true), Some(true), Some(false)]); + let new_accumulator = || { + AvgGroupsAccumulator::::new( + &Float64, + |sum: f64, count: i64| Ok(sum / count as f64), + ) + }; + + let converted = + new_accumulator().convert_to_state(&[Arc::clone(&values)], Some(&filter))?; + + let mut grouped = new_accumulator(); + grouped.update_batch(&[values], &[0, 1, 2, 3], Some(&filter), 4)?; + assert_eq!(converted, grouped.state(EmitTo::All)?); + Ok(()) + } } diff --git a/pom.xml b/pom.xml index a215e2a818f..e510c08eb66 100644 --- a/pom.xml +++ b/pom.xml @@ -1010,7 +1010,7 @@ under the License. -Ywarn-value-discard -Ywarn-unused:imports,patvars,privates,locals,params,-implicits - -Wconf:cat=unused-params&site=org\.apache\.comet\.Native\..*:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.Native(ExecutorsClient)?\..*:s -Wconf:cat=unused-params&src=.*/src/[a-z]+/spark-[^/]+/.*:s -Wconf:cat=unused-params&site=org\.apache\.comet\.rules\.CometScanContrib\.tryTransformV[12]\..*:s @@ -1055,7 +1055,7 @@ under the License. -Ywarn-numeric-widen -Ywarn-unused:imports,patvars,privates,locals,params,-implicits - -Wconf:cat=unused-params&site=org\.apache\.comet\.Native\..*:s + -Wconf:cat=unused-params&site=org\.apache\.comet\.Native(ExecutorsClient)?\..*:s -Wconf:cat=unused-params&src=.*/src/[a-z]+/spark-[^/]+/.*:s -Wconf:cat=unused-params&site=org\.apache\.comet\.rules\.CometScanContrib\.tryTransformV[12]\..*:s diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 824118f3efe..bc042152c53 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -510,6 +510,28 @@ object CometConf extends ShimCometConf { .checkValue(v => v > 0 && v <= Int.MaxValue, "Must be positive and fit in an Int.") .createWithDefault(1024 * 1024) + val COMET_NATIVE_EXECUTORS_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.nativeExecutors.enabled") + .category(CATEGORY_EXEC) + .doc( + "Experimental. When enabled, Comet runs each query on a Comet Native Executor cluster " + + "instead of on Spark's executors: the driver plans the query and submits it to the " + + "cluster's scheduler (`spark.comet.nativeExecutors.scheduler.url`), and the results " + + "come back to the driver. A query that cannot run entirely natively fails, naming " + + "the operator that cannot. See https://github.com/apache/datafusion-comet/issues/7288.") + .internal() + .booleanConf + .createWithDefault(false) + + val COMET_NATIVE_EXECUTORS_SCHEDULER_URL: ConfigEntry[String] = + conf("spark.comet.nativeExecutors.scheduler.url") + .category(CATEGORY_EXEC) + .doc("Experimental. The Arrow Flight URL of a Comet Native Executor cluster's scheduler, " + + "such as `http://host:50060`. See `spark.comet.nativeExecutors.enabled`.") + .internal() + .stringConf + .createWithDefault("http://localhost:50060") + val COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.columnarToRow.native.enabled") .category(CATEGORY_EXEC) @@ -710,8 +732,8 @@ object CometConf extends ShimCometConf { "Experimental opt-in: let a native partial aggregate stop aggregating once its input " + "looks mostly distinct, and send the rest of the task's rows to the shuffle " + "unaggregated. Only applies to partial aggregates that feed Comet native shuffle " + - "and whose aggregate functions, if any, are all single-argument COUNT. The check " + - "starts after the first 100,000 input rows of a task, and once aggregation stops " + + "and whose aggregate functions, if any, are all single-argument COUNT or AVG. The " + + "check starts after the first 100,000 input rows of a task, and once aggregation stops " + "it does not resume, so a task whose keys repeat after a mostly distinct start can " + s"shuffle many times more rows than it would with this disabled. $TUNING_GUIDE.") .booleanConf diff --git a/spark/src/main/scala/org/apache/comet/NativeExecutorsClient.scala b/spark/src/main/scala/org/apache/comet/NativeExecutorsClient.scala new file mode 100644 index 00000000000..fc4d3c98d9c --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/NativeExecutorsClient.scala @@ -0,0 +1,62 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet + +/** + * The JNI client a Spark driver uses to run queries on a Comet Native Executor cluster. See + * https://github.com/apache/datafusion-comet/issues/7288. + * + * The native methods exist only in a libcomet built with the `native-executors` Cargo feature. + */ +class NativeExecutorsClient extends NativeBase { + + /** + * Runs a query on the cluster and waits for it to finish. + * + * @param schedulerUrl + * the Arrow Flight URL of the cluster's scheduler + * @param query + * a serialized `QueryRequest`, whose Comet version and plan format the native library fills + * in + * @return + * a serialized `ResultEndpoints`: where each result partition is, in order + */ + @native def submit(schedulerUrl: String, query: Array[Byte]): Array[Byte] + + /** + * Starts reading a result partition. + * + * @return + * a handle for [[nextBatch]] and [[closePartition]] + */ + @native def openPartition(location: String, ticket: Array[Byte]): Long + + /** + * Moves the partition's next batch into the Arrow C Data Interface structs at `arrayAddrs` and + * `schemaAddrs`, one pair per column. + * + * @return + * the batch's row count, or -1 at the end of the partition + */ + @native def nextBatch(handle: Long, arrayAddrs: Array[Long], schemaAddrs: Array[Long]): Long + + /** Stops reading a partition and releases its handle. */ + @native def closePartition(handle: Long): Unit +} diff --git a/spark/src/main/scala/org/apache/comet/rules/CometRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometRule.scala index fbde4a79090..07b2eea6863 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometRule.scala @@ -25,6 +25,7 @@ import scala.util.control.NonFatal import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.trees.TreeNodeTag +import org.apache.spark.sql.comet.CometNativeExecutors import org.apache.spark.sql.execution.{ApplyColumnarRulesAndInsertTransitions, BaseSubqueryExec, ColumnarToRowExec, ExecSubqueryExpression, InputAdapter, QueryExecution, ReusedSubqueryExec, RowToColumnarExec, SparkPlan, WholeStageCodegenExec} import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, InsertAdaptiveSparkPlan, QueryStageExec} import org.apache.spark.sql.execution.exchange.{BroadcastExchangeExec, Exchange} @@ -163,6 +164,10 @@ case class CometRule(session: SparkSession, queryStagePrep: Boolean = false) if (CometRule.planOnlyApplies(conf, plan)) { reportPlanOnlyCoverage(plan) plan + } else if (CometNativeExecutors.applies(conf, plan, queryStagePrep)) { + // The whole query runs on a Comet Native Executor cluster. The conversion alone leaves + // every native block serialized, and nothing reverts a stage to Spark. + CometNativeExecutors.offload(plan, convert(plan)) } else { convert(plan) } diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeExecutorsExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeExecutorsExec.scala new file mode 100644 index 00000000000..2f024b706d9 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeExecutorsExec.scala @@ -0,0 +1,558 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.spark.sql.comet + +import java.util.IdentityHashMap + +import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ +import scala.util.Try + +import org.apache.spark.{Partition, SparkContext, TaskContext} +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, Literal} +import org.apache.spark.sql.catalyst.plans.physical.{HashPartitioning, RangePartitioning, SinglePartition} +import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec} +import org.apache.spark.sql.execution.{ColumnarToRowExec, CommandResultExec, DeserializeToObjectExec, LeafExecNode, ScalarSubquery, SparkPlan} +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, QueryStageExec} +import org.apache.spark.sql.execution.command.{DataWritingCommandExec, ExecutedCommandExec} +import org.apache.spark.sql.execution.datasources.v2.V2CommandExec +import org.apache.spark.sql.execution.exchange.{Exchange, ReusedExchangeExec} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.StructField +import org.apache.spark.sql.vectorized.ColumnarBatch + +import com.google.protobuf.ByteString +import com.google.protobuf.Descriptors.FieldDescriptor +import com.google.protobuf.Message + +import org.apache.comet.{CometConf, CometExecIterator, CometExplainInfo, CometRuntimeException, NativeExecutorsClient} +import org.apache.comet.serde.{ExprOuterClass, PartitioningOuterClass, QueryPlanSerde} +import org.apache.comet.serde.Config.ConfigMap +import org.apache.comet.serde.NativeExecutors.{QueryRequest, ResultEndpoint, ResultEndpoints, Stage, StageInput} +import org.apache.comet.serde.OperatorOuterClass.Operator +import org.apache.comet.serde.operator.schema2Proto +import org.apache.comet.vector.NativeUtil + +/** + * Runs a whole query on a Comet Native Executor cluster instead of on Spark's executors. See + * https://github.com/apache/datafusion-comet/issues/7288. + * + * This leaf replaces the plan Comet would otherwise execute. `doExecuteColumnar` submits the + * query to the cluster's scheduler and waits for it on the driver; each partition of the RDD it + * returns then reads one result partition from the executor that holds it. + * + * @param query + * the serialized `QueryRequest` built by [[CometNativeExecutors.buildQuery]] + * @param subqueryResults + * the query's scalar subqueries. They are expressions of this node, so Spark runs them before + * it executes the node, and their results replace them in the query before it is submitted. + */ +case class CometNativeExecutorsExec( + override val output: Seq[Attribute], + query: Array[Byte], + schedulerUrl: String, + subqueryResults: Seq[Expression] = Nil) + extends LeafExecNode { + + override def supportsColumnar: Boolean = true + + override protected def doExecute(): RDD[InternalRow] = + throw new UnsupportedOperationException(s"$nodeName only produces columnar output") + + override protected def doExecuteColumnar(): RDD[ColumnarBatch] = { + val request = CometNativeExecutors.withSubqueryResults( + query, + subqueryResults.collect { case sub: ScalarSubquery => sub }) + val endpoints = + ResultEndpoints.parseFrom(new NativeExecutorsClient().submit(schedulerUrl, request)) + new CometNativeExecutorsRDD( + sparkContext, + endpoints.getEndpointsList.asScala.toSeq, + output.length) + } + + override def stringArgs: Iterator[Any] = + Iterator(output, s"${query.length} bytes", schedulerUrl) +} + +private class CometNativeExecutorsPartition(override val index: Int, val endpoint: ResultEndpoint) + extends Partition + +/** + * The result partitions of a query that ran on a Comet Native Executor cluster, each read from + * the executor that holds it. + */ +private class CometNativeExecutorsRDD( + sc: SparkContext, + endpoints: Seq[ResultEndpoint], + numColumns: Int) + extends RDD[ColumnarBatch](sc, Nil) { + + override protected def getPartitions: Array[Partition] = + endpoints.zipWithIndex.map { case (endpoint, index) => + new CometNativeExecutorsPartition(index, endpoint): Partition + }.toArray + + override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = { + val partition = split.asInstanceOf[CometNativeExecutorsPartition] + val client = new NativeExecutorsClient() + val nativeUtil = new NativeUtil() + val handle = + client.openPartition( + partition.endpoint.getLocation, + partition.endpoint.getTicket.toByteArray) + val batches = new Iterator[ColumnarBatch] with AutoCloseable { + private var current: Option[ColumnarBatch] = None + private var pending: Option[ColumnarBatch] = None + private var closed = false + + override def hasNext: Boolean = { + if (pending.isEmpty && !closed) { + // The consumer has moved past the batch it last got. + current.foreach(_.close()) + current = None + pending = nativeUtil.getNextBatch( + numColumns, + (arrays, schemas) => client.nextBatch(handle, arrays, schemas)) + if (pending.isEmpty) { + close() + } + } + pending.isDefined + } + + override def next(): ColumnarBatch = { + if (!hasNext) { + throw new NoSuchElementException("No more batches in this result partition") + } + current = pending + pending = None + current.get + } + + override def close(): Unit = if (!closed) { + closed = true + current.foreach(_.close()) + current = None + pending.foreach(_.close()) + pending = None + client.closePartition(handle) + nativeUtil.close() + } + } + context.addTaskCompletionListener[Unit](_ => batches.close()) + batches + } +} + +/** + * Plans queries for Comet Native Executors: turns the plan Comet would execute into the stage + * graph a cluster runs, and refuses any query that cannot run entirely natively. + */ +object CometNativeExecutors { + + /** + * Messages of native plans that call back into the JVM, which a cluster doesn't have. A scalar + * subquery's `Subquery` is not one of them: the driver runs the subquery and puts its result in + * the plan before submitting it. The native planner could refuse them all instead + * (https://github.com/apache/datafusion-comet/issues/7375). + */ + private val jvmCallbacks = Set("JvmScalarUdf", "ArrowPythonUdf") + + /** + * Whether Comet runs `plan` on a Comet Native Executor cluster. It must be a whole query: not a + * command, not a stage AQE has split off, and not a plan already replaced. + */ + def applies(conf: SQLConf, plan: SparkPlan, queryStagePrep: Boolean): Boolean = + CometConf.COMET_NATIVE_EXECUTORS_ENABLED.get(conf) && + !(conf.adaptiveExecutionEnabled && !queryStagePrep && plan.isInstanceOf[Exchange]) && + !plan.exists { + case _: CometNativeExecutorsExec | _: AdaptiveSparkPlanExec | _: QueryStageExec | + _: CommandResultExec | _: ExecutedCommandExec | _: DataWritingCommandExec | + _: V2CommandExec => + true + case _ => false + } + + /** + * Replaces `converted`, the plan Comet would execute for `original`, with a leaf that runs it + * on a Comet Native Executor cluster. An operator on top that only turns the query's result + * into JVM objects, as `Dataset.rdd` adds, stays on the driver above the leaf. + */ + def offload(original: SparkPlan, converted: SparkPlan): SparkPlan = converted match { + case toObjects: DeserializeToObjectExec => + toObjects.withNewChildren(Seq(offload(toObjects.child, toObjects.child))) + case native => + val exec = CometNativeExecutorsExec( + native.output, + buildQuery(native).toByteArray, + CometConf.COMET_NATIVE_EXECUTORS_SCHEDULER_URL.get(), + scalarSubqueries(native)) + native.logicalLink.orElse(original.logicalLink).foreach(exec.setLogicalLink) + exec + } + + /** + * The query a Comet Native Executor cluster runs for `plan`, the plan Comet would execute: one + * stage per native block, with an input for each exchange that feeds the block. + * + * @throws CometRuntimeException + * naming the first operator that cannot run on a cluster + */ + def buildQuery(plan: SparkPlan): QueryRequest = { + val builder = new StageGraphBuilder(plan) + val root = plan match { + case ColumnarToRowExec(child) => child + case CometColumnarToRowExec(child) => child + case CometNativeColumnarToRowExec(child) => child + case other => other + } + builder.stageOf(root) + QueryRequest + .newBuilder() + .putAllConf(nativeConf()) + .addAllStages(builder.stages.asJava) + .build() + } + + /** The scalar subqueries of `plan`'s expressions, through exchanges, each once. */ + private def scalarSubqueries(plan: SparkPlan): Seq[Expression] = { + val found = ArrayBuffer.empty[ScalarSubquery] + def walk(node: SparkPlan): Unit = { + node.expressions.foreach(_.foreach { + case sub: ScalarSubquery if !found.exists(_.exprId == sub.exprId) => found += sub + case _ => + }) + node match { + case reused: ReusedExchangeExec => walk(reused.child) + case _ => node.children.foreach(walk) + } + } + walk(plan) + found.toSeq + } + + /** + * `query` with each scalar subquery's `Subquery` expression replaced by a literal of its + * result. The subqueries must have run. Sending the results alongside the plan would spare + * rewriting it (https://github.com/apache/datafusion-comet/issues/7374). + */ + def withSubqueryResults(query: Array[Byte], subqueries: Seq[ScalarSubquery]): Array[Byte] = { + if (subqueries.isEmpty) { + return query + } + val literals: Map[Long, ExprOuterClass.Expr] = subqueries.map { sub => + // The result is already in Catalyst's internal form. + val literal = Literal(sub.eval(null), sub.dataType) + val expr = QueryPlanSerde + .exprToProto(literal, Seq.empty) + .getOrElse(throw new CometRuntimeException( + s"Comet Native Executors cannot pass the result of $sub to the cluster")) + sub.exprId.id -> expr + }.toMap + val replace: PartialFunction[Message, Message] = { + case expr: ExprOuterClass.Expr if expr.hasSubquery => + literals.getOrElse( + expr.getSubquery.getId, + throw new CometRuntimeException( + s"Comet Native Executors found no result for subquery ${expr.getSubquery.getId}")) + } + val request = QueryRequest.parseFrom(query).toBuilder + (0 until request.getStagesCount).foreach { i => + val stage = request.getStages(i) + val plan = transformMessages(stage.getPlan)(replace) + // A partition's plan is the stage's with its files added, so it holds a subquery only + // when the stage's plan does. + if (plan ne stage.getPlan) { + val replaced = stage.toBuilder.setPlan(plan.asInstanceOf[Operator]) + (0 until stage.getPartitionPlansCount).foreach { p => + val partitionPlan = Operator.parseFrom(stage.getPartitionPlans(p)) + replaced.setPartitionPlans(p, transformMessages(partitionPlan)(replace).toByteString) + } + request.setStages(i, replaced) + } + } + request.build().toByteArray + } + + /** + * The messages directly inside `message`, each with how to put a replacement for it in a + * builder of `message`. + */ + private def childMessages( + message: Message): Iterator[(Message, (Message.Builder, Message) => Unit)] = + message.getAllFields.asScala.iterator + .filter { case (field, _) => field.getJavaType == FieldDescriptor.JavaType.MESSAGE } + .flatMap { case (field, value) => + if (field.isRepeated) { + value.asInstanceOf[java.util.List[Message]].asScala.iterator.zipWithIndex.map { + case (child, index) => + ( + child, + (builder: Message.Builder, m: Message) => { + builder.setRepeatedField(field, index, m) + () + }) + } + } else { + Iterator( + ( + value.asInstanceOf[Message], + (builder: Message.Builder, m: Message) => { + builder.setField(field, m) + () + })) + } + } + + /** + * `message` with every message inside it that `pf` matches replaced by what `pf` returns. A + * message nothing inside changed is returned as is. + */ + private def transformMessages(message: Message)( + pf: PartialFunction[Message, Message]): Message = + pf.applyOrElse( + message, + (m: Message) => { + var builder: Message.Builder = null + childMessages(m).foreach { case (child, set) => + val replaced = transformMessages(child)(pf) + if (replaced ne child) { + if (builder == null) { + builder = m.toBuilder + } + set(builder, replaced) + } + } + if (builder == null) m else builder.build() + }) + + /** The configs a Spark task would pass to native code with a plan. */ + private def nativeConf(): java.util.Map[String, String] = { + val conf = new java.util.HashMap[String, String]( + ConfigMap.parseFrom(CometExecIterator.serializeCometSQLConfs()).getEntriesMap) + // A Spark task passes the batch size apart from the other configs. + conf.put(CometConf.COMET_BATCH_SIZE.key, CometConf.COMET_BATCH_SIZE.get().toString) + conf + } + + private val singlePartition = PartitioningOuterClass.Partitioning + .newBuilder() + .setSinglePartition(PartitioningOuterClass.SinglePartition.newBuilder()) + .build() + + private class StageGraphBuilder(query: SparkPlan) { + val stages: ArrayBuffer[Stage] = ArrayBuffer.empty + private val stageIndexes = new IdentityHashMap[SparkPlan, Integer]() + + /** The index of the stage that runs the native block rooted at `plan`, adding it first. */ + def stageOf(plan: SparkPlan): Int = { + val known = stageIndexes.get(plan) + if (known != null) { + return known + } + val block = plan match { + case take: CometTakeOrderedAndProjectExec => return takeOrderedStage(take) + case native: CometNativeExec if native.serializedPlanOpt.isDefined => native + case other => fail(other, "it does not run natively") + } + val inputs = inputsOf(block).map { + case exchange: CometShuffleExchangeExec => + stageInput(exchange.child)(_.setPartitioning(partitioningOf(exchange))) + case broadcast => stageInput(broadcast.children.head)(_.setBroadcast(true)) + } + val blockPlan = Operator.parseFrom(block.serializedPlanOpt.plan.get) + + // A block that reads files runs one copy per partition, each naming its own files. + val (common, perPartition) = PlanDataInjector.findAllPlanData(block) + val partitionPlans = if (perPartition.isEmpty) { + Seq.empty + } else { + val base = new PlanDataInjector.CachedPlanData(blockPlan) + (0 until perPartition.values.head.length).map { i => + val partitionData = perPartition.map { case (key, data) => key -> data(i) } + PlanDataInjector.injectPlanData(base, common, partitionData).toByteString + } + } + // The cluster reads it only for a stage that reads no exchange. A stage that does has as + // many partitions as its exchanges deliver. + val numPartitions = + if (partitionPlans.nonEmpty) partitionPlans.length + else block.outputPartitioning.numPartitions + // A stage that reads a range exchange sorts each of its partitions. The cluster merges them + // in this order instead of range-partitioning the exchange. Without it, the cluster reads + // the exchange as one partition. + val outputOrdering = if (inputs.exists(_.getPartitioning.hasRangePartition)) { + val orders = block.outputOrdering.map(QueryPlanSerde.exprToProto(_, block.output)) + if (orders.nonEmpty && orders.forall(_.isDefined)) orders.flatten else Nil + } else { + Nil + } + addStage(plan, blockPlan, inputs, partitionPlans, numPartitions, outputOrdering) + } + + /** An input of a stage that reads the stage running `producer`, as `configure` says. */ + private def stageInput(producer: SparkPlan)( + configure: StageInput.Builder => StageInput.Builder): StageInput = + configure(StageInput.newBuilder().setProducer(stageOf(producer))).build() + + /** + * A stage for `take` that reads its child's whole output as one partition and keeps the first + * rows in order, as Comet's final top-K plan does after its shuffle to one partition. A top-K + * of each partition first would send fewer rows through the exchange + * (https://github.com/apache/datafusion-comet/issues/7365). + */ + private def takeOrderedStage(take: CometTakeOrderedAndProjectExec): Int = { + val input = stageInput(take.child)(_.setPartitioning(singlePartition)) + // More than one input partition, so that the plan sorts across all of them. + val blockPlan = take + .finalNativePlan(inputPartitions = 2) + .getOrElse(fail(take, "its top-K plan cannot run natively")) + addStage(take, blockPlan, Seq(input), Seq.empty, 1, Nil) + } + + private def addStage( + plan: SparkPlan, + blockPlan: Operator, + inputs: Seq[StageInput], + partitionPlans: Seq[ByteString], + numPartitions: Int, + outputOrdering: Seq[ExprOuterClass.Expr]): Int = { + findJvmCallback(blockPlan).foreach { name => + fail(plan, s"its native plan calls back into the JVM ($name)") + } + // Declared as Comet's shuffle writer declares its input, since the cluster casts the + // block's output to it as that writer does. + val outputSchema = schema2Proto( + plan.output.map(a => StructField(a.name, a.dataType, a.nullable, a.metadata))) + val stage = Stage + .newBuilder() + .setPlan(blockPlan) + .addAllInputs(inputs.asJava) + .addAllPartitionPlans(partitionPlans.asJava) + .setNumPartitions(numPartitions) + .addAllOutputSchema(outputSchema.asJava) + .addAllOutputOrdering(outputOrdering.asJava) + estimatedOutputBytes(plan).foreach(stage.setEstimatedOutputBytes) + stages += stage.build() + val index = stages.length - 1 + stageIndexes.put(plan, index) + index + } + + /** + * Spark's estimate of the size of `plan`'s output, from the logical plan of the nearest node + * in it that has one. The scheduler uses it to decide whether a join's build side is worth + * measuring before its probe side is shuffled. + */ + private def estimatedOutputBytes(plan: SparkPlan): Option[Long] = + plan + .collectFirst { case node if node.logicalLink.isDefined => node.logicalLink.get } + .flatMap(logical => Try(logical.stats.sizeInBytes).toOption) + .filter(_ >= 0) + .map(_.min(Long.MaxValue).toLong) + + /** + * The exchanges that feed `block`, in the order its native plan's input leaves read them: + * native shuffles and broadcasts, with a reused exchange resolved to the one it reuses. + */ + private def inputsOf(block: CometNativeExec): Seq[SparkPlan] = { + val inputs = ArrayBuffer.empty[SparkPlan] + def walk(plan: SparkPlan): Unit = plan match { + case reused: ReusedExchangeExec => + walk(reused.child) + case exchange: CometShuffleExchangeExec if exchange.shuffleType == CometNativeShuffle => + inputs += exchange + case broadcast: CometBroadcastExchangeExec => + inputs += broadcast + case _: CometNativeScanExec | _: CometIcebergNativeScanExec | _: CometRangeExec => + // Read inside the block. + // Part of the block, even when it has a serialized plan of its own: the conversion + // serializes some operators before their parent joins them into a larger block. + case native: CometNativeExec => + native.children.foreach(walk) + case other => fail(other, "Comet Native Executors do not support it yet") + } + walk(block) + inputs.toSeq + } + + private def partitioningOf( + exchange: CometShuffleExchangeExec): PartitioningOuterClass.Partitioning = { + def partitioning = PartitioningOuterClass.Partitioning.newBuilder() + exchange.outputPartitioning match { + case HashPartitioning(expressions, numPartitions) => + val keys = expressions.map { key => + QueryPlanSerde + .exprToProto(key, exchange.child.output) + .getOrElse(fail(exchange, s"its partitioning key $key cannot run natively")) + } + partitioning + .setHashPartition( + PartitioningOuterClass.HashPartition + .newBuilder() + .addAllHashExpression(keys.asJava) + .setNumPartitions(numPartitions)) + .build() + case RangePartitioning(ordering, numPartitions) => + // Without bounds: the cluster merges the stage's sorted partitions instead. + val sortOrders = ordering.map { order => + QueryPlanSerde + .exprToProto(order, exchange.child.output) + .getOrElse(fail(exchange, s"its sort key $order cannot run natively")) + } + partitioning + .setRangePartition( + PartitioningOuterClass.RangePartition + .newBuilder() + .addAllSortOrders(sortOrders.asJava) + .setNumPartitions(numPartitions)) + .build() + case SinglePartition => singlePartition + case other => + fail(exchange, s"Comet Native Executors do not support $other exchanges yet") + } + } + + private def fail(node: SparkPlan, reason: String): Nothing = { + val fallbackReasons = node + .getTagValue(CometExplainInfo.FALLBACK_REASONS) + .getOrElse(Set.empty[String]) + val because = if (fallbackReasons.isEmpty) { + "" + } else { + fallbackReasons.toSeq.sorted.mkString(" Comet could not convert it: ", "; ", ".") + } + throw new CometRuntimeException( + s"Comet Native Executors cannot run this query, because of ${node.nodeName}: " + + s"$reason.$because\n${query.treeString}") + } + } + + /** `message` and every message inside it, depth first. */ + private def messagesIn(message: Message): Iterator[Message] = + Iterator(message) ++ childMessages(message).flatMap { case (child, _) => messagesIn(child) } + + /** The first message in `message` that calls back into the JVM, if any. */ + private def findJvmCallback(message: Message): Option[String] = + messagesIn(message).map(_.getDescriptorForType.getName).find(jvmCallbacks.contains) +} diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeExecutorsSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeExecutorsSuite.scala new file mode 100644 index 00000000000..5d53eebd2c4 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeExecutorsSuite.scala @@ -0,0 +1,388 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.exec + +import java.io.File +import java.net.{ServerSocket, Socket} +import java.nio.file.Files +import java.util.concurrent.TimeUnit + +import scala.util.Try + +import org.apache.spark.SparkConf +import org.apache.spark.sql.{CometTestBase, DataFrame, Row} +import org.apache.spark.sql.comet.CometNativeExecutorsExec +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +import org.apache.spark.sql.functions.{col, count, max, sum} +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.CometConf + +/** + * Runs queries on a Comet Native Executor cluster (see + * https://github.com/apache/datafusion-comet/issues/7288): a comet-scheduler and a comet-executor + * started from the binaries in the directory `COMET_NATIVE_EXECUTORS_BIN_DIR` names. The tests + * are skipped when it isn't set. To run them: + * + * {{{ + * cargo build --manifest-path native/Cargo.toml -p datafusion-comet --features native-executors + * CARGO_TARGET_DIR=native/target cargo build --manifest-path native/native-executors/Cargo.toml + * COMET_NATIVE_EXECUTORS_BIN_DIR=native/target/debug ./mvnw test -Dtest=none \ + * -Dsuites=org.apache.comet.exec.CometNativeExecutorsSuite + * }}} + */ +class CometNativeExecutorsSuite extends CometTestBase with AdaptiveSparkPlanHelper { + + private val binDir: Option[File] = + sys.env.get("COMET_NATIVE_EXECUTORS_BIN_DIR").map(new File(_).getAbsoluteFile) + private var processes: Seq[Process] = Nil + private var schedulerUrl: String = _ + + override protected def sparkConf: SparkConf = { + val conf = super.sparkConf + conf.set(CometConf.COMET_SHUFFLE_MODE.key, "native") + conf.set(SQLConf.SHUFFLE_PARTITIONS.key, "4") + conf + } + + override def beforeAll(): Unit = { + super.beforeAll() + binDir.foreach(startCluster) + } + + override def afterAll(): Unit = { + try { + // The executor first, so that it doesn't report to a scheduler that has gone. + processes.reverse.foreach { process => + process.destroy() + if (!process.waitFor(10, TimeUnit.SECONDS)) { + process.destroyForcibly() + } + } + } finally { + super.afterAll() + } + } + + private def startCluster(dir: File): Unit = { + val schedulerPort = freePort() + val flightPort = freePort() + val executorPort = freePort() + val executorGrpcPort = freePort() + val logs = Files.createTempDirectory("comet-native-executors").toFile + logInfo(s"Comet Native Executor cluster logs: $logs") + start( + new File(dir, "comet-scheduler"), + Map( + "COMET_SCHEDULER_BIND_HOST" -> "127.0.0.1", + "COMET_SCHEDULER_PORT" -> schedulerPort.toString, + "COMET_SCHEDULER_FLIGHT_PORT" -> flightPort.toString), + new File(logs, "scheduler.log")) + waitForPort(schedulerPort) + waitForPort(flightPort) + start( + new File(dir, "comet-executor"), + Map( + "COMET_EXECUTOR_BIND_HOST" -> "127.0.0.1", + "COMET_EXECUTOR_PORT" -> executorPort.toString, + "COMET_EXECUTOR_GRPC_PORT" -> executorGrpcPort.toString, + "COMET_SCHEDULER_HOST" -> "127.0.0.1", + "COMET_SCHEDULER_PORT" -> schedulerPort.toString, + "COMET_EXECUTOR_TASKS" -> "4", + "COMET_EXECUTOR_WORK_DIR" -> new File(logs, "work").getPath), + new File(logs, "executor.log")) + waitForPort(executorPort) + schedulerUrl = s"http://127.0.0.1:$flightPort" + } + + private def start(binary: File, env: Map[String, String], log: File): Unit = { + val builder = new ProcessBuilder(binary.getPath) + env.foreach { case (key, value) => builder.environment().put(key, value) } + builder.redirectErrorStream(true).redirectOutput(log) + processes :+= builder.start() + } + + private def freePort(): Int = { + val socket = new ServerSocket(0) + try socket.getLocalPort + finally socket.close() + } + + private def waitForPort(port: Int): Unit = { + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(60) + while (Try(new Socket("127.0.0.1", port).close()).isFailure) { + assert(System.nanoTime() < deadline, s"Nothing is listening on port $port") + Thread.sleep(100) + } + } + + private def nativeExecutorsTest(name: String)(body: => Unit): Unit = test(name) { + assume(binDir.isDefined, "COMET_NATIVE_EXECUTORS_BIN_DIR is not set") + body + } + + private def onCluster(body: => Unit): Unit = withSQLConf( + CometConf.COMET_NATIVE_EXECUTORS_ENABLED.key -> "true", + CometConf.COMET_NATIVE_EXECUTORS_SCHEDULER_URL.key -> schedulerUrl)(body) + + /** + * Checks that the cluster runs `query` and returns what Spark returns, in order if `ordered`. + */ + private def checkOnCluster(query: => DataFrame, ordered: Boolean = false): Unit = + checkClusterReturns(query, query.collect().toSeq, ordered) + + /** Checks that the cluster runs `query` and returns `expected`, in order if `ordered`. */ + private def checkClusterReturns( + query: => DataFrame, + expected: Seq[Row], + ordered: Boolean): Unit = onCluster { + // A new Dataset, since a Dataset keeps the plan it was first run with. + val df = query + if (ordered) { + assert(df.collect().toSeq == expected, "the cluster returned other rows or another order") + } else { + checkAnswer(df, expected) + } + val plan = df.queryExecution.executedPlan + assert(collect(plan) { case exec: CometNativeExecutorsExec => exec }.size == 1, plan) + } + + /** A Parquet table with several files, read as several partitions. */ + private def withTable(body: String => Unit): Unit = withTempPath { dir => + spark + .range(0, 10000, 1, 8) + .selectExpr("id", "id % 7 AS k", "CAST(id AS DOUBLE) * 1.5 AS v") + .write + .parquet(dir.getCanonicalPath) + withSQLConf(SQLConf.FILES_MAX_PARTITION_BYTES.key -> "16384") { + body(dir.getCanonicalPath) + } + } + + nativeExecutorsTest("scan, filter and project") { + withTable { path => + checkOnCluster(spark.read.parquet(path).where("k > 2").selectExpr("id", "v * 2 AS w")) + } + } + + nativeExecutorsTest("aggregate across a hash shuffle") { + withTable { path => + checkOnCluster(spark.read.parquet(path).groupBy("k").agg(sum("v"), count("*"))) + } + } + + nativeExecutorsTest("aggregate across a hash shuffle without AQE") { + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withTable { path => + checkOnCluster(spark.read.parquet(path).groupBy("k").agg(sum("v"), max("id"))) + } + } + } + + nativeExecutorsTest("global aggregate across a single-partition shuffle") { + withTable { path => + checkOnCluster(spark.read.parquet(path).agg(sum("id"), max("v"), count("*"))) + } + } + + Seq(true, false).foreach { sortMerge => + val join = if (sortMerge) "sort-merge" else "hash" + nativeExecutorsTest(s"shuffled $join join") { + withSQLConf( + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.PREFER_SORTMERGEJOIN.key -> sortMerge.toString) { + withTable { path => + val left = spark.read.parquet(path).select("id", "k") + val right = spark.read + .parquet(path) + .where("id < 10") + .select(col("k").as("k2"), col("v")) + checkOnCluster(left.join(right, left("k") === right("k2")).select("id", "v")) + } + } + } + } + + /** + * TPC-H's `lineitem`: the table under `SPARK_TPCH_DATA` when it is set, as for + * `CometTPCHQuerySuite`, and otherwise generated data with the same schema. + */ + private def withLineitem(body: => Unit): Unit = sys.env.get("SPARK_TPCH_DATA") match { + case Some(path) => + spark.read.parquet(s"$path/lineitem").createOrReplaceTempView("lineitem") + withSQLConf(SQLConf.FILES_MAX_PARTITION_BYTES.key -> "32m")(body) + case None => + withTempPath { dir => + spark + .range(0, 100000, 1, 8) + .selectExpr( + "CAST(id % 50 + 1 AS DECIMAL(15, 2)) AS l_quantity", + "CAST((id % 1000) * 1.37 + 900 AS DECIMAL(15, 2)) AS l_extendedprice", + "CAST((id % 11) / 100.0 AS DECIMAL(15, 2)) AS l_discount", + "CAST((id % 9) / 100.0 AS DECIMAL(15, 2)) AS l_tax", + "ELT(CAST(id % 3 AS INT) + 1, 'A', 'N', 'R') AS l_returnflag", + "ELT(CAST(id % 2 AS INT) + 1, 'F', 'O') AS l_linestatus", + "DATE_ADD(DATE '1992-01-01', CAST(id % 2500 AS INT)) AS l_shipdate") + .write + .parquet(dir.getCanonicalPath) + spark.read.parquet(dir.getCanonicalPath).createOrReplaceTempView("lineitem") + withSQLConf(SQLConf.FILES_MAX_PARTITION_BYTES.key -> "131072")(body) + } + } + + private val q1Aggregate = """ + |select + | l_returnflag, + | l_linestatus, + | sum(l_quantity) as sum_qty, + | sum(l_extendedprice) as sum_base_price, + | sum(l_extendedprice * (1 - l_discount)) as sum_disc_price, + | sum(l_extendedprice * (1 - l_discount) * (1 + l_tax)) as sum_charge, + | avg(l_quantity) as avg_qty, + | avg(l_extendedprice) as avg_price, + | avg(l_discount) as avg_disc, + | count(*) as count_order + |from + | lineitem + |where + | l_shipdate <= date '1998-12-01' - interval '90' day + |group by + | l_returnflag, + | l_linestatus""".stripMargin + + nativeExecutorsTest("TPC-H q1 without its order by") { + withLineitem { + checkOnCluster(sql(q1Aggregate)) + } + } + + nativeExecutorsTest("TPC-H q1") { + withLineitem { + checkOnCluster( + sql(s"$q1Aggregate\norder by\n l_returnflag,\n l_linestatus"), + ordered = true) + } + } + + private val tpchTables = + Seq("customer", "lineitem", "nation", "orders", "part", "partsupp", "region", "supplier") + + /** + * The 22 TPC-H queries, on the tables under `SPARK_TPCH_DATA`. `TPCH_QUERIES`, a + * comma-separated list of query numbers, runs only those. + */ + private val tpchQueries = { + val only = sys.env.get("TPCH_QUERIES").map(_.split(",").map(_.trim.toInt).toSet) + (1 to 22).filter(q => only.forall(_.contains(q))) + } + + tpchQueries.foreach { q => + nativeExecutorsTest(s"TPC-H query $q") { + val data = sys.env.get("SPARK_TPCH_DATA") + assume(data.isDefined, "SPARK_TPCH_DATA is not set") + tpchTables.foreach { table => + spark.read.parquet(s"${data.get}/$table").createOrReplaceTempView(table) + } + val file = + new File(getClass.getResource("/").toURI.resolve("../../../benchmarks/tpc/queries/tpch/")) + val text = new String(Files.readAllBytes(new File(file, s"q$q.sql").toPath), "UTF-8") + val statements = text + .split("\n") + .filterNot(_.trim.startsWith("--")) + .mkString("\n") + .split(";") + .map(_.trim) + .filter(_.nonEmpty) + // q15 creates a view, queries it and drops it. + val (setup, rest) = statements.span(!_.toLowerCase.startsWith("select")) + val query = rest.head + val teardown = rest.tail + // The tables are temporary views, which only a temporary view can read. + setup.map(_.replaceFirst("(?i)^create view", "create temp view")).foreach(sql) + try { + // Spark's own answer. Comet on the driver would run the whole query in one JVM, in the + // test's small native memory pool, and Spark's broadcast relations don't fit in the test + // JVM's heap either. + // withSQLConf returns Unit on Spark 3.x, so the rows come out through a var. + var expected: Seq[Row] = Nil + withSQLConf( + CometConf.COMET_ENABLED.key -> "false", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ADAPTIVE_AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1") { + expected = sql(query).collect().toSeq + } + checkClusterReturns(sql(query), expected, query.toLowerCase.contains("order by")) + } finally { + teardown.foreach(sql) + } + } + } + + nativeExecutorsTest("broadcast hash join") { + withTable { path => + val left = spark.read.parquet(path).select("id", "k") + val right = spark.read.parquet(path).where("id < 10").select(col("k").as("k2"), col("v")) + checkOnCluster(left.join(right, left("k") === right("k2")).select("id", "v")) + } + } + + nativeExecutorsTest("scalar subquery") { + withTable { path => + spark.read.parquet(path).createOrReplaceTempView("t") + checkOnCluster(sql("select id, v from t where v > (select avg(v) from t where k = 3)")) + } + } + + nativeExecutorsTest("a command's plan stays in Spark") { + withTable { path => + onCluster { + withTempView("t") { + spark.read.parquet(path).createOrReplaceTempView("t") + // Explaining a command explains its result, which isn't a query to run on the cluster. + val df = sql("create or replace temp view v as select * from t") + assert( + df.queryExecution.explainString(org.apache.spark.sql.execution.SimpleMode).nonEmpty) + } + } + } + } + + nativeExecutorsTest("a query the cluster cannot run fails, naming the operator") { + withTable { path => + // With Comet's sort off, the sort stays in Spark. + withSQLConf(CometConf.COMET_EXEC_SORT_ENABLED.key -> "false") { + onCluster { + val error = intercept[Exception](spark.read.parquet(path).orderBy("v").collect()) + val messages = Iterator + .iterate[Throwable](error)(_.getCause) + .takeWhile(_ != null) + .map(e => String.valueOf(e.getMessage)) + .toSeq + assert( + messages.exists(message => + message.contains("Comet Native Executors cannot run this query") && + message.contains("Sort")), + error) + } + } + } + } +}