From c9c10a164d4ab32ad0d7cd472898471c2e99531d Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 6 Nov 2026 13:02:56 -0700 Subject: [PATCH 1/4] feat: add the JNI client of Comet Native Executors behind a Cargo feature Part of #7288. The client submits a query to a cluster's scheduler over Arrow Flight and reads the result partitions from the executors that hold them. It is compiled only with the native-executors feature, which is off by default, so the standard build links no gRPC client. The wire format gains the result endpoints the client hands to the JVM and the stage fields the driver fills in. --- native/Cargo.lock | 143 +++++++++ native/core/Cargo.toml | 6 + native/core/src/execution/jni_api.rs | 2 +- .../src/execution/native_executors/client.rs | 271 ++++++++++++++++++ .../src/execution/native_executors/mod.rs | 2 + native/proto/src/proto/native_executors.proto | 25 ++ 6 files changed, 448 insertions(+), 1 deletion(-) create mode 100644 native/core/src/execution/native_executors/client.rs diff --git a/native/Cargo.lock b/native/Cargo.lock index eb93db1252b..f82fa1611cb 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -347,6 +347,27 @@ dependencies = [ "num-traits", ] +[[package]] +name = "arrow-flight" +version = "59.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43f998a8cc9ffdc37731b66c0414dccf4f4f11a03d9136d383ef10f95181cd7d" +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.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d5848f5a9e8117cf94f2f90faac24a87ff81600ff8628c1e3deaab572e76bfa1" + [[package]] name = "miniz_oxide" version = "0.8.9" @@ -6883,6 +6981,47 @@ dependencies = [ "tokio", ] +[[package]] +name = "tonic" +version = "0.14.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2cdfc673aabd30fe8fee92e65837b644c37c5d47f0f97ae6d5e0f1237d15e1b" +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.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eba464054e2226f182a34489b8b079bf4c6e3e7ab5fd7fc30e024f24aa92a484" +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..fe842678cbf 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -1234,7 +1234,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/mod.rs b/native/core/src/execution/native_executors/mod.rs index fc0b9d6b2f4..f7c152d3fba 100644 --- a/native/core/src/execution/native_executors/mod.rs +++ b/native/core/src/execution/native_executors/mod.rs @@ -24,6 +24,8 @@ //! 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}; diff --git a/native/proto/src/proto/native_executors.proto b/native/proto/src/proto/native_executors.proto index 6c0ba50afb4..bb8df9489ae 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,16 @@ 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. The driver sends no range bounds, since + // Spark would sample for them. A cluster that can read the producer's partitions as they are, + // sort each, and merge them in this order may do that; one that ignores this field reads the + // exchange as one partition, which the block then sorts. + repeated spark.spark_expression.Expr output_ordering = 6; + // Spark's estimate of the stage's output size in bytes, when it has one. A scheduler may use + // it to decide, for example, whether measuring a join's build side before shuffling its probe + // side pays off. A scheduler that has no use for it ignores it. + optional uint64 estimated_output_bytes = 7; } // How a scheduler sends a stage to its executors, inside the task plans Ballista ships. The @@ -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; +} From f9a57866c064e53ae98c216ede4d5b7ae48b680e Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 6 Nov 2026 13:14:02 -0700 Subject: [PATCH 2/4] feat: add the Comet Native Executors driver, off by default Part of #7288. With spark.comet.nativeExecutors.enabled, CometRule replaces the plan Comet would execute with CometNativeExecutorsExec. The driver turns the plan into a stage graph, submits it to a cluster's scheduler through the JNI client, and Spark tasks read the result partitions from the executors over Arrow Flight. A query that cannot run entirely in the Comet pipeline fails, naming the operator. The two configs are internal, and the suite starts a cluster from COMET_NATIVE_EXECUTORS_BIN_DIR and skips without it. --- pom.xml | 4 +- .../scala/org/apache/comet/CometConf.scala | 27 + .../apache/comet/NativeExecutorsClient.scala | 93 +++ .../org/apache/comet/rules/CometRule.scala | 5 + .../sql/comet/CometNativeExecutorsExec.scala | 622 ++++++++++++++++++ .../exec/CometNativeExecutorsSuite.scala | 451 +++++++++++++ 6 files changed, 1200 insertions(+), 2 deletions(-) create mode 100644 spark/src/main/scala/org/apache/comet/NativeExecutorsClient.scala create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/CometNativeExecutorsExec.scala create mode 100644 spark/src/test/scala/org/apache/comet/exec/CometNativeExecutorsSuite.scala 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..663aa73939e --- /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. A cluster that supports + // it merges them in this order instead of range-partitioning the exchange. Without it, or + // on a cluster that ignores it, the exchange is read 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. A scheduler can use 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 stage's sort order tells the cluster how to merge 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") + } + } + } +} From cf20a8ecb2d6f4e2c823636691656e164c536dec Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Fri, 6 Nov 2026 13:15:42 -0700 Subject: [PATCH 3/4] docs: describe running the Comet Native Executors driver, and check it in CI Part of #7288. Adds the contributor guide page for the driver, with how to run it and its current limitations. Registers the manual suite with the suite check, and lints and tests the native-executors feature in the Rust job. --- .github/actions/rust-test/action.yaml | 14 ++ dev/ci/check-suites.py | 1 + docs/source/contributor-guide/index.md | 1 + .../contributor-guide/native_executors.md | 124 ++++++++++++++++++ native/native-executors/README.md | 7 +- 5 files changed, 144 insertions(+), 3 deletions(-) create mode 100644 docs/source/contributor-guide/native_executors.md 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/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..03c46d53308 --- /dev/null +++ b/docs/source/contributor-guide/native_executors.md @@ -0,0 +1,124 @@ + + +# 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 covers the Spark driver side: how to run a query on a cluster and +what does not work yet. The driver never links Ballista. It sends the stage graph to the +scheduler's Arrow Flight service through a small client in libcomet, and the Spark tasks that read +the results use the same client. The cluster features that some queries rely on are in +[#7289](https://github.com/apache/datafusion-comet/pull/7289) and are not part of the cluster in +this repository yet. [Current limitations](#current-limitations) says which. + +## 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, with native shuffle (`spark.comet.shuffle.mode=native`, as in the + [installation guide](../user-guide/latest/installation.md)), because the driver turns only + native shuffle exchanges into the edges of the stage graph. + +4. Check that the query ran on the cluster: its physical plan has a `CometNativeExecutors` leaf + (the `CometNativeExecutorsExec` class) 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. It is not part of CI. Its +tests that need a cluster are skipped when the variable is not set. The tests that check what the +driver refuses to plan run without a cluster, but they still need a libcomet built with the +`native-executors` feature. + +## 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). +- Only native shuffle exchanges. A query whose exchanges use JVM shuffle is refused. +- `LIMIT` and `show()`: [#7365](https://github.com/apache/datafusion-comet/issues/7365). +- Global sorts. The driver sends a range exchange's sort order but no range bounds. The cluster in + this repository reads the exchange as one partition, so every global sort runs in a single task. + [#7289](https://github.com/apache/datafusion-comet/pull/7289) adds a cluster that sorts the + partitions in parallel and merges them in the driver's order, except for sorts by floats or + nested types, which still run in a single task: + [#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). The driver sends Spark's + estimate of each stage's output size, which the cluster in this repository ignores. Moving a join + whose build side turns out small into the stage that produces its probe side is in + [#7289](https://github.com/apache/datafusion-comet/pull/7289). +- 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. diff --git a/native/native-executors/README.md b/native/native-executors/README.md index a271ba863f9..0265ac08dfc 100644 --- a/native/native-executors/README.md +++ b/native/native-executors/README.md @@ -25,9 +25,10 @@ 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 Spark driver side is in Comet's Spark plugin. See +[Comet Native Executors](https://github.com/apache/datafusion-comet/blob/main/docs/source/contributor-guide/native_executors.md) +for how to run a query on the cluster and what does not work yet. ## The scheduler's Flight service From f1819d5ba5f7fe23be1023dae1c8d110e139ae06 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Tue, 10 Nov 2026 08:27:29 -0700 Subject: [PATCH 4/4] fix: plan Comet Native Executors queries without side effects, and trim what no cluster here uses Review feedback on the driver: - Planning only validates the stage graph. The files a scan reads are listed, and the subqueries in its filters run, when the query executes, so EXPLAIN no longer contacts the cluster and the per-file plans no longer stay on the plan node. - A scan that lists no files gets one empty partition, so a query over an empty table or a fully pruned partition returns no rows instead of failing in the scheduler. - The offload follows the conversion's own guards: Comet loaded, exec enabled, not streaming. - A hash exchange keyed by an expression is refused naming the exchange, since the cluster hashes columns. - Remove the Stage fields output_ordering and estimated_output_bytes, which nothing in this repository reads, and the client's handling of a scheduler that serves partitions itself. - Ask libcomet whether it was built with the feature through NativeBase.isFeatureEnabled. - Add CometNativeExecutorsPlanSuite, which checks the stage graph, the refusals and the subquery rewrite in the regular CI matrix, and compare the cluster suite with Spark's answer with Comet off. - Document the blocked driver thread on cancellation and the generic exception for cluster errors, and say which TPC-H tests need the cluster of #7289. --- .github/actions/rust-test/action.yaml | 13 +- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + .../contributor-guide/native_executors.md | 29 ++- .../src/execution/native_executors/client.rs | 46 ++-- native/core/src/lib.rs | 5 + native/proto/src/proto/native_executors.proto | 11 - .../java/org/apache/comet/NativeBase.java | 3 +- .../apache/comet/NativeExecutorsClient.scala | 15 +- .../org/apache/comet/rules/CometRule.scala | 26 ++- .../sql/comet/CometNativeExecutorsExec.scala | 143 ++++++++----- .../exec/CometNativeExecutorsSuite.scala | 33 ++- .../comet/CometNativeExecutorsPlanSuite.scala | 197 ++++++++++++++++++ 13 files changed, 389 insertions(+), 134 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/sql/comet/CometNativeExecutorsPlanSuite.scala diff --git a/.github/actions/rust-test/action.yaml b/.github/actions/rust-test/action.yaml index 6a4d430edba..649be3eee36 100644 --- a/.github/actions/rust-test/action.yaml +++ b/.github/actions/rust-test/action.yaml @@ -39,18 +39,13 @@ runs: cd native cargo clippy --color=never --all-targets --workspace -- -D warnings - - name: Check native Python UDF feature + # Both features are off by default, so nothing above compiles them. The native-executors one + # is the JNI client for Comet Native Executors. + - name: Check native Python UDF and native executors features shell: bash run: | 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 + cargo clippy --color=never -p datafusion-comet --all-targets --features python-udf,native-executors -- -D warnings - name: Check compilation shell: bash diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 7f70f81122c..b69a0a85385 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -589,6 +589,7 @@ jobs: org.apache.comet.exec.CometMergeRowsSuite org.apache.comet.exec.CometTopKSuite org.apache.comet.exec.CometEmptyRelationExecSuite + org.apache.spark.sql.comet.CometNativeExecutorsPlanSuite org.apache.comet.exec.CometRangeExecSuite org.apache.comet.exec.CometInMemoryCacheSuite org.apache.comet.exec.CometInMemoryCachePruningSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 79cf33713ec..d3a7132c78f 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -224,6 +224,7 @@ jobs: org.apache.comet.exec.CometMergeRowsSuite org.apache.comet.exec.CometTopKSuite org.apache.comet.exec.CometEmptyRelationExecSuite + org.apache.spark.sql.comet.CometNativeExecutorsPlanSuite org.apache.comet.exec.CometRangeExecSuite org.apache.comet.exec.CometInMemoryCacheSuite org.apache.comet.exec.CometInMemoryCachePruningSuite diff --git a/docs/source/contributor-guide/native_executors.md b/docs/source/contributor-guide/native_executors.md index 03c46d53308..e10ac242d2c 100644 --- a/docs/source/contributor-guide/native_executors.md +++ b/docs/source/contributor-guide/native_executors.md @@ -68,7 +68,11 @@ this repository yet. [Current limitations](#current-limitations) says which. `COMET_NATIVE_EXECUTORS_BIN_DIR`. Its header comment gives the commands. It is not part of CI. Its tests that need a cluster are skipped when the variable is not set. The tests that check what the driver refuses to plan run without a cluster, but they still need a libcomet built with the -`native-executors` feature. +`native-executors` feature. Its TPC-H tests need the cluster of +[#7289](https://github.com/apache/datafusion-comet/pull/7289): on the cluster in this repository, +only q1, q6, q11 and q15 pass. `CometNativeExecutorsPlanSuite` runs in CI. It takes the plan Comet +would execute and checks the stage graph, the queries the driver refuses, and how a scalar +subquery's result goes into the plan, with neither a cluster nor the feature. ## Current limitations @@ -78,25 +82,28 @@ Comet Native Executors are a first version. The sub-issues of - Whole queries only. See [below](#whole-queries-only). - Only native shuffle exchanges. A query whose exchanges use JVM shuffle is refused. - `LIMIT` and `show()`: [#7365](https://github.com/apache/datafusion-comet/issues/7365). -- Global sorts. The driver sends a range exchange's sort order but no range bounds. The cluster in - this repository reads the exchange as one partition, so every global sort runs in a single task. +- Global sorts. The cluster in this repository reads a range exchange as one partition, so every + global sort runs in a single task. [#7289](https://github.com/apache/datafusion-comet/pull/7289) adds a cluster that sorts the - partitions in parallel and merges them in the driver's order, except for sorts by floats or - nested types, which still run in a single task: - [#7366](https://github.com/apache/datafusion-comet/issues/7366). + partitions in parallel and merges them, except for sorts by floats or nested types, which still + run in a single task: [#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). The driver sends Spark's - estimate of each stage's output size, which the cluster in this repository ignores. Moving a join - whose build side turns out small into the stage that produces its probe side is in + [#7369](https://github.com/apache/datafusion-comet/issues/7369). Moving a join whose build side + turns out small into the stage that produces its probe side is in [#7289](https://github.com/apache/datafusion-comet/pull/7289). - 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. +- Cancellation: [#7372](https://github.com/apache/datafusion-comet/issues/7372). The driver thread + waits for the cluster in native code until the query finishes, so cancelling the Spark job or + interrupting the thread does not stop the wait, and a killed Spark job leaves its query running + on the cluster. +- Errors. A failure on the cluster reaches Spark as a `CometNativeException` with the cluster's + message, not as the exception Spark would raise. For example, a division by zero with ANSI mode on + is not a `SparkArithmeticException` with the `DIVIDE_BY_ZERO` error class. - 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. diff --git a/native/core/src/execution/native_executors/client.rs b/native/core/src/execution/native_executors/client.rs index 7d3843acef3..35be17f0aa1 100644 --- a/native/core/src/execution/native_executors/client.rs +++ b/native/core/src/execution/native_executors/client.rs @@ -60,18 +60,14 @@ const KEEP_ALIVE_INTERVAL: Duration = Duration::from_secs(30); 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. +/// message that names the setting to change: 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, - }, + FlightError::Tonic(status) if status.code() == Code::Internal => { + let message = status.message(); + message.contains("Failed to open partition file") && message.contains("NotFound") + } _ => false, }; if expired { @@ -132,11 +128,11 @@ async fn submit(scheduler_url: &str, mut query: Vec) -> CometResult bool { match feature { "jemalloc" => cfg!(feature = "jemalloc"), "hdfs-opendal" => cfg!(feature = "hdfs-opendal"), + "native-executors" => cfg!(feature = "native-executors"), _ => false, // Unknown features return false } } @@ -376,6 +377,10 @@ mod tests { is_feature_enabled("hdfs-opendal"), cfg!(feature = "hdfs-opendal") ); + assert_eq!( + is_feature_enabled("native-executors"), + cfg!(feature = "native-executors") + ); assert!(!is_feature_enabled("no-such-feature")); } diff --git a/native/proto/src/proto/native_executors.proto b/native/proto/src/proto/native_executors.proto index bb8df9489ae..a324dc690a1 100644 --- a/native/proto/src/proto/native_executors.proto +++ b/native/proto/src/proto/native_executors.proto @@ -23,7 +23,6 @@ syntax = "proto3"; package spark.spark_native_executors; -import "expr.proto"; import "operator.proto"; import "partitioning.proto"; @@ -59,16 +58,6 @@ 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. The driver sends no range bounds, since - // Spark would sample for them. A cluster that can read the producer's partitions as they are, - // sort each, and merge them in this order may do that; one that ignores this field reads the - // exchange as one partition, which the block then sorts. - repeated spark.spark_expression.Expr output_ordering = 6; - // Spark's estimate of the stage's output size in bytes, when it has one. A scheduler may use - // it to decide, for example, whether measuring a join's build side before shuffling its probe - // side pays off. A scheduler that has no use for it ignores it. - optional uint64 estimated_output_bytes = 7; } // How a scheduler sends a stage to its executors, inside the task plans Ballista ships. The diff --git a/spark/src/main/java/org/apache/comet/NativeBase.java b/spark/src/main/java/org/apache/comet/NativeBase.java index b9349057c99..282d0e81b5b 100644 --- a/spark/src/main/java/org/apache/comet/NativeBase.java +++ b/spark/src/main/java/org/apache/comet/NativeBase.java @@ -363,7 +363,8 @@ public static void releaseNative() throws Throwable { /** * Check if a specific feature is enabled in the native library. * - * @param featureName The name of the feature to check (e.g., "jemalloc", "hdfs-opendal") + * @param featureName The name of the feature to check (e.g., "jemalloc", "hdfs-opendal", + * "native-executors") * @return true if the feature is enabled, false otherwise */ public static native boolean isFeatureEnabled(String featureName); diff --git a/spark/src/main/scala/org/apache/comet/NativeExecutorsClient.scala b/spark/src/main/scala/org/apache/comet/NativeExecutorsClient.scala index 1da701743c0..3b53f41d7fa 100644 --- a/spark/src/main/scala/org/apache/comet/NativeExecutorsClient.scala +++ b/spark/src/main/scala/org/apache/comet/NativeExecutorsClient.scala @@ -67,19 +67,8 @@ object NativeExecutorsClient { "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 - } - } + /** Whether the loaded libcomet was built with the feature, which adds the client's methods. */ + private lazy val hasNativeSupport: Boolean = NativeBase.isFeatureEnabled("native-executors") /** * Fails with a clear error unless libcomet was built with the `native-executors` Cargo feature. 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 07b2eea6863..bea93ec4c52 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometRule.scala @@ -53,16 +53,28 @@ object CometRule { CometCacheColumnarRule(preview)) /** - * Whether plan-only mode applies to `plan`, so that Comet only reports the plan it would - * execute and Spark executes `plan` unchanged. Mirrors the conversion rules' own guards; - * plan-only is scoped to exec being enabled. + * Whether the conversion rules convert `plan`: they mirror these guards. Plan-only mode and + * Comet Native Executors, which each rely on the conversion, are scoped to it. */ - private[comet] def planOnlyApplies(conf: SQLConf, plan: SparkPlan): Boolean = - CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf) && - isCometLoaded(conf) && + private def conversionApplies(conf: SQLConf, plan: SparkPlan): Boolean = + isCometLoaded(conf) && !ShimCometStreaming.isStreamingPlan(plan) && CometConf.COMET_EXEC_ENABLED.get(conf) + /** + * Whether plan-only mode applies to `plan`, so that Comet only reports the plan it would + * execute and Spark executes `plan` unchanged. + */ + private[comet] def planOnlyApplies(conf: SQLConf, plan: SparkPlan): Boolean = + CometConf.COMET_EXPLAIN_PLAN_ONLY_ENABLED.get(conf) && conversionApplies(conf, plan) + + /** Whether `plan` runs on a Comet Native Executor cluster instead of Spark's executors. */ + private def nativeExecutorsApply( + conf: SQLConf, + plan: SparkPlan, + queryStagePrep: Boolean): Boolean = + conversionApplies(conf, plan) && CometNativeExecutors.applies(conf, plan, queryStagePrep) + /** * Canonical hashes of the subquery plans reported for the query this thread is preparing. Spark * prepares a query's subqueries synchronously, before the query itself, so the scope is reset @@ -164,7 +176,7 @@ case class CometRule(session: SparkSession, queryStagePrep: Boolean = false) if (CometRule.planOnlyApplies(conf, plan)) { reportPlanOnlyCoverage(plan) plan - } else if (CometNativeExecutors.applies(conf, plan, queryStagePrep)) { + } else if (CometRule.nativeExecutorsApply(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)) 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 index 663aa73939e..b503e43f442 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeExecutorsExec.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometNativeExecutorsExec.scala @@ -23,7 +23,6 @@ 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 @@ -48,6 +47,7 @@ import org.apache.comet.{CometConf, CometExecIterator, CometExplainInfo, CometRu 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 import org.apache.comet.serde.OperatorOuterClass.Operator import org.apache.comet.serde.operator.schema2Proto import org.apache.comet.vector.NativeUtil @@ -56,19 +56,22 @@ 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. + * This leaf replaces the plan Comet would otherwise execute. Planning only checks that the plan + * can run on a cluster ([[CometNativeExecutors.validate]]). `doExecuteColumnar` builds the query, + * which lists the files the scans read, submits it 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 nativePlan + * the plan Comet would execute, which [[CometNativeExecutors.buildQuery]] turns into the + * `QueryRequest` the cluster runs * @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], + nativePlan: SparkPlan, schedulerUrl: String, subqueryResults: Seq[Expression] = Nil) extends LeafExecNode { @@ -81,7 +84,7 @@ case class CometNativeExecutorsExec( override protected def doExecuteColumnar(): RDD[ColumnarBatch] = { NativeExecutorsClient.requireNativeSupport() val request = CometNativeExecutors.withSubqueryResults( - query, + CometNativeExecutors.buildQuery(nativePlan).toByteArray, subqueryResults.collect { case sub: ScalarSubquery => sub }) val endpoints = ResultEndpoints.parseFrom(new NativeExecutorsClient().submit(schedulerUrl, request)) @@ -92,7 +95,7 @@ case class CometNativeExecutorsExec( } override def stringArgs: Iterator[Any] = - Iterator(output, s"${query.length} bytes", schedulerUrl) + Iterator(output, schedulerUrl) } private class CometNativeExecutorsPartition(override val index: Int, val endpoint: ResultEndpoint) @@ -194,7 +197,8 @@ object CometNativeExecutors { /** * 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. + * command, not a stage AQE has split off, and not a plan already replaced. The caller checks + * that Comet executes the plan at all. */ def applies(conf: SQLConf, plan: SparkPlan, queryStagePrep: Boolean): Boolean = CometConf.COMET_NATIVE_EXECUTORS_ENABLED.get(conf) && @@ -222,25 +226,46 @@ object CometNativeExecutors { case toObjects: DeserializeToObjectExec => toObjects.withNewChildren(Seq(offloadPlan(toObjects.child, toObjects.child))) case native => + validate(native) val exec = CometNativeExecutorsExec( native.output, - buildQuery(native).toByteArray, + native, CometConf.COMET_NATIVE_EXECUTORS_SCHEDULER_URL.get(), scalarSubqueries(native)) native.logicalLink.orElse(original.logicalLink).foreach(exec.setLogicalLink) exec } + /** + * Checks that a cluster can run `plan`, the plan Comet would execute, without running anything: + * it lists no files and resolves no subquery, so it is safe at planning time. + * + * @throws CometRuntimeException + * naming the first operator that cannot run on a cluster + */ + def validate(plan: SparkPlan): Unit = stageGraph(plan, withFiles = false) + /** * 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. + * that feeds the block. A stage that reads files has a plan for each partition, naming its + * files, so this lists them and resolves the subqueries in the scans' filters. It belongs at + * execution time. * * @throws CometRuntimeException * naming the first operator that cannot run on a cluster */ def buildQuery(plan: SparkPlan): QueryRequest = { - val builder = new StageGraphBuilder(plan) + val stages = stageGraph(plan, withFiles = true) + QueryRequest + .newBuilder() + .putAllConf(nativeConf()) + .addAllStages(stages.asJava) + .build() + } + + private def stageGraph(plan: SparkPlan, withFiles: Boolean): Seq[Stage] = { + val builder = new StageGraphBuilder(plan, withFiles) val root = plan match { case ColumnarToRowExec(child) => child case CometColumnarToRowExec(child) => child @@ -248,15 +273,11 @@ object CometNativeExecutors { case other => other } builder.stageOf(root) - QueryRequest - .newBuilder() - .putAllConf(nativeConf()) - .addAllStages(builder.stages.asJava) - .build() + builder.stages.toSeq } /** The scalar subqueries of `plan`'s expressions, through exchanges, each once. */ - private def scalarSubqueries(plan: SparkPlan): Seq[Expression] = { + private[comet] def scalarSubqueries(plan: SparkPlan): Seq[Expression] = { val found = ArrayBuffer.empty[ScalarSubquery] def walk(node: SparkPlan): Unit = { node.expressions.foreach(_.foreach { @@ -381,7 +402,14 @@ object CometNativeExecutors { .setSinglePartition(PartitioningOuterClass.SinglePartition.newBuilder()) .build() - private class StageGraphBuilder(query: SparkPlan) { + /** The planning data of a partition that reads no files. */ + private val noFiles = OperatorOuterClass.NativeScan + .newBuilder() + .setFilePartition(OperatorOuterClass.SparkFilePartition.getDefaultInstance) + .build() + .toByteArray + + private class StageGraphBuilder(query: SparkPlan, withFiles: Boolean) { val stages: ArrayBuffer[Stage] = ArrayBuffer.empty private val stageIndexes = new IdentityHashMap[SparkPlan, Integer]() @@ -399,7 +427,8 @@ object CometNativeExecutors { 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 { + val scans = ArrayBuffer.empty[CometNativeScanExec] + val inputs = inputsOf(block, scans).map { case exchange: CometShuffleExchangeExec => stageInput(exchange.child)(_.setPartitioning(partitioningOf(exchange))) case broadcast => stageInput(broadcast.children.head)(_.setBroadcast(true)) @@ -407,7 +436,9 @@ object CometNativeExecutors { 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 (common, perPartition) = + if (withFiles) filesOf(block, scans.toSeq) + else (Map.empty[String, Array[Byte]], Map.empty[String, Array[Array[Byte]]]) val partitionPlans = if (perPartition.isEmpty) { Seq.empty } else { @@ -426,16 +457,24 @@ object CometNativeExecutors { val numPartitions = if (partitionPlans.nonEmpty) partitionPlans.length else block.outputPartitioning.numPartitions - // A stage that reads a range exchange sorts each of its partitions. A cluster that supports - // it merges them in this order instead of range-partitioning the exchange. Without it, or - // on a cluster that ignores it, the exchange is read 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 + addStage(plan, blockPlan, inputs, partitionPlans, numPartitions) + } + + /** + * The planning data of the scans in `block`. A block whose scans list no files at all, such + * as a table with no data or a partition filter that prunes everything, gets one empty + * partition for each, as Spark would return no rows from them. + */ + private def filesOf(block: CometNativeExec, scans: Seq[CometNativeScanExec]) + : (Map[String, Array[Byte]], Map[String, Array[Array[Byte]]]) = { + val found @ (_, perPartition) = PlanDataInjector.findAllPlanData(block) + if (perPartition.nonEmpty || scans.isEmpty) { + found } else { - Nil + ( + scans.map(scan => scan.sourceKey -> scan.commonData).toMap, + scans.map(scan => scan.sourceKey -> Array(noFiles)).toMap) } - addStage(plan, blockPlan, inputs, partitionPlans, numPartitions, outputOrdering) } /** An input of a stage that reads the stage running `producer`, as `configure` says. */ @@ -455,7 +494,7 @@ object CometNativeExecutors { 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) + addStage(take, blockPlan, Seq(input), Seq.empty, 1) } private def addStage( @@ -463,8 +502,7 @@ object CometNativeExecutors { blockPlan: Operator, inputs: Seq[StageInput], partitionPlans: Seq[ByteString], - numPartitions: Int, - outputOrdering: Seq[ExprOuterClass.Expr]): Int = { + numPartitions: Int): Int = { findJvmCallback(blockPlan).foreach { name => fail(plan, s"its native plan calls back into the JVM ($name)") } @@ -479,31 +517,20 @@ object CometNativeExecutors { .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. A scheduler can use 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. + * native shuffles and broadcasts, with a reused exchange resolved to the one it reuses. Adds + * the block's native Parquet scans to `scans`. */ - private def inputsOf(block: CometNativeExec): Seq[SparkPlan] = { + private def inputsOf( + block: CometNativeExec, + scans: ArrayBuffer[CometNativeScanExec]): Seq[SparkPlan] = { val inputs = ArrayBuffer.empty[SparkPlan] def walk(plan: SparkPlan): Unit = plan match { case reused: ReusedExchangeExec => @@ -515,6 +542,7 @@ object CometNativeExecutors { // These read or generate their rows inside the block. case scan: CometNativeScanExec => rejectDynamicPruning(scan, scan.partitionFilters) + scans += scan case scan: CometIcebergNativeScanExec => rejectDynamicPruning(scan, scan.runtimeFilters) case _: CometRangeExec => @@ -565,10 +593,16 @@ object CometNativeExecutors { 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")) + val keys = expressions.map { + // A cluster hashes columns of the stage's input. + case key: Attribute => + QueryPlanSerde + .exprToProto(key, exchange.child.output) + .getOrElse(fail(exchange, s"its partitioning key $key is not supported by Comet")) + case key => + fail( + exchange, + s"its partitioning key $key is not a column, which a cluster requires") } partitioning .setHashPartition( @@ -578,7 +612,8 @@ object CometNativeExecutors { .setNumPartitions(numPartitions)) .build() case RangePartitioning(ordering, numPartitions) => - // Without bounds: the stage's sort order tells the cluster how to merge instead. + // Without bounds, since Spark would sample for them: the cluster reads the exchange as + // one partition, which the block then sorts. val sortOrders = ordering.map { order => QueryPlanSerde .exprToProto(order, exchange.child.output) diff --git a/spark/src/test/scala/org/apache/comet/exec/CometNativeExecutorsSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometNativeExecutorsSuite.scala index 1287230dce8..041abd15b89 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometNativeExecutorsSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometNativeExecutorsSuite.scala @@ -47,6 +47,12 @@ import org.apache.comet.CometConf * COMET_NATIVE_EXECUTORS_BIN_DIR=native/target/debug ./mvnw test -Dtest=none \ * -Dsuites=org.apache.comet.exec.CometNativeExecutorsSuite * }}} + * + * The TPC-H tests need the cluster of https://github.com/apache/datafusion-comet/pull/7289. On + * the cluster in this repository, q1, q6, q11 and q15 pass and the other queries, which join, + * fail while the scheduler plans them. + * + * `CometNativeExecutorsPlanSuite` checks the stage graphs without a cluster, and runs in CI. */ class CometNativeExecutorsSuite extends CometTestBase with AdaptiveSparkPlanHelper { @@ -148,7 +154,17 @@ class CometNativeExecutorsSuite extends CometTestBase with AdaptiveSparkPlanHelp * 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) + checkClusterReturns(query, sparkAnswer(query), ordered) + + /** Spark's own answer, with Comet off, so that a bug the cluster shares with Comet shows. */ + private def sparkAnswer(query: => DataFrame): Seq[Row] = { + // withSQLConf returns Unit on Spark 3.x, so the rows come out through a var. + var rows: Seq[Row] = Nil + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + rows = query.collect().toSeq + } + rows + } /** Checks that the cluster runs `query` and returns `expected`, in order if `ordered`. */ private def checkClusterReturns( @@ -346,6 +362,21 @@ class CometNativeExecutorsSuite extends CometTestBase with AdaptiveSparkPlanHelp } } + nativeExecutorsTest("a scan that lists no files returns no rows") { + withTempPath { dir => + spark + .range(0, 100) + .selectExpr("id", "id % 7 AS k") + .write + .partitionBy("k") + .parquet(dir.getCanonicalPath) + // No partition has k = 99, so the scan has no files to read. + def table = spark.read.parquet(dir.getCanonicalPath).where("k = 99") + checkOnCluster(table) + checkOnCluster(table.groupBy().count()) + } + } + nativeExecutorsTest("scalar subquery") { withTable { path => spark.read.parquet(path).createOrReplaceTempView("t") diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeExecutorsPlanSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeExecutorsPlanSuite.scala new file mode 100644 index 00000000000..ae229f9f84c --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometNativeExecutorsPlanSuite.scala @@ -0,0 +1,197 @@ +/* + * 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 scala.jdk.CollectionConverters._ + +import org.apache.spark.SparkConf +import org.apache.spark.sql.{CometTestBase, DataFrame} +import org.apache.spark.sql.execution.{InputAdapter, ScalarSubquery, SparkPlan, WholeStageCodegenExec} +import org.apache.spark.sql.functions.{col, sum} +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.{CometConf, CometRuntimeException} +import org.apache.comet.serde.NativeExecutors.{QueryRequest, Stage} + +/** + * Plans queries for Comet Native Executors (see + * https://github.com/apache/datafusion-comet/issues/7288) without a cluster. The tests take the + * plan Comet would execute and turn it into the stage graph a cluster runs, so they need neither + * a libcomet built with the `native-executors` feature nor the cluster binaries. + * `CometNativeExecutorsSuite` runs the queries on a cluster. + */ +class CometNativeExecutorsPlanSuite extends CometTestBase { + + override protected def sparkConf: SparkConf = { + val conf = super.sparkConf + conf.set(CometConf.COMET_SHUFFLE_MODE.key, "native") + conf.set(SQLConf.SHUFFLE_PARTITIONS.key, "4") + // Without AQE, the plan is final once it is planned. + conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "false") + conf + } + + /** 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) + } + } + + /** + * The plan Comet would execute, as the planner sees it before Spark adds whole-stage codegen. + */ + private def cometPlan(df: DataFrame): SparkPlan = df.queryExecution.executedPlan match { + case codegen: WholeStageCodegenExec => + codegen.child.transform { case adapter: InputAdapter => adapter.child } + case plan => plan + } + + private def stagesOf(df: DataFrame): Seq[Stage] = + CometNativeExecutors.buildQuery(cometPlan(df)).getStagesList.asScala.toSeq + + private def refusal(df: DataFrame): String = + intercept[CometRuntimeException](CometNativeExecutors.validate(cometPlan(df))).getMessage + + test("an aggregate is a scan stage and a stage that reads its hash exchange") { + withTable { path => + val stages = stagesOf(spark.read.parquet(path).groupBy("k").agg(sum("v"))) + assert(stages.size == 2) + val Seq(scan, aggregate) = stages + // The scan stage runs one copy of its plan for each partition, naming that partition's files. + assert(scan.getInputsCount == 0) + assert(scan.getPartitionPlansCount > 1) + assert(scan.getNumPartitions == scan.getPartitionPlansCount) + assert(aggregate.getInputsCount == 1) + val input = aggregate.getInputs(0) + assert(input.getProducer == 0) + assert(input.getPartitioning.hasHashPartition) + assert(input.getPartitioning.getHashPartition.getNumPartitions == 4) + } + } + + test("a global aggregate reads a single-partition exchange") { + withTable { path => + val stages = stagesOf(spark.read.parquet(path).agg(sum("id"))) + assert(stages.size == 2) + assert(stages.last.getInputs(0).getPartitioning.hasSinglePartition) + } + } + + test("a broadcast join reads the build side whole") { + 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")) + val stages = stagesOf(left.join(right, left("k") === right("k2")).select("id", "v")) + assert(stages.flatMap(_.getInputsList.asScala).exists(_.getBroadcast), stages) + } + } + + test("a shuffled join reads both sides' hash exchanges") { + withSQLConf( + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.PREFER_SORTMERGEJOIN.key -> "true") { + withTable { path => + val left = spark.read.parquet(path).select("id", "k") + val right = spark.read.parquet(path).select(col("k").as("k2"), col("v")) + val stages = stagesOf(left.join(right, left("k") === right("k2")).select("id", "v")) + val join = stages.last + assert(join.getInputsCount == 2) + assert(join.getInputsList.asScala.forall(_.getPartitioning.hasHashPartition)) + } + } + } + + test("a scan that lists no files gets one empty partition") { + withTempPath { dir => + spark + .range(0, 100) + .selectExpr("id", "id % 7 AS k") + .write + .partitionBy("k") + .parquet(dir.getCanonicalPath) + val df = spark.read.parquet(dir.getCanonicalPath).where("k = 99").groupBy().count() + val scan = stagesOf(df).head + assert(scan.getPartitionPlansCount == 1) + assert(scan.getNumPartitions == 1) + } + } + + test("a scalar subquery's result replaces it in the plan") { + withTable { path => + spark.read.parquet(path).createOrReplaceTempView("t") + val df = sql("select id from t where v > (select avg(v) from t where k = 3)") + // Spark runs the subquery before it executes the plan. + df.collect() + val plan = cometPlan(df) + val subqueries = CometNativeExecutors.scalarSubqueries(plan).collect { + case sub: ScalarSubquery => sub + } + assert(subqueries.size == 1) + def plans(query: Array[Byte]): String = + QueryRequest.parseFrom(query).getStagesList.asScala.map(_.getPlan.toString).mkString + val query = CometNativeExecutors.buildQuery(plan).toByteArray + assert(plans(query).contains("subquery")) + val rewritten = CometNativeExecutors.withSubqueryResults(query, subqueries) + assert(!plans(rewritten).contains("subquery")) + } + } + + test("a query Comet does not execute runs in Spark") { + withSQLConf( + CometConf.COMET_NATIVE_EXECUTORS_ENABLED.key -> "true", + CometConf.COMET_ENABLED.key -> "false") { + val df = spark.range(10).selectExpr("id + 1 AS i") + assert(df.collect().map(_.getLong(0)).toSeq == (1L to 10L)) + assert(df.queryExecution.executedPlan.collect { case e: CometNativeExecutorsExec => + e + }.isEmpty) + } + } + + test("a sort that stays in Spark is refused, naming the operator") { + withTable { path => + withSQLConf(CometConf.COMET_EXEC_SORT_ENABLED.key -> "false") { + val message = refusal(spark.read.parquet(path).orderBy("v")) + assert(message.contains("Comet Native Executors cannot run this query"), message) + assert(message.contains("Sort"), message) + } + } + } + + test("a hash exchange on an expression is refused, naming the exchange") { + withSQLConf( + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.PREFER_SORTMERGEJOIN.key -> "true") { + withTable { path => + val left = spark.read.parquet(path).select("id", "k") + val right = spark.read.parquet(path).select(col("k").as("k2"), col("v")) + val message = refusal(left.join(right, left("k") + 1 === right("k2"))) + assert(message.contains("Exchange"), message) + assert(message.contains("is not a column"), message) + } + } + } +}