diff --git a/.github/actions/rust-test/action.yaml b/.github/actions/rust-test/action.yaml index ab5af0b68ca..6a4d430edba 100644 --- a/.github/actions/rust-test/action.yaml +++ b/.github/actions/rust-test/action.yaml @@ -45,6 +45,13 @@ runs: cd native cargo clippy --color=never -p datafusion-comet --all-targets --features python-udf -- -D warnings + # The JNI client for Comet Native Executors is off by default, so nothing above compiles it. + - name: Check native executors feature + shell: bash + run: | + cd native + cargo clippy --color=never -p datafusion-comet --all-targets --features native-executors -- -D warnings + - name: Check compilation shell: bash run: | @@ -106,6 +113,13 @@ runs: export LD_LIBRARY_PATH=${JAVA_HOME}/lib/server:${LD_LIBRARY_PATH} RUST_BACKTRACE=1 cargo nextest run -p datafusion-comet --lib --features python-udf python_udf + - name: Run native executors client tests + shell: bash + run: | + cd native + export LD_LIBRARY_PATH=${JAVA_HOME}/lib/server:${LD_LIBRARY_PATH} + RUST_BACKTRACE=1 cargo nextest run -p datafusion-comet --lib --features native-executors native_executors + # The steps above lint and test the accounting allocator over the system allocator. Lint and # test it over jemalloc too, since that is a different allocator backend and nothing else in CI # builds it. 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/_static/images/comet-native-executors.svg b/docs/source/_static/images/comet-native-executors.svg new file mode 100644 index 00000000000..7b51ab13041 --- /dev/null +++ b/docs/source/_static/images/comet-native-executors.svg @@ -0,0 +1,91 @@ + + How Comet Native Executors run a Spark query + A Spark driver plans the query with Catalyst and Comet, then sends it as a stage graph of Comet native blocks over Arrow Flight to a cluster with no JVM. The comet-scheduler checks each block with Comet's planner and schedules the stages on Apache DataFusion Ballista. Each comet-executor runs its tasks with Comet's operators on Apache DataFusion, reads Parquet directly, and passes data to the other executors between stages. The results return to the driver as Arrow over Flight. + + + + + + + + + + + + Spark driver with Comet + Plans each query with Catalyst and Comet, then hands all of it to the cluster + + JVM + + + + Comet Native Executors cluster + + No JVM + + + + comet-scheduler + + Checks each block with Comet's planner and schedules + the stages with Apache DataFusion Ballista + + + + + + comet-executor + + Comet operators on + Apache DataFusion + + + + comet-executor + + Comet operators on + Apache DataFusion + + + + comet-executor + + Comet operators on + Apache DataFusion + + + ยทยทยท + + + + + + Shuffle files or streamed batches between stages + + + + + Stage graph of native blocks + Comet protobuf over Arrow Flight + + + Tasks + + + + + Results as Arrow + over Arrow Flight + + + + + + + + Parquet files, such as on Amazon S3 + Read directly by every executor + + diff --git a/docs/source/contributor-guide/index.md b/docs/source/contributor-guide/index.md index 1be327be39e..41a4701426e 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..fce25af4f68 --- /dev/null +++ b/docs/source/contributor-guide/native_executors.md @@ -0,0 +1,231 @@ + + +# 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 graph of stages +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). Each stage is a +_block_: a run of Comet operators that Comet serializes as one Rust plan, with the exchanges between +blocks as the edges of the graph. The design is in +[issue #7288](https://github.com/apache/datafusion-comet/issues/7288), and +[`native/native-executors/README.md`](https://github.com/apache/datafusion-comet/blob/main/native/native-executors/README.md) +describes the cluster, how to build it and how to test it. + +This page is for contributors. It explains how to run a query on a cluster, lists what does not work +yet, and sketches ways to run part of a query in Spark. + +![A Spark driver with Comet sends each query as a graph of stages, each a block of Comet operators, to the comet-scheduler, which schedules its tasks on comet-executors with no JVM. The executors read Parquet directly, pass data to each other between stages, and Spark tasks read the results from them as Arrow over Flight.](../_static/images/comet-native-executors.svg) + +## Running it + +1. Build libcomet with the `native-executors` Cargo feature. It is off by default, so the standard + build links no gRPC client. A libcomet without it fails the first query with an error that names + the feature. For example, `make COMET_FEATURES=native-executors` builds libcomet and the Comet + jar that bundles it. The Spark executors that read the results need this libcomet too. +2. Build and start the cluster. The `comet-scheduler` and `comet-executor` binaries are their own + Cargo workspace, so the main Comet build does not depend on Ballista. The cluster and the driver + must come from the same Comet build, which the scheduler checks. See + [`native/native-executors/README.md`](https://github.com/apache/datafusion-comet/blob/main/native/native-executors/README.md) + for how to build the binaries and for the environment variables that configure them. +3. Point Spark at the scheduler's Flight service and turn the feature on: + + ```shell + --conf spark.comet.nativeExecutors.enabled=true \ + --conf spark.comet.nativeExecutors.scheduler.url=http://scheduler-host:50060 + ``` + + Both configs are internal while the feature is experimental. Comet itself must be enabled as + usual. + +4. Check that the query ran on the cluster: its physical plan has a `CometNativeExecutorsExec` leaf + where the Comet plan would otherwise be. A query that cannot run entirely in the Comet pipeline + and on the cluster fails instead, with a message that names the operator and Comet's reasons + for it. + +`CometNativeExecutorsSuite` runs queries against a cluster it starts from the binaries in +`COMET_NATIVE_EXECUTORS_BIN_DIR`. Its header comment gives the commands. + +## Current limitations + +Comet Native Executors are a first version. The sub-issues of +[#7288](https://github.com/apache/datafusion-comet/issues/7288) track what is missing: + +- Whole queries only. See [below](#whole-queries-only). +- `LIMIT` and `show()`: [#7365](https://github.com/apache/datafusion-comet/issues/7365). +- Global sorts by floats or nested types sort in a single task rather than in parallel: + [#7366](https://github.com/apache/datafusion-comet/issues/7366). +- Hash partitioning is DataFusion's, not Spark-compatible: + [#7368](https://github.com/apache/datafusion-comet/issues/7368). +- A small broadcast side runs in every task that reads it: + [#7369](https://github.com/apache/datafusion-comet/issues/7369). +- Dynamic partition pruning: [#7370](https://github.com/apache/datafusion-comet/issues/7370). A scan + with a dynamic pruning filter is refused. +- Writes: [#7371](https://github.com/apache/datafusion-comet/issues/7371). Commands and writes stay + in Spark. +- Cancellation: [#7372](https://github.com/apache/datafusion-comet/issues/7372). A killed Spark job + leaves its query running on the cluster. +- Spark UI metrics: [#7373](https://github.com/apache/datafusion-comet/issues/7373). +- Security. The scheduler and the executors have no authentication and no TLS, and a query names + the files and object stores the executors read. Run a cluster only on a trusted network. +- Result expiry. A query's results stay on the cluster for + `COMET_SCHEDULER_RESULT_RETENTION_SECONDS` after it ends (five minutes by default), whether or + not Spark has read them. A Spark task that reads a partition after that fails with a message that + names the setting. + +### Whole queries only + +The driver offloads a query only if every operator and expression in it runs in the Comet pipeline +(`CometNativeExecutors` in `CometNativeExecutorsExec.scala`). Otherwise it refuses the query, naming +the operator that could not 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. It +also includes a scan that stays in Spark, such as one whose rows Comet converts to Arrow batches, +because the cluster cannot read from a 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.** Each Spark task reads one result partition from the executor that holds it over + Arrow Flight, one endpoint per partition (`ResultEndpoint`). Data flows from the cluster to Spark, + but only at the root of the plan. + +## Possible future research: running part of a query in Spark + +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 on the cluster. 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 subtree that runs on the cluster as a leaf of the Spark plan + +The driver could offload the largest subtrees of Comet operators under an operator that must run in +Spark, rather than requiring the whole plan to run in the Comet pipeline. 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 offloaded subtree should end. + +It suits the common shape where the scans, joins and partial aggregates run in the Comet pipeline, +and an unsupported operator or UDF sits near the top of the plan. + +### Small Spark outputs as inputs to stages on the cluster + +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 stage 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 the executors + +Instead of splitting the plan, `comet-executor` could embed a JVM to evaluate the expressions and +UDFs that have no Rust implementation, reusing Comet's existing bridges for calling back into the +JVM. The whole plan would still run on the cluster 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 +tasks on the cluster 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 Spark executor without a JVM.** 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 data plane in Rust, 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 stage on the cluster 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 stage on + the cluster 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 block fetch client in Rust 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 the Rust side, and reports map statuses and accumulators. This works, but the + result is a Spark executor with a small heap whose tasks run in Rust: Comet on Spark with a data + plane that lives more in Rust. 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 Rust implementation, as described above, where Spark's protocols are + not involved. + +The data-plane option combines with either way of placing the work. The cluster's executors could +join a Spark application as executors, which needs the embedded JVM for their control plane, with +stage-level scheduling sending the stages that run entirely in the Comet pipeline to them through a +`ResourceProfile` that asks for a resource only they offer. Or the driver could keep offloading +whole subgraphs of stages to the cluster, which keeps its own scheduling and pipelining, and use a +Spark-compatible shuffle as the boundary between Spark stages and stages on the cluster. 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 eb93db1252b..20c9d19861b 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.2" @@ -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 4ba78cdfa43..cdff06aee94 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 } @@ -133,6 +136,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 a01789dead3..91ad8a8ef68 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -1058,6 +1058,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, @@ -1152,9 +1155,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), ); @@ -1234,7 +1253,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, 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..7d3843acef3 --- /dev/null +++ b/native/core/src/execution/native_executors/client.rs @@ -0,0 +1,271 @@ +// 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 std::time::Duration; + +use arrow_flight::decode::FlightRecordBatchStream; +use arrow_flight::error::FlightError; +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 tonic::Code; + +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}")) +} + +/// How long to wait for a connection to a scheduler or an executor. +const CONNECT_TIMEOUT: Duration = Duration::from_secs(30); + +/// How often to probe a connection that carries no traffic, at the TCP level and at the HTTP/2 +/// level. A query can run for a long time before the scheduler answers its submission, so a +/// request has no overall timeout. These probes are what notice a scheduler or an executor that +/// died or became unreachable without closing the connection. +const KEEP_ALIVE_INTERVAL: Duration = Duration::from_secs(30); + +/// How long to wait for an answer to an HTTP/2 probe before giving up on the connection. +const KEEP_ALIVE_TIMEOUT: Duration = Duration::from_secs(20); + +/// The error for a failed read of a result partition. A partition the cluster has deleted gets a +/// message that names the setting to change. The scheduler says that a partition it stores has +/// expired, and an executor fails to open the file of a partition it has deleted. +fn read_error(context: &str, error: FlightError) -> CometError { + let expired = match &error { + FlightError::Tonic(status) => match status.code() { + Code::NotFound => status.message().contains("expired"), + Code::Internal => { + let message = status.message(); + message.contains("Failed to open partition file") && message.contains("NotFound") + } + _ => false, + }, + _ => false, + }; + if expired { + client_error( + context, + format!( + "the result partition expired on the cluster; raise \ + COMET_SCHEDULER_RESULT_RETENTION_SECONDS ({error})" + ), + ) + } else { + client_error(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_timeout(CONNECT_TIMEOUT) + .tcp_keepalive(Some(KEEP_ALIVE_INTERVAL)) + .http2_keep_alive_interval(KEEP_ALIVE_INTERVAL) + .keep_alive_timeout(KEEP_ALIVE_TIMEOUT) + .keep_alive_while_idle(true) + .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| read_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| read_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(()) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use tonic::Status; + + fn message(status: Status) -> String { + read_error("cannot read a partition", FlightError::from(status)).to_string() + } + + #[test] + fn an_expired_result_names_the_setting_to_raise() { + for status in [ + Status::not_found("The results of job abc have expired"), + Status::internal( + "Ballista Error: General(\"Failed to open partition file at \\\"p\\\": \ + Os { code: 2, kind: NotFound, message: \\\"No such file or directory\\\" }\")", + ), + ] { + assert!( + message(status).contains("COMET_SCHEDULER_RESULT_RETENTION_SECONDS"), + "no hint" + ); + } + } + + #[test] + fn other_failures_do_not_blame_expiry() { + for status in [ + Status::not_found("Job abc has no result partition 7"), + Status::internal("Ballista Error: something else"), + Status::unavailable("connection reset"), + ] { + assert!(!message(status).contains("expired")); + } + } +} diff --git a/native/core/src/execution/native_executors/fragment.rs b/native/core/src/execution/native_executors/fragment.rs index 4e4aec790e0..3e542265e4d 100644 --- a/native/core/src/execution/native_executors/fragment.rs +++ b/native/core/src/execution/native_executors/fragment.rs @@ -22,8 +22,9 @@ use std::sync::Arc; use arrow::array::{Array, AsArray, BinaryArray, RecordBatch, RecordBatchOptions, UInt32Array}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef, UInt32Type}; use datafusion::common::tree_node::TreeNodeRecursion; -use datafusion::common::{internal_datafusion_err, DataFusionError, Result}; +use datafusion::common::{internal_datafusion_err, plan_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, }; @@ -52,6 +55,14 @@ use crate::execution::spark_config::{ const DEFAULT_BATCH_SIZE: usize = 8192; +/// The most partitions a stage or an exchange may have. A request names the counts, and the +/// scheduler allocates for each partition, so a request can't ask for more than this. Spark's +/// own default is 200 shuffle partitions, and large queries use a few thousand. +pub const MAX_PARTITIONS: usize = 1 << 16; + +/// The largest `spark.comet.batchSize` a request may set. Each operator sizes its buffers by it. +const MAX_BATCH_SIZE: usize = 1 << 20; + fn descriptor_schema() -> SchemaRef { Arc::new(Schema::new(vec![ Field::new("partition", DataType::UInt32, false), @@ -71,6 +82,11 @@ pub fn partition_descriptor( num_partitions: usize, plans: &[Vec], ) -> Result> { + if num_partitions == 0 || num_partitions > MAX_PARTITIONS { + return Err(plan_datafusion_err!( + "A stage must have between 1 and {MAX_PARTITIONS} partitions, not {num_partitions}" + )); + } let schema = descriptor_schema(); let partitions = (0..num_partitions) .map(|index| { @@ -88,6 +104,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 +165,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 +181,41 @@ 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(); + if self.children.len() != self.fragment.broadcast.len() + 1 { + return Err(internal_datafusion_err!( + "CometFragmentExec has {} inputs but {} broadcast flags", + self.children.len() - 1, + self.fragment.broadcast.len() + )); + } + 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 +251,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 +298,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,14 +412,16 @@ fn fragment_session( ) -> Result> { let mut runtime = comet_runtime_env(memory_pool); runtime.disk_manager = disk_manager; - let session = comet_session_context( - runtime, - conf.get_usize(COMET_BATCH_SIZE, DEFAULT_BATCH_SIZE), - 1, - conf, - spark_plan, - ) - .map_err(|e| DataFusionError::External(Box::new(e)))?; + let mut conf = conf.clone(); + conf.insert(NATIVE_EXECUTORS_FRAGMENT.to_string(), "true".to_string()); + let batch_size = conf.get_usize(COMET_BATCH_SIZE, DEFAULT_BATCH_SIZE); + if batch_size == 0 || batch_size > MAX_BATCH_SIZE { + return Err(plan_datafusion_err!( + "{COMET_BATCH_SIZE} must be between 1 and {MAX_BATCH_SIZE}, not {batch_size}" + )); + } + let session = comet_session_context(runtime, batch_size, 1, &conf, spark_plan) + .map_err(|e| DataFusionError::External(Box::new(e)))?; Ok(Arc::new(session)) } diff --git a/native/core/src/execution/native_executors/mod.rs b/native/core/src/execution/native_executors/mod.rs index fc0b9d6b2f4..7e079be7298 100644 --- a/native/core/src/execution/native_executors/mod.rs +++ b/native/core/src/execution/native_executors/mod.rs @@ -24,9 +24,14 @@ //! 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, + MAX_PARTITIONS, +}; /// 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 47a671f95bd..c20f0300da1 100644 --- a/native/native-executors/Cargo.lock +++ b/native/native-executors/Cargo.lock @@ -999,8 +999,7 @@ dependencies = [ [[package]] name = "ballista-api-types" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b97d9eea12c759eda5f828b4f0eba2b42784d96a4e7858afe5624a9522507841" +source = "git+https://github.com/apache/datafusion-ballista?branch=comet-native#0f092ee494d47d611f6e83f3660e60a4e4038d8c" dependencies = [ "serde", "utoipa", @@ -1009,8 +1008,7 @@ dependencies = [ [[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", @@ -1040,8 +1038,7 @@ dependencies = [ [[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", @@ -1070,8 +1067,7 @@ dependencies = [ [[package]] name = "ballista-history" version = "55.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "926ed4783d3b4d24fc28e899c7aef2f0127a511238320eadfc06c65ba2772bf8" +source = "git+https://github.com/apache/datafusion-ballista?branch=comet-native#0f092ee494d47d611f6e83f3660e60a4e4038d8c" dependencies = [ "ballista-api-types", "log", @@ -1083,8 +1079,7 @@ dependencies = [ [[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", @@ -4110,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", ] @@ -5105,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", @@ -5127,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", diff --git a/native/native-executors/Cargo.toml b/native/native-executors/Cargo.toml index 388c06189aa..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,11 +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"] } -# `rest-api` compiles the scheduler's event-log writer (COMET_SCHEDULER_EVENT_LOG_DIR) and its REST -# API. In Ballista 55.0.0 it logs through `tracing` without enabling it, so that is enabled here. -ballista-scheduler = { version = "55.0.0", default-features = false, features = ["rest-api", "tracing"] } +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..16496ed89d1 100644 --- a/native/native-executors/README.md +++ b/native/native-executors/README.md @@ -25,9 +25,14 @@ Spark driver plans each query with Catalyst and Comet, then submits it to a clus [Apache DataFusion Ballista](https://github.com/apache/datafusion-ballista). The design is in [issue #7288](https://github.com/apache/datafusion-comet/issues/7288). -This directory holds the cluster. The Spark driver side is in progress in -[#7289](https://github.com/apache/datafusion-comet/pull/7289). It is its own Cargo workspace, so the -main Comet build never depends on Ballista. +This directory holds the cluster, which is its own Cargo workspace, so the main Comet build never +depends on Ballista. The other half is in Comet's Spark plugin: with +`spark.comet.nativeExecutors.enabled` set, the driver submits each query it can offload to the +scheduler's Flight service (`spark.comet.nativeExecutors.scheduler.url`) and reads the results back. +The two must come from the same Comet build, which the scheduler checks. + +**Security.** The scheduler and the executors have no authentication and no TLS, and a query names +the files and object stores the executors read. Run a cluster only on a trusted network. ## The scheduler's Flight service @@ -39,34 +44,78 @@ A client runs a query through the scheduler's Arrow Flight service: from another Comet build and plans each block with Comet's planner to check it. That builds no object store for the block's files, so the scheduler uses no storage credentials. It then builds a Ballista physical plan (`src/assembly.rs`), runs it on the Ballista scheduler it - wraps, and answers with one endpoint per result partition. + wraps, and answers with one endpoint per result partition. It refuses a request with more + partitions or a larger batch size than a cluster plausibly needs, a stage whose block reads a + different number of inputs than the stage declares, and a block that writes files. 2. `DoGet` with each endpoint's ticket, at the endpoint's location: the executor that holds the partition. A query's results stay there for `COMET_SCHEDULER_RESULT_RETENTION_SECONDS` after the query ends (five minutes by default), whether or not a client has read them, and each endpoint's expiration time says until when. +Ballista runs a query in one of two modes, set for every query with +`COMET_SCHEDULER_BALLISTA_CONFIG` (`ballista.execution.mode`, blocking by default). In pipelined +mode the executors stream a result once and cannot hold it, so the scheduler reads each query's +result into its own memory and serves it from its own Flight service. Its endpoints have no +location: read them from the scheduler. The scheduler buffers at most +`COMET_SCHEDULER_PIPELINED_RESULT_MAX_BYTES` for a query and +`COMET_SCHEDULER_PIPELINED_STORE_MAX_BYTES` for all queries, and fails a query whose result is +larger than the first. When a new result does not fit in the second, the scheduler drops the oldest +results, which a client then finds expired. With a retention of `0` that is the only thing that +drops them. Run a query with a large result in blocking mode. + +The scheduler cancels a job whose client disconnects, whose result it cannot read, or that runs +longer than `COMET_SCHEDULER_QUERY_TIMEOUT_SECONDS`. + 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: -- A stage can't call into the JVM: scalar subqueries and JVM UDFs fail to plan, and a scan whose - object store needs a JVM credential provider fails when it runs. +- A stage can't call into the JVM: JVM UDFs and a scalar subquery the driver has not replaced with + its value fail to plan, and a scan whose object store needs a JVM credential provider fails when + it runs. +- A stage can't write files: a Parquet writer, an Iceberg write or a shuffle writer in a block is + refused. The cluster moves data between stages itself. - A stage can't load a native UDF library. These processes serve every client, and a plan names the library path, so a client would choose what the processes load. - 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 @@ -77,7 +126,16 @@ native/target/debug/comet-executor & Sharing `native/target` with the main workspace saves building the common dependencies twice. Both binaries read their settings from environment variables, listed at the top of -`src/bin/comet-scheduler.rs` and `src/bin/comet-executor.rs`. +`src/bin/comet-scheduler.rs` and `src/bin/comet-executor.rs`. The scheduler's include: + +| Variable | Meaning | +| -------------------------------------------- | ------------------------------------------------------------------------------ | +| `COMET_SCHEDULER_BALLISTA_CONFIG` | Ballista settings for every query, such as `ballista.execution.mode=pipelined` | +| `COMET_SCHEDULER_RESULT_RETENTION_SECONDS` | How long a query's results stay readable after it ends (default 300) | +| `COMET_SCHEDULER_FLIGHT_MAX_REQUEST_BYTES` | The largest query the Flight service accepts (default 512 MiB) | +| `COMET_SCHEDULER_QUERY_TIMEOUT_SECONDS` | The longest a query may take (default 0: no limit) | +| `COMET_SCHEDULER_PIPELINED_RESULT_MAX_BYTES` | The most Arrow data buffered for one pipelined result (default 1 GiB) | +| `COMET_SCHEDULER_PIPELINED_STORE_MAX_BYTES` | The most Arrow data kept for all pipelined results together (default 4 GiB) | ## Testing @@ -90,8 +148,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..e837c3cf624 --- /dev/null +++ b/native/native-executors/src/adaptive.rs @@ -0,0 +1,517 @@ +// 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))) = ( + inputs + .get(join.probe_input) + .copied() + .and_then(hash_exchange), + inputs + .get(join.build_input) + .copied() + .and_then(hash_exchange), + ) 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, _))) = ( + inputs + .get(join.probe_input) + .copied() + .and_then(hash_exchange), + inputs + .get(join.build_input) + .copied() + .and_then(hash_exchange), + ) 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) = inputs + .get(join.probe_input) + .copied() + .and_then(hash_exchange) + .ok_or_else(|| internal_datafusion_err!("The join's probe side is not an exchange"))?; + let (build, build_stage) = inputs + .get(join.build_input) + .copied() + .and_then(hash_exchange) + .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..ab6a9147a91 100644 --- a/native/native-executors/src/assembly.rs +++ b/native/native-executors/src/assembly.rs @@ -21,19 +21,27 @@ 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, + MAX_PARTITIONS, }; +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_operator::operator::OpStruct; +use datafusion_comet_proto::spark_operator::Operator; 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 +82,19 @@ 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"))?; + // The block's input leaves are fed from the stage's inputs, in order, so the two must agree. + check_block(plan, stage.inputs.len())?; + for (index, partition_plan) in stage.partition_plans.iter().enumerate() { + let partition_plan = Operator::decode(partition_plan.as_slice()).map_err(|e| { + plan_datafusion_err!("Cannot decode the plan of partition {index}: {e}") + })?; + check_block(&partition_plan, stage.inputs.len()) + .map_err(|e| e.context(format!("Plan of partition {index}")))?; + } // 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 +116,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,28 +130,43 @@ 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"); } + if num_partitions > MAX_PARTITIONS { + return plan_err!( + "The stage has {num_partitions} partitions, more than the {MAX_PARTITIONS} allowed" + ); + } if !stage.partition_plans.is_empty() && stage.partition_plans.len() != num_partitions { return plan_err!( "The stage has {num_partitions} partitions but {} partition plans", 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 +182,54 @@ 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, + }) +} + +/// Checks a stage's block, `block`, before it is planned: it must read as many inputs as the +/// stage declares, and it must not write files. Like a native UDF library or a call into the JVM, +/// writing is something a driver's stage never does on the cluster, so a request that asks for it +/// is refused rather than run by processes that serve every client. +fn check_block(block: &Operator, inputs: usize) -> Result<()> { + let mut leaves = 0; + let mut pending = vec![block]; + while let Some(op) = pending.pop() { + let writer = match op.op_struct { + Some(OpStruct::ParquetWriter(_)) => Some("ParquetWriter"), + Some(OpStruct::IcebergWrite(_)) => Some("IcebergWrite"), + Some(OpStruct::ShuffleWriter(_)) => Some("ShuffleWriter"), + Some(OpStruct::Scan(_)) | Some(OpStruct::ShuffleScan(_)) => { + leaves += 1; + None + } + _ => None, + }; + if let Some(writer) = writer { + return plan_err!( + "A stage can't contain a {writer} operator: Comet Native Executors don't write \ + files, and the cluster moves data between stages itself" + ); + } + pending.extend(&op.children); + } + if leaves != inputs { + return plan_err!( + "The stage's block reads {leaves} inputs but the stage declares {inputs}" + ); + } + Ok(()) } /// The exchange that carries `producer`'s output to a stage, as a node Ballista's planner turns @@ -167,13 +238,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 @@ -181,32 +257,111 @@ fn exchange( .and_then(|p| p.partitioning_struct.as_ref()); match partitioning { Some(PartitioningStruct::HashPartition(hash)) => { + let num_partitions = usize::try_from(hash.num_partitions) + .ok() + .filter(|count| (1..=MAX_PARTITIONS).contains(count)) + .ok_or_else(|| { + plan_datafusion_err!( + "A hash exchange must have between 1 and {MAX_PARTITIONS} partitions, \ + not {}", + hash.num_partitions + ) + })?; let schema = producer.schema(); 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( producer, - Partitioning::Hash(keys, hash.num_partitions as usize), + Partitioning::Hash(keys, num_partitions), )?)) } - // 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-scheduler.rs b/native/native-executors/src/bin/comet-scheduler.rs index ba68c19a1d6..fdbf88ebc57 100644 --- a/native/native-executors/src/bin/comet-scheduler.rs +++ b/native/native-executors/src/bin/comet-scheduler.rs @@ -29,9 +29,23 @@ //! - `COMET_GRPC_MAX_MESSAGE_SIZE`: the largest message, in bytes, of Ballista's gRPC service //! between the scheduler and its executors, such as a task with many partitions of a scan //! (default: Ballista's 16 MiB). Set it on the executors as well, for their side of the same -//! messages. The Flight service for drivers has no limit. Shuffle reads between executors keep -//! the limit of the query's `ballista.client.grpc_max_message_size` (16 MiB by default), which -//! this does not change. +//! messages. Shuffle reads between executors keep the limit of the query's +//! `ballista.client.grpc_max_message_size` (16 MiB by default), which this does not change. +//! - `COMET_SCHEDULER_FLIGHT_MAX_REQUEST_BYTES`: the largest query, in bytes, the Flight service +//! accepts from a driver. A query carries every stage's plan and the files of every scan +//! partition, so a large one is big (default `536870912`, 512 MiB). +//! - `COMET_SCHEDULER_BALLISTA_CONFIG`: Ballista settings for every query, as comma-separated +//! `key=value` pairs, such as `ballista.execution.mode=pipelined` +//! - `COMET_SCHEDULER_QUERY_TIMEOUT_SECONDS`: the longest a query may take. The scheduler fails +//! and cancels a query that runs longer. A query is also cancelled when the driver that +//! submitted it disconnects. (default `0`: no limit) +//! - `COMET_SCHEDULER_PIPELINED_RESULT_MAX_BYTES`: in pipelined mode, the scheduler reads each +//! query's result into its memory and serves it to the driver, and fails a query whose result +//! has more Arrow data than this (default `1073741824`, 1 GiB) +//! - `COMET_SCHEDULER_PIPELINED_STORE_MAX_BYTES`: the most Arrow data the scheduler keeps for +//! all pipelined results together. A new result that doesn't fit drops the oldest results, +//! which a driver then finds expired. With `COMET_SCHEDULER_RESULT_RETENTION_SECONDS` at `0` +//! this is the only thing that drops them (default `4294967296`, 4 GiB) //! - `COMET_SCHEDULER_EVENT_LOG_DIR`: a directory to write one event log per query to, as //! Ballista's scheduler does (default: no event logs). The scheduler creates it if needed and //! doesn't start if it can't write there. A log has the query's plan and a record for each @@ -46,17 +60,23 @@ //! results, not to logs. //! //! Ballista's REST API, which lists and cancels jobs without authentication, is not served. +//! +//! Neither the Flight service nor Ballista's gRPC service authenticates its callers or uses TLS, +//! and a query names the files and object stores executors read. Run a cluster only on a trusted +//! network. use std::net::SocketAddr; use std::path::Path; use std::sync::Arc; +use std::time::Duration; 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::frontend::{CometFlightService, Limits}; use comet_native_executors::{env_var, init_logging}; use datafusion_proto::protobuf::{LogicalPlanNode, PhysicalPlanNode}; use log::info; @@ -66,6 +86,7 @@ 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)), // Compiling the event-log writer compiles the REST API in with it, on the gRPC port. disable_rest_api: true, ..Default::default() @@ -98,7 +119,21 @@ async fn main() -> Result<(), Box> { })?; config.event_log_dir = Some(dir); } + let max_request_bytes: usize = + env_var("COMET_SCHEDULER_FLIGHT_MAX_REQUEST_BYTES")?.unwrap_or(512 << 20); + let mut limits = Limits::default(); + if let Some(seconds) = env_var::("COMET_SCHEDULER_QUERY_TIMEOUT_SECONDS")? { + limits.query_timeout = (seconds > 0).then(|| Duration::from_secs(seconds)); + } + if let Some(bytes) = env_var("COMET_SCHEDULER_PIPELINED_RESULT_MAX_BYTES")? { + limits.max_result_bytes = bytes; + } + if let Some(bytes) = env_var("COMET_SCHEDULER_PIPELINED_STORE_MAX_BYTES")? { + limits.max_stored_bytes = bytes; + } 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); @@ -112,9 +147,11 @@ 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)) - .max_decoding_message_size(usize::MAX) - .max_encoding_message_size(usize::MAX), + FlightServiceServer::new( + CometFlightService::try_new(scheduler, &ballista_settings)?.with_limits(limits), + ) + .max_decoding_message_size(max_request_bytes) + .max_encoding_message_size(usize::MAX), ) .serve(flight_addr); tokio::select! { @@ -124,6 +161,21 @@ 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() +} + /// Creates `dir` if it doesn't exist, and checks that the scheduler can create files in it: a /// directory that already exists, such as a volume mounted with another owner, may not allow it. fn check_writable(dir: &Path) -> std::io::Result<()> { diff --git a/native/native-executors/src/frontend.rs b/native/native-executors/src/frontend.rs index b6077669819..9f7545f0806 100644 --- a/native/native-executors/src/frontend.rs +++ b/native/native-executors/src/frontend.rs @@ -18,29 +18,48 @@ //! The Arrow Flight service that Spark drivers submit queries to. //! //! `GetFlightInfo` takes a [`QueryRequest`] as its command, runs the query, and returns one -//! endpoint per result partition. Each endpoint names the executor that holds the partition and -//! carries an opaque ticket for its `DoGet`, so the driver never speaks Ballista's protocol. It -//! 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). +//! endpoint per result partition. Each endpoint carries an opaque ticket for its `DoGet`, so the +//! driver never speaks Ballista's protocol, and 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). +//! +//! In Ballista's blocking mode each endpoint names the executor that holds the partition. In its +//! pipelined mode the executors stream a result once, so the scheduler reads the result into its +//! own memory, within the limits of [`Limits`], and its endpoints have no location: the driver +//! reads them from this service. -use std::sync::Arc; -use std::time::{Duration, UNIX_EPOCH}; +use std::collections::{HashMap, VecDeque}; +use std::future::Future; +use std::sync::atomic::{AtomicUsize, Ordering}; +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_core::JobId; 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 log::info; +use futures::TryStreamExt; +use log::{info, warn}; use prost::Message; use prost_types::Timestamp; use tonic::{Request, Response, Status, Streaming}; @@ -51,37 +70,194 @@ use crate::assembly::{build_plan, check_version}; /// executors cache a runtime per session, so sharing one lets queries share those runtimes. const SESSION_ID: &str = "comet"; +/// How long a query may run, and how much of a pipelined query's result the scheduler buffers. +#[derive(Clone, Debug)] +pub struct Limits { + /// The longest a query may take, from submission to its end. The scheduler fails and cancels a + /// query that runs longer. `None` waits for as long as it takes. + pub query_timeout: Option, + /// The most bytes of Arrow data the scheduler buffers for one pipelined query's result. It + /// fails and cancels a query whose result is larger. + pub max_result_bytes: usize, + /// The most bytes of Arrow data the scheduler keeps for all pipelined results together. When + /// a new result doesn't fit it drops the oldest results, even before they expire. + pub max_stored_bytes: usize, +} + +impl Default for Limits { + fn default() -> Self { + Self { + query_timeout: None, + max_result_bytes: 1 << 30, + max_stored_bytes: 4 << 30, + } + } +} + /// The Flight service of a Comet scheduler, which runs queries on the Ballista scheduler it /// 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, + limits: Limits, } 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, + limits: Limits::default(), + }) } - /// 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 + /// The service with `limits` instead of the defaults. + pub fn with_limits(mut self, limits: Limits) -> Self { + self.limits = limits; + self + } + + /// 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) + ) + } + + /// Keeps a pipelined query's result until it expires, as long as a blocking job's. With no + /// retention it is kept until the limit on the results kept pushes it out. Fails if the + /// result alone is over that limit. + fn store(&self, job_id: String, result: StoredResult) -> Result<(), Status> { + self.stored_results.lock().unwrap().insert( + job_id.clone(), + Arc::new(result), + self.limits.max_stored_bytes, + )?; + let retention = 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 + .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); + }); + } + Ok(()) + } + + /// 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. + fn results_expire_at(&self, ended_at: u64) -> Option { + let retention = 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}"); + .state + .config + .finished_job_data_clean_up_interval_seconds; + if retention == 0 { + return None; + } + UNIX_EPOCH + .checked_add(Duration::from_millis(ended_at))? + .checked_add(Duration::from_secs(retention)) + .map(Timestamp::from) + } +} + +/// Cancels its job when dropped before [`finished`](Self::finished), as when the driver +/// disconnects and the request's future is dropped, or a wait times out or a read fails. +struct JobGuard { + scheduler: SchedulerServer, + job_id: Option, +} + +impl JobGuard { + /// The job has ended, so there is nothing to cancel. + fn finished(&mut self) { + self.job_id = None; + } +} + +impl Drop for JobGuard { + fn drop(&mut self) { + let Some(job_id) = self.job_id.take() else { + return; + }; + let Ok(runtime) = tokio::runtime::Handle::try_current() else { + return; + }; + let scheduler = self.scheduler.clone(); + runtime.spawn(async move { + info!("Cancelling job {job_id}, which nothing is waiting for"); + if let Err(e) = scheduler.cancel_job(job_id.clone()).await { + warn!("Cannot cancel job {job_id}: {e}"); + } + }); + } +} + +/// Waits for `wait`, the end of job `job_id`, for at most `timeout`. +async fn within( + timeout: Option, + job_id: &JobId, + wait: impl Future>, +) -> Result { + let Some(timeout) = timeout else { + return wait.await; + }; + tokio::time::timeout(timeout, wait) + .await + .unwrap_or_else(|_| { + Err(Status::deadline_exceeded(format!( + "Job {job_id} did not finish within {} seconds \ + (COMET_SCHEDULER_QUERY_TIMEOUT_SECONDS) and was cancelled", + timeout.as_secs() + ))) + }) +} + +fn failed(job_id: &JobId, error: &str) -> Status { + Status::internal(format!("Job {job_id} failed: {error}")) +} + +/// 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, + limits: &Limits, +) -> Result { + let (job_id, mut receiver, mut guard) = submit(scheduler, session, plan).await?; + let wait = async { while let Some(status) = receiver.recv().await { match status.status { Some(job_status::Status::Successful(mut job)) => { + guard.finished(); job.partition_location.sort_by_key(|location| { let partition = location.partition_id.as_ref().map(|p| p.partition_id); (partition, location.map_partition_id) @@ -89,39 +265,279 @@ impl CometFlightService { return Ok(job); } Some(job_status::Status::Failed(job)) => { - return Err(Status::internal(format!( - "Job {job_id} failed: {}", - job.error - ))) + guard.finished(); + return Err(failed(&job_id, &job.error)); } _ => {} } } + guard.finished(); Err(Status::internal(format!( "Job {job_id} ended without a final status" ))) + }; + within(limits.query_timeout, &job_id, wait).await +} + +/// Submits `plan` as a Ballista job on `scheduler`, and returns its id, the channel its status +/// changes arrive on, and a guard that cancels it if the caller stops waiting. +async fn submit( + scheduler: &SchedulerServer, + session: &SessionConfig, + plan: Arc, +) -> Result<(JobId, tokio::sync::mpsc::Receiver, JobGuard), 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}")))?; + // Ballista drops a status it cannot fit in the channel, and a dropped final status would + // leave the caller waiting for ever, so the channel is far larger than the statuses of a job. + let (sender, receiver) = tokio::sync::mpsc::channel(1000); + 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}"); + let guard = JobGuard { + scheduler: scheduler.clone(), + job_id: Some(job_id.clone()), + }; + Ok((job_id, receiver, guard)) +} + +/// 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. +/// It buffers at most [`Limits::max_result_bytes`] of them, and fails the query beyond that. +/// +/// 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, + limits: &Limits, +) -> 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, mut guard) = submit(scheduler, session, plan).await?; + let wait = async { + let (locations, mut succeeded) = loop { + let Some(status) = receiver.recv().await else { + guard.finished(); + 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, + .. + })) if !pipelined_output_locations.is_empty() => { + break (pipelined_output_locations, false) + } + Some(job_status::Status::Successful(SuccessfulJob { + pipelined_output_locations, + .. + })) if !pipelined_output_locations.is_empty() => { + break (pipelined_output_locations, true) + } + Some(job_status::Status::Failed(job)) => { + guard.finished(); + return Err(failed(&job_id, &job.error)); + } + _ => {} + } + }; + // Every partition at once: the final stage can't finish until all of them are read. + let buffered = AtomicUsize::new(0); + let fetch = futures::future::try_join_all( + locations + .into_iter() + .map(|location| fetch_pipelined(location, &config, &buffered, limits)), + ); + tokio::pin!(fetch); + // Keep reading the job's statuses meanwhile, so that its failure ends the wait, and a + // failed read ends it at once, with the job cancelled. + let mut watching = true; + let partitions = loop { + tokio::select! { + partitions = &mut fetch => break partitions?, + status = receiver.recv(), if watching => match status.and_then(|s| s.status) { + Some(job_status::Status::Successful(_)) => succeeded = true, + Some(job_status::Status::Failed(job)) => { + guard.finished(); + return Err(failed(&job_id, &job.error)); + } + Some(_) => {} + None => watching = false, + }, + } + }; + while watching && !succeeded { + match receiver.recv().await.and_then(|status| status.status) { + Some(job_status::Status::Successful(_)) => succeeded = true, + Some(job_status::Status::Failed(job)) => { + guard.finished(); + return Err(failed(&job_id, &job.error)); + } + Some(_) => {} + None => watching = false, + } + } + guard.finished(); + Ok(StoredResult::new(schema, partitions)) + }; + let result = within(limits.query_timeout, &job_id, wait).await?; + Ok((job_id.to_string(), result)) +} + +/// Reads one result partition of a pipelined job from the executor that streams it. +/// +/// `buffered` counts the bytes read for the whole result, which may not pass the limit. +async fn fetch_pipelined( + location: ballista_core::serde::protobuf::PipelinedPartitionLocation, + config: &BallistaConfig, + buffered: &AtomicUsize, + limits: &Limits, +) -> 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 mut stream = client + .fetch_pipelined_partition(&location, config.pipelined_registration_timeout()) + .await + .map_err(failed)?; + let max_bytes = limits.max_result_bytes.min(limits.max_stored_bytes); + let mut batches = vec![]; + while let Some(batch) = stream + .try_next() + .await + .map_err(|e| failed(BallistaError::DataFusionError(Box::new(e))))? + { + let size = batch.get_array_memory_size(); + if buffered + .fetch_add(size, Ordering::Relaxed) + .saturating_add(size) + > max_bytes + { + return Err(Status::resource_exhausted(format!( + "The query's result is larger than the {max_bytes} bytes the scheduler buffers \ + for a pipelined query (COMET_SCHEDULER_PIPELINED_RESULT_MAX_BYTES and \ + COMET_SCHEDULER_PIPELINED_STORE_MAX_BYTES). Run it in blocking mode instead." + ))); + } + batches.push(batch); } + Ok(batches) +} - /// 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. - fn results_expire_at(&self, ended_at: u64) -> Option { - let retention = self - .scheduler - .state - .config - .finished_job_data_clean_up_interval_seconds; - if retention == 0 { - return None; +/// The results of a pipelined query, which the scheduler serves itself. +struct StoredResult { + schema: SchemaRef, + partitions: Vec>, + /// The Arrow memory the batches hold. + bytes: usize, +} + +impl StoredResult { + fn new(schema: SchemaRef, partitions: Vec>) -> Self { + let bytes = partitions + .iter() + .flatten() + .map(RecordBatch::get_array_memory_size) + .sum(); + Self { + schema, + partitions, + bytes, } - UNIX_EPOCH - .checked_add(Duration::from_millis(ended_at))? - .checked_add(Duration::from_secs(retention)) - .map(Timestamp::from) } } +/// The stored results of pipelined queries, oldest first, within a limit on their total size. +#[derive(Default)] +struct ResultStore { + results: HashMap>, + /// The job ids in `results`, oldest first. + order: VecDeque, + bytes: usize, +} + +impl ResultStore { + /// Stores `result`, after dropping the oldest results as far as it takes to keep the total + /// within `max_bytes`. + fn insert( + &mut self, + job_id: String, + result: Arc, + max_bytes: usize, + ) -> Result<(), Status> { + if result.bytes > max_bytes { + return Err(Status::resource_exhausted(format!( + "The query's result is larger than the {max_bytes} bytes the scheduler keeps for \ + pipelined results (COMET_SCHEDULER_PIPELINED_STORE_MAX_BYTES)" + ))); + } + while self.bytes + result.bytes > max_bytes { + let Some(oldest) = self.order.front().cloned() else { + break; + }; + warn!("Dropping the results of job {oldest} to make room for job {job_id}'s"); + self.remove(&oldest); + } + self.bytes += result.bytes; + self.order.push_back(job_id.clone()); + self.results.insert(job_id, result); + Ok(()) + } + + fn remove(&mut self, job_id: &str) { + if let Some(result) = self.results.remove(job_id) { + self.bytes -= result.bytes; + self.order.retain(|id| id != job_id); + } + } + + fn get(&self, job_id: &str) -> Option> { + self.results.get(job_id).cloned() + } +} + +/// 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 +597,34 @@ 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, &self.limits).await?; + let partitions = result.partitions.len(); + self.store(job_id.clone(), result)?; + 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..partitions { + 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); + } + } else { + let job = run(&self.scheduler, &self.session, plan, &self.limits).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 +657,37 @@ 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) + .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..d0e86e7724f 100644 --- a/native/native-executors/tests/assembly.rs +++ b/native/native-executors/tests/assembly.rs @@ -20,20 +20,25 @@ mod common; +use comet::execution::native_executors::MAX_PARTITIONS; use comet_native_executors::assembly::build_plan; use common::*; +use datafusion::arrow::array::{AsArray, RecordBatch}; 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::{ expr::ExprStruct, literal::Value, Expr, IfExpr, Literal, Subquery, }; -use datafusion_comet_proto::spark_native_executors::{QueryRequest, Stage}; +use datafusion_comet_proto::spark_native_executors::{QueryRequest, Stage, StageInput}; use datafusion_comet_proto::spark_operator::{ - operator::OpStruct, NativeScan, NativeScanCommon, Operator, Projection, SparkFilePartition, - SparkPartitionedFile, + operator::OpStruct, IcebergWrite, NativeScan, NativeScanCommon, Operator, ParquetWriter, + Projection, Scan, ShuffleWriter, SparkFilePartition, SparkPartitionedFile, }; use prost::Message; +use std::cmp::Ordering; use std::sync::Arc; async fn assemble(request: QueryRequest) -> datafusion::common::Result> { @@ -76,17 +81,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); @@ -278,3 +329,348 @@ async fn file_stage_is_checked_without_its_object_store() { .unwrap_err(); assert!(error.to_string().contains("no-such-store"), "{error}"); } + +/// The values of column `column` of `batches`, in order. +fn longs(batches: &[RecordBatch], column: usize) -> Vec> { + batches + .iter() + .flat_map(|batch| { + batch + .column(column) + .as_primitive::() + .iter() + .collect::>() + }) + .collect() +} + +/// A stage of `n` rows in `slices` partitions with two columns: a key that is null for every +/// fifth row and otherwise `id % 7`, so it has many duplicates, and `id`. +fn keyed_stage(n: i64, slices: i32) -> Stage { + let id = || bound(0, long_type()); + let key = if_else( + equal(remainder(id(), long_literal(5)), long_literal(0)), + null_long(), + remainder(id(), long_literal(7)), + ); + projecting_stage( + range_stage(n, slices), + vec![key, id()], + vec![long_type(), long_type()], + ) +} + +fn keyed_rows(n: i64) -> Vec<(Option, i64)> { + (0..n) + .map(|id| ((id % 5 != 0).then_some(id % 7), id)) + .collect() +} + +fn compare_keys(a: &Option, b: &Option, descending: bool, nulls_first: bool) -> Ordering { + match (a, b) { + (None, None) => Ordering::Equal, + (None, Some(_)) if nulls_first => Ordering::Less, + (None, Some(_)) => Ordering::Greater, + (Some(_), None) if nulls_first => Ordering::Greater, + (Some(_), None) => Ordering::Less, + (Some(a), Some(b)) if descending => b.cmp(a), + (Some(a), Some(b)) => a.cmp(b), + } +} + +async fn run_in_order(request: QueryRequest) -> (Arc, Vec) { + let plan = assemble(request).await.unwrap(); + let batches = collect(Arc::clone(&plan), SessionContext::new().task_ctx()) + .await + .unwrap(); + (plan, batches) +} + +#[tokio::test(flavor = "multi_thread")] +async fn merge_honors_every_null_ordering() { + for descending in [false, true] { + for nulls_first in [false, true] { + let order = sort_key(0, long_type(), descending, nulls_first); + let request = query(vec![ + keyed_stage(500, 4), + sorting_stage_by( + 0, + range_partitioning(4, vec![]), + vec![order], + vec![long_type(), long_type()], + ), + ]); + let (plan, batches) = run_in_order(request).await; + assert!(plan.is::()); + let mut expected = keyed_rows(500) + .into_iter() + .map(|(key, _)| key) + .collect::>(); + expected.sort_by(|a, b| compare_keys(a, b, descending, nulls_first)); + assert_eq!( + longs(&batches, 0), + expected, + "descending={descending} nulls_first={nulls_first}" + ); + } + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn merge_honors_a_multi_key_order() { + // The first key has many duplicates and nulls, so the second decides most of the order. + let orders = vec![ + sort_key(0, long_type(), false, true), + sort_key(1, long_type(), true, false), + ]; + let request = query(vec![ + keyed_stage(500, 4), + sorting_stage_by( + 0, + range_partitioning(4, vec![]), + orders, + vec![long_type(), long_type()], + ), + ]); + let (plan, batches) = run_in_order(request).await; + assert!(plan.is::()); + let mut expected = keyed_rows(500); + expected.sort_by(|a, b| compare_keys(&a.0, &b.0, false, true).then(b.1.cmp(&a.1))); + let rows = longs(&batches, 0) + .into_iter() + .zip(longs(&batches, 1).into_iter().map(Option::unwrap)) + .collect::>(); + assert_eq!(rows, expected); +} + +#[tokio::test(flavor = "multi_thread")] +async fn merge_orders_strings_as_bytes() { + // As text "10" < "2", which differs from the numbers' order. + let key = cast(bound(0, long_type()), string_type()); + let producer = projecting_stage(range_stage(200, 4), vec![key], vec![string_type()]); + let order = sort_key(0, string_type(), false, false); + let request = query(vec![ + producer, + sorting_stage(0, range_partitioning(4, vec![]), order, string_type()), + ]); + let (plan, batches) = run_in_order(request).await; + assert!(plan.is::()); + let values = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_string::() + .iter() + .map(|value| value.unwrap().to_string()) + .collect::>() + }) + .collect::>(); + let mut expected = (0..200).map(|id| id.to_string()).collect::>(); + expected.sort(); + assert_eq!(values, expected); +} + +#[tokio::test(flavor = "multi_thread")] +async fn merge_tolerates_empty_producer_partitions() { + // Two rows in four partitions leave at least two of them empty. + let order = first_column_descending(long_type()); + let request = query(vec![ + range_stage(2, 4), + sorting_stage(0, range_partitioning(4, vec![]), order, long_type()), + ]); + let (plan, batches) = run_in_order(request).await; + assert!(plan.is::()); + assert_eq!(longs(&batches, 0), vec![Some(1), Some(0)]); +} + +#[tokio::test(flavor = "multi_thread")] +async fn ordered_stage_with_one_partition_needs_no_merge() { + let order = first_column_descending(long_type()); + let request = query(vec![ + range_stage(100, 1), + sorting_stage(0, range_partitioning(1, vec![]), order, long_type()), + ]); + let (plan, batches) = run_in_order(request).await; + assert!(!plan.is::()); + assert_eq!(plan.output_partitioning().partition_count(), 1); + assert_eq!( + longs(&batches, 0), + (0..100).rev().map(Some).collect::>() + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn hash_exchange_partition_counts_are_bounded() { + for count in [i32::MIN, -1, 0, MAX_PARTITIONS as i32 + 1, i32::MAX] { + let request = query(vec![ + range_stage(100, 4), + reading_stage(0, hash_on_first_column(count), 3), + ]); + let error = assemble(request).await.unwrap_err().to_string(); + assert!( + error.contains(&format!( + "between 1 and {MAX_PARTITIONS} partitions, not {count}" + )), + "{error}" + ); + } + let request = query(vec![ + range_stage(100, 4), + reading_stage(0, hash_on_first_column(MAX_PARTITIONS as i32), 3), + ]); + let plan = assemble(request).await.unwrap(); + assert_eq!(plan.output_partitioning().partition_count(), MAX_PARTITIONS); +} + +#[tokio::test(flavor = "multi_thread")] +async fn stage_partition_counts_are_bounded() { + for count in [MAX_PARTITIONS as u32 + 1, 1 << 31, u32::MAX] { + let mut stage = range_stage(100, 4); + stage.num_partitions = count; + let error = assemble(query(vec![stage])).await.unwrap_err().to_string(); + assert!( + error.contains(&format!("more than the {MAX_PARTITIONS} allowed")), + "{error}" + ); + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn batch_size_is_bounded() { + for batch_size in ["0", "2000000", "18446744073709551615"] { + let mut request = query(vec![range_stage(100, 4)]); + request + .conf + .insert("spark.comet.batchSize".to_string(), batch_size.to_string()); + let error = assemble(request).await.unwrap_err().to_string(); + assert!( + error.contains("spark.comet.batchSize must be between 1 and"), + "{error}" + ); + } + let mut request = query(vec![range_stage(100, 4)]); + request + .conf + .insert("spark.comet.batchSize".to_string(), "4096".to_string()); + assemble(request).await.unwrap(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn block_must_read_as_many_inputs_as_the_stage_declares() { + let leaf_less_input = StageInput { + producer: 0, + partitioning: Some(single_partition()), + broadcast: false, + }; + // A block with no input leaves, and an input. + let mut stage = range_stage(100, 4); + stage.inputs = vec![leaf_less_input.clone()]; + let error = assemble(query(vec![range_stage(100, 4), stage])) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("reads 0 inputs but the stage declares 1"), + "{error}" + ); + // A block with one input leaf, and two inputs. + let mut stage = reading_stage(0, single_partition(), 1); + stage.inputs.push(leaf_less_input); + let error = assemble(query(vec![range_stage(100, 4), stage])) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("reads 1 inputs but the stage declares 2"), + "{error}" + ); + // A join's two input leaves, and one input. + let mut stage = joining_stage(0, 1, 3); + stage.inputs.pop(); + let error = assemble(query(vec![range_stage(100, 4), range_stage(100, 2), stage])) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("reads 2 inputs but the stage declares 1"), + "{error}" + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn every_partition_plan_must_read_as_many_inputs_as_the_stage_declares() { + let range = range_stage(100, 2); + let block = range.plan.as_ref().unwrap().encode_to_vec(); + let with_leaf = Operator { + op_struct: Some(OpStruct::Scan(Scan { + fields: vec![long_type()], + source: "test".to_string(), + })), + ..Default::default() + } + .encode_to_vec(); + let mut stage = range; + stage.partition_plans = vec![block, with_leaf]; + let error = assemble(query(vec![stage])).await.unwrap_err(); + assert!( + error.to_string().contains("Plan of partition 1") + && error + .to_string() + .contains("reads 1 inputs but the stage declares 0"), + "{error}" + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn writer_operators_are_refused() { + let writers = [ + ( + "ParquetWriter", + OpStruct::ParquetWriter(ParquetWriter::default()), + ), + ( + "IcebergWrite", + OpStruct::IcebergWrite(IcebergWrite::default()), + ), + ( + "ShuffleWriter", + OpStruct::ShuffleWriter(ShuffleWriter::default()), + ), + ]; + for (name, writer) in writers { + let writing = |child: Operator| Operator { + op_struct: Some(writer.clone()), + children: vec![child], + ..Default::default() + }; + let expected = format!("can't contain a {name} operator"); + let range = range_stage(100, 1); + let block = range.plan.clone().unwrap(); + + // At the root of the block. + let mut stage = range.clone(); + stage.plan = Some(writing(block.clone())); + let error = assemble(query(vec![stage])).await.unwrap_err(); + assert!(error.to_string().contains(&expected), "{error}"); + + // Below other operators. + let mut stage = range.clone(); + stage.plan = Some(Operator { + op_struct: Some(OpStruct::Projection(Projection::default())), + children: vec![writing(block.clone())], + ..Default::default() + }); + let error = assemble(query(vec![stage])).await.unwrap_err(); + assert!(error.to_string().contains(&expected), "{error}"); + + // In a partition's own copy of the block. + let mut stage = range; + stage.partition_plans = vec![writing(block).encode_to_vec()]; + let error = assemble(query(vec![stage])).await.unwrap_err(); + assert!(error.to_string().contains(&expected), "{error}"); + } +} diff --git a/native/native-executors/tests/cluster.rs b/native/native-executors/tests/cluster.rs index 2ed61c2cebf..aa82acf1ac5 100644 --- a/native/native-executors/tests/cluster.rs +++ b/native/native-executors/tests/cluster.rs @@ -31,9 +31,9 @@ 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 datafusion_comet_proto::spark_native_executors::{QueryRequest, Stage}; use futures::TryStreamExt; use prost::Message; use tonic::transport::Endpoint; @@ -44,6 +44,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, /// The port of Ballista's gRPC service, for executors. scheduler_port: u16, } @@ -70,6 +72,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() + )), scheduler_port, }; let mut env = vec![ @@ -78,7 +84,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(); @@ -99,10 +111,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); @@ -121,17 +138,33 @@ impl Cluster { let info = self.get_flight_info(request).await.unwrap(); let mut partitions = vec![]; for endpoint in &info.endpoint { - partitions.push(read(endpoint).await.unwrap()); + partitions.push(self.read(endpoint).await.unwrap()); } partitions } + + /// Reads a result partition from where its endpoint says, or from the scheduler if the + /// endpoint names no location, as for the results of a pipelined query. + async fn read(&self, endpoint: &FlightEndpoint) -> Result, FlightError> { + match endpoint.location.first() { + Some(location) => read_from(&location.uri, endpoint).await, + None => read_from(&self.flight_url, endpoint).await, + } + } } /// Reads a result partition from the executor that holds it. async fn read(endpoint: &FlightEndpoint) -> Result, FlightError> { - let mut executor = connect(&endpoint.location[0].uri).await; + read_from(&endpoint.location[0].uri, endpoint).await +} + +async fn read_from( + location: &str, + endpoint: &FlightEndpoint, +) -> Result, FlightError> { + let mut server = connect(location).await; let ticket = endpoint.ticket.clone().unwrap(); - within_time_limit(async { executor.do_get(ticket).await?.try_collect().await }).await + within_time_limit(async { server.do_get(ticket).await?.try_collect().await }).await } impl Drop for Cluster { @@ -142,6 +175,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); } } @@ -222,7 +256,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); @@ -232,6 +266,267 @@ 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}"); +} + +/// The stages of `broadcasts_a_build_side_that_measures_small_into_the_probe_side_stage`, with +/// `build` for the build side and `join` for the stage that joins. Spark's estimates make the +/// build side worth measuring before the probe side is shuffled. +fn held_join_query(build: Stage, join: Stage) -> QueryRequest { + let mut probe = range_stage(10_000, 4); + probe.estimated_output_bytes = Some(1 << 40); + let mut build = build; + build.estimated_output_bytes = Some(1 << 30); + query(vec![probe, build, join]) +} + +/// `id % 10` of `range(0, 100)`: a build side whose join keys repeat ten times each. +fn duplicate_keys_build_side() -> Stage { + let key = remainder(bound(0, long_type()), long_literal(10)); + projecting_stage(range_stage(100, 2), vec![key], vec![long_type()]) +} + +#[tokio::test(flavor = "multi_thread")] +async fn broadcast_join_keeps_every_match_of_duplicate_build_keys() { + let cluster = Cluster::start(); + let request = held_join_query(duplicate_keys_build_side(), joining_stage(0, 1, 3)); + + let rows = cluster.run(&request).await.concat(); + + // Probe ids 0..10 each meet ten build rows. + assert_eq!( + rows.iter().map(|batch| batch.num_rows()).sum::(), + 100 + ); + let keys = sorted_values::(&rows); + assert_eq!(keys, (0..10).flat_map(|key| [key; 10]).collect::>()); + for batch in &rows { + assert_eq!(batch.column(0), batch.column(1)); + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn broadcast_join_that_is_not_the_root_of_its_block() { + // The join's consumer then holds the join's output in a leaf and runs a projection over it. + let cluster = Cluster::start(); + let mut join = joining_stage(0, 1, 3); + join = projecting_stage( + join, + vec![bound(0, long_type()), bound(1, long_type())], + vec![long_type(), long_type()], + ); + let request = held_join_query(range_stage(100, 2), join); + + let rows = cluster.run(&request).await.concat(); + + assert_eq!( + sorted_values::(&rows), + (0..100).collect::>() + ); + for batch in &rows { + assert_eq!(batch.column(0), batch.column(1)); + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn join_with_a_build_side_that_measures_empty_returns_no_rows() { + let cluster = Cluster::start(); + let request = held_join_query(range_stage(0, 2), joining_stage(0, 1, 3)); + + let rows = cluster.run(&request).await.concat(); + + assert_eq!(rows.iter().map(|batch| batch.num_rows()).sum::(), 0); +} + +/// Settings that run queries in Ballista's pipelined mode. +fn pipelined() -> (&'static str, String) { + ( + "COMET_SCHEDULER_BALLISTA_CONFIG", + "ballista.execution.mode=pipelined".to_string(), + ) +} + +#[tokio::test(flavor = "multi_thread")] +async fn pipelined_results_are_served_by_the_scheduler_as_often_as_they_are_read() { + let cluster = Cluster::start_with(&[pipelined()]); + let request = query(vec![ + range_stage(1000, 4), + reading_stage(0, hash_on_first_column(3), 3), + ]); + let info = cluster.get_flight_info(&request).await.unwrap(); + assert!(!info.endpoint.is_empty()); + assert!(info + .endpoint + .iter() + .all(|endpoint| endpoint.location.is_empty())); + + // Spark can compute an RDD partition again, so a partition can be read more than once. + for _ in 0..2 { + let mut batches = vec![]; + for endpoint in &info.endpoint { + batches.extend(cluster.read(endpoint).await.unwrap()); + } + assert_eq!( + sorted_values::(&batches), + (0..1000).collect::>() + ); + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn pipelined_result_over_the_buffer_limit_fails_the_query() { + let cluster = Cluster::start_with(&[ + pipelined(), + ( + "COMET_SCHEDULER_PIPELINED_RESULT_MAX_BYTES", + (64 << 10).to_string(), + ), + ]); + // 800,000 bytes of longs. + match cluster + .get_flight_info(&query(vec![range_stage(100_000, 4)])) + .await + { + Err(FlightError::Tonic(status)) => { + assert_eq!(status.code(), Code::ResourceExhausted, "{status}"); + assert!( + status + .message() + .contains("COMET_SCHEDULER_PIPELINED_RESULT_MAX_BYTES"), + "{status}" + ); + } + other => panic!("Expected the query to fail, got {other:?}"), + } + // The cluster still runs queries that fit: the failed query's job was cancelled. + let partitions = cluster.run(&query(vec![range_stage(100, 2)])).await; + assert_eq!( + sorted_values::(&partitions.concat()), + (0..100).collect::>() + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn pipelined_results_without_retention_are_dropped_oldest_first_at_the_limit() { + // Retention 0 never expires a result, so only the limit on the stored results drops one. Each + // result is about 16 KB, so two don't fit. + let cluster = Cluster::start_with(&[ + pipelined(), + ("COMET_SCHEDULER_RESULT_RETENTION_SECONDS", "0".to_string()), + ( + "COMET_SCHEDULER_PIPELINED_STORE_MAX_BYTES", + (24 << 10).to_string(), + ), + ]); + let request = query(vec![range_stage(2000, 2)]); + let first = cluster.get_flight_info(&request).await.unwrap(); + let second = cluster.get_flight_info(&request).await.unwrap(); + + for endpoint in &second.endpoint { + cluster.read(endpoint).await.unwrap(); + } + match cluster.read(&first.endpoint[0]).await { + Err(FlightError::Tonic(status)) => assert_eq!(status.code(), Code::NotFound, "{status}"), + other => panic!("Expected the first result to be dropped, got {other:?}"), + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn refuses_a_query_larger_than_the_request_limit() { + let cluster = Cluster::start_with(&[( + "COMET_SCHEDULER_FLIGHT_MAX_REQUEST_BYTES", + "4096".to_string(), + )]); + cluster.run(&query(vec![range_stage(10, 1)])).await; + let mut large = query(vec![range_stage(10, 1)]); + large.conf.insert("padding".to_string(), "x".repeat(8192)); + match cluster.get_flight_info(&large).await { + Err(FlightError::Tonic(status)) => assert!( + matches!(status.code(), Code::OutOfRange | Code::ResourceExhausted), + "{status}" + ), + other => panic!("Expected the query to be refused, got {other:?}"), + } +} + #[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..c17ec8785e1 100644 --- a/native/native-executors/tests/common/mod.rs +++ b/native/native-executors/tests/common/mod.rs @@ -25,10 +25,14 @@ 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, literal::Value, BinaryExpr, BoundReference, Cast, Expr, IfExpr, + Literal, MathExpr, 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, Projection, RangeScan, Scan, + ShuffleScan, Sort, SparkStructField, }; use datafusion_comet_proto::spark_partitioning::{ partitioning::PartitioningStruct, HashPartition, Partitioning, RangePartition, SinglePartition, @@ -60,6 +64,14 @@ pub fn int_type() -> spark_expression::DataType { data_type(DataTypeId::Int32) } +pub fn double_type() -> spark_expression::DataType { + data_type(DataTypeId::Double) +} + +pub fn string_type() -> spark_expression::DataType { + data_type(DataTypeId::String) +} + pub fn bool_type() -> spark_expression::DataType { data_type(DataTypeId::Bool) } @@ -125,6 +137,148 @@ 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 { + sorting_stage_by(producer, partitioning, vec![order], vec![data_type]) +} + +/// Like [`sorting_stage`], for a producer with columns of `types`, sorted by `orders`. +pub fn sorting_stage_by( + producer: u32, + partitioning: Partitioning, + orders: Vec, + types: Vec, +) -> Stage { + let scan = Operator { + op_struct: Some(OpStruct::Scan(Scan { + fields: types.clone(), + source: "test".to_string(), + })), + ..Default::default() + }; + Stage { + plan: Some(Operator { + children: vec![scan], + op_struct: Some(OpStruct::Sort(Sort { + sort_orders: orders.clone(), + ..Default::default() + })), + ..Default::default() + }), + inputs: vec![StageInput { + producer, + partitioning: Some(partitioning), + broadcast: false, + }], + num_partitions: 1, + output_schema: types.into_iter().map(field).collect(), + output_ordering: orders, + ..Default::default() + } +} + +/// `input` with a projection of `exprs` on top of its block, producing columns of `types`. +pub fn projecting_stage( + mut input: Stage, + exprs: Vec, + types: Vec, +) -> Stage { + input.plan = Some(Operator { + op_struct: Some(OpStruct::Projection(Projection { + project_list: exprs, + })), + children: vec![input.plan.take().unwrap()], + ..Default::default() + }); + input.output_schema = types.into_iter().map(field).collect(); + input +} + +fn expr(expr_struct: ExprStruct) -> Expr { + Expr { + expr_struct: Some(expr_struct), + ..Default::default() + } +} + +pub fn long_literal(value: i64) -> Expr { + expr(ExprStruct::Literal(Literal { + value: Some(Value::LongVal(value)), + datatype: Some(long_type()), + is_null: false, + })) +} + +pub fn null_long() -> Expr { + expr(ExprStruct::Literal(Literal { + value: None, + datatype: Some(long_type()), + is_null: true, + })) +} + +pub fn remainder(left: Expr, right: Expr) -> Expr { + expr(ExprStruct::Remainder(Box::new(MathExpr { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + return_type: Some(long_type()), + ..Default::default() + }))) +} + +pub fn equal(left: Expr, right: Expr) -> Expr { + expr(ExprStruct::Eq(Box::new(BinaryExpr { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + }))) +} + +pub fn if_else(condition: Expr, then: Expr, otherwise: Expr) -> Expr { + expr(ExprStruct::If(Box::new(IfExpr { + if_expr: Some(Box::new(condition)), + true_expr: Some(Box::new(then)), + false_expr: Some(Box::new(otherwise)), + }))) +} + +pub fn cast(child: Expr, data_type: spark_expression::DataType) -> Expr { + expr(ExprStruct::Cast(Box::new(Cast { + child: Some(Box::new(child)), + datatype: Some(data_type), + timezone: "UTC".to_string(), + ..Default::default() + }))) +} + +/// A sort key: column `index`, of type `data_type`. +pub fn sort_key( + index: i32, + data_type: spark_expression::DataType, + descending: bool, + nulls_first: bool, +) -> Expr { + expr(ExprStruct::SortOrder(Box::new(SortOrder { + child: Some(Box::new(bound(index, data_type))), + direction: if descending { + SortDirection::Descending + } else { + SortDirection::Ascending + } as i32, + null_ordering: if nulls_first { + NullOrdering::NullsFirst + } else { + NullOrdering::NullsLast + } as i32, + }))) +} + pub fn hash_on_first_column(num_partitions: i32) -> Partitioning { Partitioning { partitioning_struct: Some(PartitioningStruct::HashPartition(HashPartition { @@ -140,16 +294,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 +329,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/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 a80cdd76f1b..27badac27a1 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -518,6 +518,33 @@ 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 in Spark's tasks. The driver plans the query and submits it to the " + + "cluster's scheduler (`spark.comet.nativeExecutors.scheduler.url`), and Spark tasks " + + "then read the result partitions from the cluster's executors, so the Spark " + + "executors need a libcomet built with the `native-executors` Cargo feature and " + + "network access to the cluster. Commands, writes and `CommandResultExec` stay in " + + "Spark. A query that cannot run entirely in the Comet pipeline and on the cluster " + + "fails, naming the operator that cannot. Cancellation is not supported: a killed " + + "Spark job leaves its query running on the cluster. " + + "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) 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..1da701743c0 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/NativeExecutorsClient.scala @@ -0,0 +1,93 @@ +/* + * 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 +} + +object NativeExecutorsClient { + + private val missingFeature = + "spark.comet.nativeExecutors.enabled needs a libcomet built with the `native-executors` " + + "Cargo feature" + + /** Whether the loaded libcomet has the client's native methods, which only the feature adds. */ + private lazy val hasNativeSupport: Boolean = { + val client = new NativeExecutorsClient() + try { + // A handle of 0 is a no-op once the method exists. + client.closePartition(0) + true + } catch { + case e: UnsatisfiedLinkError + if String.valueOf(e.getMessage).contains(classOf[NativeExecutorsClient].getName) => + false + } + } + + /** + * Fails with a clear error unless libcomet was built with the `native-executors` Cargo feature. + * Without it, the first call into the client would fail with a bare `UnsatisfiedLinkError`. + */ + def requireNativeSupport(): Unit = { + if (!hasNativeSupport) { + throw new CometRuntimeException(missingFeature) + } + } +} 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..23461908277 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeExecutorsExec.scala @@ -0,0 +1,622 @@ +/* + * 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, DynamicPruning, DynamicPruningExpression, 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] = { + NativeExecutorsClient.requireNativeSupport() + 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] + // A Spark executor needs a libcomet built with the feature too. + NativeExecutorsClient.requireNativeSupport() + val client = new NativeExecutorsClient() + val nativeUtil = new NativeUtil() + // The task's completion listener closes nativeUtil, but only once there is one to register. + val handle = + try { + client.openPartition( + partition.endpoint.getLocation, + partition.endpoint.getTicket.toByteArray) + } catch { + case t: Throwable => + nativeUtil.close() + throw t + } + 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 + try { + current.foreach(_.close()) + current = None + pending.foreach(_.close()) + pending = None + } finally { + // Each step runs even if an earlier one throws. + try client.closePartition(handle) + finally 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 in the Comet pipeline. + */ +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 = { + NativeExecutorsClient.requireNativeSupport() + offloadPlan(original, converted) + } + + private def offloadPlan(original: SparkPlan, converted: SparkPlan): SparkPlan = + converted match { + case toObjects: DeserializeToObjectExec => + toObjects.withNewChildren(Seq(offloadPlan(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 block of Comet operators that run as one Rust plan, 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 block of Comet operators 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 is not a Comet operator that runs in Rust") + } + 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) + val counts = perPartition.map { case (key, data) => key -> data.length } + require( + counts.values.toSet.size == 1, + s"The scans of one block must have the same number of partitions, but found $counts") + (0 until counts.values.head).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 be built")) + 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 + // These read or generate their rows inside the block. + case scan: CometNativeScanExec => + rejectDynamicPruning(scan, scan.partitionFilters) + case scan: CometIcebergNativeScanExec => + rejectDynamicPruning(scan, scan.runtimeFilters) + case _: CometRangeExec => + // These stand for rows that the JVM feeds to the plan, which a cluster has no way to do. + case _: CometScanWrapper | _: CometSinkPlaceHolder => + fail(plan, "its input comes from the JVM, which a cluster cannot read") + case _: CometNativeWriteExec | _: CometWriteFilesExec | _: CometIcebergWriteExec => + fail( + plan, + "Comet Native Executors do not support writes yet " + + "(https://github.com/apache/datafusion-comet/issues/7371)") + // Any other scan: its files are not described to the cluster. + case leaf: CometLeafExec => + fail(leaf, "Comet Native Executors do not support it yet") + // 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 + } + + /** + * A scan that prunes partitions with a dynamic filter lists its files only once the filter's + * subquery has run, which the driver cannot do before it builds the query. That includes the + * placeholder AQE leaves in a filter until it plans the subquery. + */ + private def rejectDynamicPruning(scan: SparkPlan, filters: Seq[Expression]): Unit = { + val dynamic = CometScanUtils + .filterUnusedDynamicPruningExpressions(filters) + .exists(_.exists { + case _: DynamicPruningExpression | _: DynamicPruning => true + case _ => false + }) + if (dynamic) { + fail( + scan, + "it prunes partitions with a dynamic filter (dynamic partition pruning), which " + + "Comet Native Executors do not support yet " + + "(https://github.com/apache/datafusion-comet/issues/7370)") + } + } + + 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 is not supported by Comet")) + } + 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 is not supported by Comet")) + } + 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..1287230dce8 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeExecutorsSuite.scala @@ -0,0 +1,451 @@ +/* + * 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", + // Without a cluster, only tests that the planner refuses run, and those never connect. + CometConf.COMET_NATIVE_EXECUTORS_SCHEDULER_URL.key -> + Option(schedulerUrl).getOrElse("http://localhost:50060"))(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) + } + } + } + } + + /** The messages of the exception that planning `df` on the cluster throws and its causes. */ + private def refusalMessages(df: => DataFrame): Seq[String] = { + // withSQLConf returns Unit on Spark 3.x, so the messages come out through a var. + var messages: Seq[String] = Nil + onCluster { + // A new Dataset, since a Dataset keeps the plan it was first run with. + val error = intercept[Exception](df.queryExecution.executedPlan) + messages = Iterator + .iterate[Throwable](error)(_.getCause) + .takeWhile(_ != null) + .map(e => String.valueOf(e.getMessage)) + .toSeq + } + messages + } + + /** Checks that the cluster refuses `df`, naming `operator` and giving `reason`. */ + private def checkRefused(df: => DataFrame, operator: String, reason: String): Unit = { + val messages = refusalMessages(df) + assert( + messages.exists(message => + message.contains("Comet Native Executors cannot run this query") && + message.contains(operator) && message.contains(reason)), + messages) + } + + 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") { + checkRefused(spark.read.parquet(path).orderBy("v"), "Sort", "") + } + } + } + + // The tests below fail while planning, so they need no cluster. + + test("a scan that stays on the JVM fails, naming the operator that reads from it") { + withTable { path => + // The scan stays in Spark, and Comet converts its columnar batches to Arrow batches for the + // operators above it. + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + CometConf.COMET_NATIVE_SCAN_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_PARQUET_ENABLED.key -> "true") { + checkRefused( + spark.read.parquet(path).selectExpr("id + 1 AS i"), + "CometSparkColumnarToColumnar", + "do not support it yet") + } + } + } + + test("a scan with dynamic partition pruning fails, naming the scan") { + withTempDir { dir => + val factPath = s"${dir.getAbsolutePath}/fact.parquet" + val dimPath = s"${dir.getAbsolutePath}/dim.parquet" + withSQLConf(CometConf.COMET_EXEC_ENABLED.key -> "false") { + val fact = spark + .range(0, 100) + .selectExpr("id AS fact_id", "date_add(date'2020-01-01', cast(id % 10 AS int)) AS d") + fact.write.partitionBy("d").parquet(factPath) + spark + .range(0, 10) + .selectExpr("id AS dim_id", "date_add(date'2020-01-01', cast(id AS int)) AS dim_d") + .write + .parquet(dimPath) + } + // Without AQE, the pruning filter is in the plan when Comet converts it. + withSQLConf( + SQLConf.USE_V1_SOURCE_LIST.key -> "parquet", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + def query = { + spark.read.parquet(factPath).createOrReplaceTempView("dpp_fact") + spark.read.parquet(dimPath).createOrReplaceTempView("dpp_dim") + sql("select * from dpp_fact join dpp_dim on d = dim_d where dim_id > 7") + } + checkRefused(query, "CometNativeScan", "dynamic partition pruning") + } + } + } +}