From 769268863e3181fb8e6964746e9bc1be5a781f04 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Mon, 12 Oct 2026 00:22:57 -0700 Subject: [PATCH 01/12] feat: run admitted local queries as one native graph Add an experimental local execution mode for applications with a local master, enabled by the internal spark.comet.exec.local.enabled option. A query is admitted only as a whole, from the plan Comet's ordinary conversion produces: one native block with no JVM input whose only leaf is a range. Such plans have no exchange, so they are admitted with AQE enabled. The single Spark task of the result creates one DataFusion graph for all partitions. Its memory pool, spill directories and session come from create_execution_context, so the pool acquires memory through the task's CometTaskMemoryManager. Each partition's operators are created by the planner for that partition, as a Spark task would create them. The native/local crate runs the partitions concurrently and delivers them in partition order, so results keep Spark's row order and limits return the same rows as Spark. Closing a query drops its streams on the calling thread, as releasePlan does, since every Tokio worker can be waiting for memory those streams hold. CometLocalExec reports UnknownPartitioning(1) so that a query over a cached local result still gets its required exchange. CometLocalResultExec hands collect and take results to the driver within the JVM, enforcing spark.driver.maxResultSize while copying. --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + dev/ci/check-ci-config.py | 4 + dev/ci/compute-changes.py | 12 + docs/source/contributor-guide/index.md | 1 + .../contributor-guide/local-execution.md | 119 ++++ .../contributor-guide/plugin_overview.md | 4 + native/Cargo.lock | 10 + native/Cargo.toml | 4 +- native/core/Cargo.toml | 1 + native/core/src/execution/jni_api.rs | 76 +-- native/core/src/execution/mod.rs | 2 +- native/core/src/lib.rs | 1 + native/core/src/local.rs | 249 +++++++++ native/local/Cargo.toml | 36 ++ native/local/src/handle.rs | 226 ++++++++ native/local/src/lib.rs | 297 ++++++++++ native/local/tests/bridge.rs | 303 ++++++++++ native/local/tests/execution.rs | 525 ++++++++++++++++++ pom.xml | 3 + .../apache/comet/local/CometLocalExec.scala | 88 +++ .../comet/local/CometLocalResultExec.scala | 122 ++++ .../apache/comet/local/CometLocalRule.scala | 135 +++++ .../comet/local/LocalQueryIterator.scala | 143 +++++ .../org/apache/comet/local/NativeLocal.scala | 58 ++ .../local/CometLocalExecutionSuite.scala | 410 ++++++++++++++ spark/pom.xml | 17 + .../scala/org/apache/comet/CometConf.scala | 10 + .../org/apache/comet/rules/CometRule.scala | 9 +- 29 files changed, 2828 insertions(+), 39 deletions(-) create mode 100644 docs/source/contributor-guide/local-execution.md create mode 100644 native/core/src/local.rs create mode 100644 native/local/Cargo.toml create mode 100644 native/local/src/handle.rs create mode 100644 native/local/src/lib.rs create mode 100644 native/local/tests/bridge.rs create mode 100644 native/local/tests/execution.rs create mode 100644 spark-local/src/main/scala/org/apache/comet/local/CometLocalExec.scala create mode 100644 spark-local/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala create mode 100644 spark-local/src/main/scala/org/apache/comet/local/CometLocalRule.scala create mode 100644 spark-local/src/main/scala/org/apache/comet/local/LocalQueryIterator.scala create mode 100644 spark-local/src/main/scala/org/apache/comet/local/NativeLocal.scala create mode 100644 spark-local/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 192dc8ffb32..0db0da6835e 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -539,6 +539,7 @@ jobs: org.apache.spark.CometRuntimeShutdownSuite org.apache.spark.CometTaskMemoryManagerSuite org.apache.spark.CometExecIteratorLifecycleSuite + org.apache.comet.local.CometLocalExecutionSuite org.apache.spark.CometPluginsDefaultSuite org.apache.spark.CometPluginsMemoryOverheadWarningSuite org.apache.spark.CometPluginsMemoryPoolFractionWarningSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 0729cb26f67..45124eb3137 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -243,6 +243,7 @@ jobs: org.apache.spark.CometRuntimeShutdownSuite org.apache.spark.CometTaskMemoryManagerSuite org.apache.spark.CometExecIteratorLifecycleSuite + org.apache.comet.local.CometLocalExecutionSuite org.apache.spark.CometPluginsDefaultSuite org.apache.spark.CometPluginsMemoryOverheadWarningSuite org.apache.spark.CometPluginsMemoryPoolFractionWarningSuite diff --git a/dev/ci/check-ci-config.py b/dev/ci/check-ci-config.py index 3832fda92dc..d65f14ab983 100644 --- a/dev/ci/check-ci-config.py +++ b/dev/ci/check-ci-config.py @@ -127,6 +127,10 @@ MVN_JOBS = BUILD_JOBS | {"delta_gate", "pyarrow_udf"} ROUTING_CASES = [ + (["spark-local/src/main/scala/org/apache/comet/local/CometLocalRule.scala"], + BUILD_JOBS | {"delta_gate"}), + (["spark-local/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala"], + {"build_linux", "build_linux_full", "build_linux_all_profiles", "build_macos", "build_macos_full"}), # The Maven wrapper and its config feed every job that runs ./mvnw: the # Linux/macOS builds, setup-spark-builder, the Iceberg `mvnw install`, the # Delta gate's effective-pom check and the PyArrow suite's `mvnw install`. diff --git a/dev/ci/compute-changes.py b/dev/ci/compute-changes.py index 4d62020bb83..87b9b4ef730 100644 --- a/dev/ci/compute-changes.py +++ b/dev/ci/compute-changes.py @@ -43,6 +43,7 @@ "native/**", "common/**", "spark/**", + "spark-local/**", "spark-integration/**", "pom.xml", "**/pom.xml", @@ -80,6 +81,7 @@ "native/**", "common/**", "spark/**", + "spark-local/**", "spark-integration/**", "pom.xml", "**/pom.xml", @@ -116,6 +118,7 @@ "native/**", "common/src/main/**", "spark/src/main/**", + "spark-local/src/main/**", "contrib/delta/**", "pom.xml", "**/pom.xml", @@ -197,6 +200,7 @@ "common/src/main/**", "common/pom.xml", "spark/src/main/**", + "spark-local/src/main/**", "!spark/src/main/spark-3.5/**", "!spark/src/main/spark-4.0/**", "!spark/src/main/spark-4.1/**", @@ -225,6 +229,7 @@ "common/src/main/**", "common/pom.xml", "spark/src/main/**", + "spark-local/src/main/**", "!spark/src/main/spark-3.4/**", "!spark/src/main/spark-4.0/**", "!spark/src/main/spark-4.1/**", @@ -253,6 +258,7 @@ "common/src/main/**", "common/pom.xml", "spark/src/main/**", + "spark-local/src/main/**", "!spark/src/main/spark-3.4/**", "!spark/src/main/spark-3.5/**", "!spark/src/main/spark-3.x/**", @@ -281,6 +287,7 @@ "common/src/main/**", "common/pom.xml", "spark/src/main/**", + "spark-local/src/main/**", "!spark/src/main/spark-3.4/**", "!spark/src/main/spark-3.5/**", "!spark/src/main/spark-3.x/**", @@ -309,6 +316,7 @@ "common/src/main/**", "common/pom.xml", "spark/src/main/**", + "spark-local/src/main/**", "!spark/src/main/spark-3.4/**", "!spark/src/main/spark-3.5/**", "!spark/src/main/spark-3.x/**", @@ -342,6 +350,7 @@ "common/src/main/**", "common/pom.xml", "spark/src/main/**", + "spark-local/src/main/**", "!spark/src/main/scala/org/apache/comet/GenerateDocs.scala", "spark/pom.xml", "dev/diffs/iceberg/**", @@ -367,6 +376,7 @@ "common/src/main/**", "common/pom.xml", "spark/src/main/**", + "spark-local/src/main/**", "!spark/src/main/scala/org/apache/comet/GenerateDocs.scala", "spark/pom.xml", "dev/diffs/iceberg/**", @@ -392,6 +402,7 @@ "common/src/main/**", "common/pom.xml", "spark/src/main/**", + "spark-local/src/main/**", "!spark/src/main/scala/org/apache/comet/GenerateDocs.scala", "spark/pom.xml", "dev/diffs/iceberg/**", @@ -417,6 +428,7 @@ "common/src/main/**", "common/pom.xml", "spark/src/main/**", + "spark-local/src/main/**", "!spark/src/main/scala/org/apache/comet/GenerateDocs.scala", "spark/pom.xml", "dev/diffs/iceberg/**", diff --git a/docs/source/contributor-guide/index.md b/docs/source/contributor-guide/index.md index 1abd8891d20..5642e1040e9 100644 --- a/docs/source/contributor-guide/index.md +++ b/docs/source/contributor-guide/index.md @@ -45,6 +45,7 @@ Development Guide :hidden: Comet Plugin Overview +Local Execution Arrow FFI Timezone Handling Floating-Point Semantics diff --git a/docs/source/contributor-guide/local-execution.md b/docs/source/contributor-guide/local-execution.md new file mode 100644 index 00000000000..87d9250fe82 --- /dev/null +++ b/docs/source/contributor-guide/local-execution.md @@ -0,0 +1,119 @@ + + +# Local Execution + +Local execution is an experimental mode for Spark applications that run with a local master. +Instead of running a query as Spark tasks that each create and run a native plan, it runs the +whole query as one DataFusion graph inside a single Spark task, for every partition at once. + +The internal `spark.comet.exec.local.enabled` option, false by default, turns it on. A query +is admitted only as a whole. Every other query uses ordinary Comet execution. + +## What is admitted + +Admission runs in `CometRule`, on the plan that Comet's ordinary conversion produces. A query is +admitted when: + +- the application's `SparkContext` is local (`local[N]`), so the query runs in this process; +- the plan is a batch query and not a subquery being prepared; +- the converted plan is one native block, a tree of Comet native operators whose root carries + the block's serialized plan, with no JVM input, broadcast or shuffle inside it; +- the block's only leaf is a `CometRangeExec`, so the range must be converted to Comet + (`spark.comet.exec.range.enabled`); +- its expressions are deterministic, contain no subquery, and do not call back into the JVM, + which the serialized plan is also checked for; +- it has between 1 and 1024 partitions. + +A root `CollectLimitExec`, which Spark plans for `take`, `head` and `show`, is kept above the +local query. Such plans have no exchange, so AQE does not wrap them and they are admitted +whether or not AQE is enabled. + +## Execution + +The admitted plan becomes `CometLocalResultExec` over `CometLocalExec`, a leaf whose RDD has a +single partition. That partition's task creates the native graph: + +- The memory pool, spill directories and DataFusion session are those `Native.createPlan` creates + for any Comet plan of the task: `create_execution_context` builds them, with Comet's function + registrations and session settings, and the pool acquires memory from Spark through the task's + `CometTaskMemoryManager`. +- The block's operators are created once per partition, by the planner for that partition, as a + Spark task would create them. Partition `i` of the graph is the output of the operators created + for partition `i`, so per-partition state such as a range slice is the same as with ordinary + Comet execution. +- `native/local` runs all partitions concurrently on Comet's Tokio runtime and delivers them in + partition order, each up to two batches ahead. The result therefore has Spark's row order, and + a limit returns the same rows as Spark's `executeTake`. The first error of any partition ends + the query at once. + +`CometLocalExec` reports `UnknownPartitioning(1)`: the result is one partition, but its rows +are not clustered by anything, so a query over a cached local result still gets the exchange +its distribution requires. It reports the block's ordering, which delivery in partition order +preserves. + +Explain shows the native block under `CometLocalExec`. Its native operators do not report SQL +metrics; `CometLocalExec` counts the rows and batches it delivers. + +## Result delivery + +The native side hands batches to the JVM through a queue that holds at most one batch. Exported +Arrow arrays use `prepare_output` and `NativeUtil`; an imported batch stays valid after its +native query is closed, until the JVM closes it. A pull returns "pending" after 50 ms, so that +the task can check whether it was killed. + +`CometLocalResultExec` sits above Spark's columnar-to-row transition. For `collect` and `take` +it still runs the single task as a Spark job, which keeps cancellation, job groups and SQL +metrics, but the task copies its rows into a driver-side slot instead of returning an encoded +task result. Spark would otherwise encode and compress every row in that one task and decode +them on the driver, all on one thread. While copying, the task counts uncompressed `UnsafeRow` +bytes and stops reading, which closes the native query, once they exceed +`spark.driver.maxResultSize`; the driver then fails the action. Spark counts compressed bytes, +so this check is the stricter one. On a non-local master the node uses Spark's ordinary +collect. Other actions, such as `toLocalIterator`, keep Spark's behavior. + +## Lifecycle and cancellation + +The JNI registry holds numeric, never reused IDs for at most 1024 live queries. EOF, errors, +early termination and task completion remove an entry. A concurrent pull holds an `Arc`, so a +close cannot free an entry that is being read, and cancellation does not take the reader lock. +A query records when it has delivered its last batch; if the runtime stops it before that, the +reader reports an error instead of a truncated result. + +Closing a query drops its streams on the calling thread, as `releasePlan` does for an ordinary +plan without JVM input: cancelling a Tokio task only takes effect once a worker is free, and +every worker can be waiting, for instance in Spark's `acquireMemory` for memory those streams +hold. Tasks that DataFusion operators spawn themselves still end asynchronously. + +## Testing + +Build native code before the JVM suite, and never use Maven `-pl`: + +```shell +cd native +cargo test -p datafusion-comet-local +cd .. +./mvnw test -Dsuites=org.apache.comet.local.CometLocalExecutionSuite +``` + +The native tests cover partition-order delivery, early failure, cancellation with busy workers, +the one-batch handoff and runtime shutdown. `CometLocalExecutionSuite` compares admitted queries +with Spark's results and order, and covers ranges Spark would overflow, declined expressions, +AQE, limits, caching, the result size limit, job and task cancellation, and native errors. After +each test it checks that no native query, imported Arrow memory or result slot is left. diff --git a/docs/source/contributor-guide/plugin_overview.md b/docs/source/contributor-guide/plugin_overview.md index 072f59e01a2..ff734a710b2 100644 --- a/docs/source/contributor-guide/plugin_overview.md +++ b/docs/source/contributor-guide/plugin_overview.md @@ -82,6 +82,10 @@ Execution is disabled. (`CometExecRule`). The order matters, because operator conversion builds its native plan up from the nodes that scan conversion produces. Each phase is described below. +With `spark.comet.exec.local.enabled`, the columnar `CometRule` first offers the converted plan to +local execution, which runs a whole qualifying query as one native graph in a single Spark task +when the application uses a local master. See [Local Execution](local-execution.md). + ### Phase 1: CometScanRule `CometScanRule` replaces any Parquet scans with Comet operators. There are different paths for Spark v1 and v2 data sources. diff --git a/native/Cargo.lock b/native/Cargo.lock index e4d2aac7e1d..b9b383f6a7e 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -2001,6 +2001,7 @@ dependencies = [ "datafusion", "datafusion-comet-common", "datafusion-comet-jni-bridge", + "datafusion-comet-local", "datafusion-comet-operators", "datafusion-comet-proto", "datafusion-comet-shuffle", @@ -2080,6 +2081,15 @@ dependencies = [ "thiserror 2.0.20", ] +[[package]] +name = "datafusion-comet-local" +version = "1.2.0" +dependencies = [ + "datafusion", + "futures", + "tokio", +] + [[package]] name = "datafusion-comet-operators" version = "1.2.0" diff --git a/native/Cargo.toml b/native/Cargo.toml index defbf41ae27..f5ea8161085 100644 --- a/native/Cargo.toml +++ b/native/Cargo.toml @@ -16,8 +16,8 @@ # under the License. [workspace] -default-members = ["core", "spark-expr", "common", "proto", "jni-bridge", "shuffle", "operators", "comet-udf-sdk", "comet-test-udfs"] -members = ["core", "spark-expr", "common", "proto", "jni-bridge", "shuffle", "operators", "comet-udf-sdk", "comet-test-udfs"] +default-members = ["core", "spark-expr", "common", "proto", "jni-bridge", "shuffle", "operators", "comet-udf-sdk", "comet-test-udfs", "local"] +members = ["core", "spark-expr", "common", "proto", "jni-bridge", "shuffle", "operators", "comet-udf-sdk", "comet-test-udfs", "local"] # Crates under ../contrib are intentionally NOT workspace members. Core pulls them in as # optional path dependencies when their corresponding contrib feature is enabled. exclude = ["../contrib"] diff --git a/native/core/Cargo.toml b/native/core/Cargo.toml index b21c5bc75a8..bf0baf56afd 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -61,6 +61,7 @@ datafusion-datasource = { workspace = true } datafusion-spark = { workspace = true } once_cell = "1.18.0" datafusion-comet-common = { workspace = true } +datafusion-comet-local = { path = "../local" } datafusion-comet-spark-expr = { workspace = true } datafusion-comet-jni-bridge = { workspace = true } datafusion-comet-proto = { workspace = true } diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 262d8b8fd41..11eaa2bcb1a 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -522,7 +522,7 @@ fn collect_op_names<'a>(op: &'a Operator, names: &mut std::collections::BTreeSet } /// Comet native execution context. Kept alive across JNI calls. -struct ExecutionContext { +pub(crate) struct ExecutionContext { /// The id of the execution context. pub id: i64, /// The deserialized Spark plan @@ -705,8 +705,21 @@ pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( }) } +impl ExecutionContext { + /// The planner that creates this plan's operators for `partition`. + pub(crate) fn planner(&self, partition: i32) -> PhysicalPlanner { + PhysicalPlanner::new(Arc::clone(&self.session_ctx), partition) + .with_exec_id(self.id) + .with_strict_timestamp_labels(self.strict_timestamp_labels) + .with_sql_text_pool(&self.spark_plan) + .with_task_context(self.task_context.clone()) + .with_class_loader(self.class_loader.clone()) + .with_shuffle_partition_pusher(self.shuffle_partition_pusher.clone()) + } +} + /// What `createPlan` takes from the Spark configs. -struct PlanSettings { +pub(crate) struct PlanSettings { spark_config: HashMap, debug_native: bool, strict_timestamp_labels: bool, @@ -718,7 +731,7 @@ struct PlanSettings { /// Deserializes the Spark configs of a plan and initializes the Tokio runtime from them. The /// first step of `Native.createPlan`, before its trace span opens. -fn prepare_plan_settings(serialized_spark_configs: &[u8]) -> CometResult { +pub(crate) fn prepare_plan_settings(serialized_spark_configs: &[u8]) -> CometResult { let spark_configs = serde::deserialize_config(serialized_spark_configs)?; let spark_config: HashMap = spark_configs.entries.into_iter().collect(); @@ -741,33 +754,34 @@ fn prepare_plan_settings(serialized_spark_configs: &[u8]) -> CometResult, - partition_count: usize, - batch_size: usize, - off_heap_mode: bool, - memory_pool_type: String, - memory_limit: i64, - task_attempt_id: i64, - task_cpus: usize, - local_dirs: Vec, +pub(crate) struct PlanInputs { + pub(crate) id: i64, + pub(crate) settings: PlanSettings, + pub(crate) serialized_plan: Vec, + pub(crate) partition_count: usize, + pub(crate) batch_size: usize, + pub(crate) off_heap_mode: bool, + pub(crate) memory_pool_type: String, + pub(crate) memory_limit: i64, + pub(crate) task_attempt_id: i64, + pub(crate) task_cpus: usize, + pub(crate) local_dirs: Vec, /// In milliseconds; zero or less publishes metrics only when the plan is released. - metrics_update_interval: i64, + pub(crate) metrics_update_interval: i64, /// When plan creation started, the start of `ExecutionContext::plan_creation_time`. - start: Instant, - jvm: PlanJvmRefs, + pub(crate) start: Instant, + pub(crate) jvm: PlanJvmRefs, } /// The JVM objects a plan calls back into, as JNI global references. Each is `None` or empty /// only without a JVM (unit tests) or when Spark passes none. -struct PlanJvmRefs { - metrics: Option>>>, - input_sources: Vec>>>, - key_unwrapper: Option>>>, - task_context: Option>>>, - class_loader: Option>>>, +#[derive(Default)] +pub(crate) struct PlanJvmRefs { + pub(crate) metrics: Option>>>, + pub(crate) input_sources: Vec>>>, + pub(crate) key_unwrapper: Option>>>, + pub(crate) task_context: Option>>>, + pub(crate) class_loader: Option>>>, } /// Creates the execution context of a native plan: deserializes the plan, sets up its memory @@ -776,7 +790,7 @@ struct PlanJvmRefs { /// /// `memory_pool_for` builds the pool every reservation of the plan goes through from its /// config; `Native.createPlan` builds one that acquires memory from Spark. -fn create_execution_context( +pub(crate) fn create_execution_context( inputs: PlanInputs, memory_pool_for: impl FnOnce(&MemoryPoolConfig) -> Arc, ) -> CometResult> { @@ -1122,7 +1136,7 @@ fn register_datafusion_spark_function(session_ctx: &SessionContext) { } /// Prepares arrow arrays for output. -fn prepare_output( +pub(crate) fn prepare_output( env: &mut Env, array_addrs: JLongArray, schema_addrs: JLongArray, @@ -1454,20 +1468,12 @@ fn execute_plan( partition: i32, publish_metrics: &mut PublishMetrics, ) -> CometResult> { - let exec_context_id = exec_context.id; - // Initialize the execution stream. // Because we don't know if input arrays are dictionary-encoded when we create // query plan, we need to defer stream initialization to first time execution. if exec_context.root_op.is_none() { let start = Instant::now(); - let planner = PhysicalPlanner::new(Arc::clone(&exec_context.session_ctx), partition) - .with_exec_id(exec_context_id) - .with_strict_timestamp_labels(exec_context.strict_timestamp_labels) - .with_sql_text_pool(&exec_context.spark_plan) - .with_task_context(exec_context.task_context.clone()) - .with_class_loader(exec_context.class_loader.clone()) - .with_shuffle_partition_pusher(exec_context.shuffle_partition_pusher.clone()); + let planner = exec_context.planner(partition); let (scans, shuffle_scans, root_op) = planner.create_plan( &exec_context.spark_plan, &mut exec_context.input_sources.clone(), diff --git a/native/core/src/execution/mod.rs b/native/core/src/execution/mod.rs index 9a2e54783b4..bdea0f9825c 100644 --- a/native/core/src/execution/mod.rs +++ b/native/core/src/execution/mod.rs @@ -26,7 +26,7 @@ pub mod operators; pub(crate) mod planner; pub mod serde; pub use datafusion_comet_shuffle as shuffle; -mod memory_pools; +pub(crate) mod memory_pools; pub(crate) mod sort; pub(crate) mod spark_aggregate_state; pub(crate) mod spark_config; diff --git a/native/core/src/lib.rs b/native/core/src/lib.rs index 5eb56aaec01..2a27591d980 100644 --- a/native/core/src/lib.rs +++ b/native/core/src/lib.rs @@ -60,6 +60,7 @@ pub mod alloc_accounting; pub mod cloud; pub mod comet_native_udf_bridge; pub mod execution; +mod local; pub mod parquet; // this module is for non release only. Intended for debugging/profiling purposes #[cfg(debug_assertions)] diff --git a/native/core/src/local.rs b/native/core/src/local.rs new file mode 100644 index 00000000000..60c836f66c9 --- /dev/null +++ b/native/core/src/local.rs @@ -0,0 +1,249 @@ +// 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. + +//! JNI adapter for local execution, which runs a whole admitted Comet plan as one DataFusion +//! graph inside the result task of a local Spark application. The registry holds only live +//! queries, never plans for reuse. EOF, errors and explicit close remove entries; Java task +//! completion also closes. + +use std::collections::HashMap; +use std::sync::atomic::{AtomicI64, Ordering}; +use std::sync::{Arc, LazyLock}; +use std::time::{Duration, Instant}; + +use datafusion::physical_plan::union::UnionExec; +use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; +use datafusion_comet_jni_bridge::{jni_new_global_ref, JVMClasses}; +use datafusion_comet_local::handle::{QueryHandle, QueryPoll}; +use datafusion_comet_local::LocalQuery; +use jni::objects::{JByteArray, JLongArray, JObject, JObjectArray, JString, ReleaseMode}; +use jni::sys::{jboolean, jint, jlong, JNI_FALSE}; +use jni::{Env, EnvUnowned}; +use parking_lot::Mutex; + +use crate::errors::{try_unwrap_or_throw, CometError, CometResult}; +use crate::execution::jni_api::{ + create_execution_context, get_runtime, prepare_output, prepare_plan_settings, ExecutionContext, + PlanInputs, PlanJvmRefs, +}; +use crate::execution::memory_pools::create_memory_pool; + +struct Entry { + query: QueryHandle, + columns: usize, + // Holds the plan's memory pool and session until the query is closed. Dropped after the + // query, which cancels the graph that reserves from the pool. Only the creating thread + // uses it; the mutex makes the entry shareable with result pulls. + _context: Mutex>, +} + +// JNI carries numeric IDs instead of raw pointers, so close racing with a result pull +// cannot free an object the pull is using. Last reader releases the removed entry. +static QUERIES: LazyLock>>> = LazyLock::new(Mutex::default); +static NEXT_ID: AtomicI64 = AtomicI64::new(1); +const MAX_LIVE_QUERIES: usize = 1024; + +fn close(id: i64) { + let entry = QUERIES.lock().remove(&id); + if let Some(entry) = entry { + entry.query.cancel(); + } +} + +/// Creates and starts a local query over `partitions` partitions of a serialized Comet plan +/// without JVM inputs. Its memory pool and DataFusion session are those `Native.createPlan` +/// creates for a plan of the calling Spark task. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_local_NativeLocal_createQuery( + e: EnvUnowned, + _: JObject, + serialized_query: JByteArray, + serialized_spark_configs: JByteArray, + partitions: jint, + columns: jint, + comet_task_memory_manager_obj: JObject, + local_dirs: JObjectArray, + batch_size: jint, + off_heap_mode: jboolean, + memory_pool_type: JString, + memory_limit: jlong, + task_attempt_id: jlong, + task_cpus: jlong, +) -> jlong { + try_unwrap_or_throw(&e, |env| { + if !(1..=MAX_PARTITIONS).contains(&partitions) || columns < 1 { + return Err(CometError::Internal( + "Invalid local query parameters".into(), + )); + } + let start = Instant::now(); + let settings = prepare_plan_settings(&env.convert_byte_array(serialized_spark_configs)?)?; + JVMClasses::init(env); + let serialized_plan = env.convert_byte_array(serialized_query)?; + let task_memory_manager = + Arc::new(jni_new_global_ref!(env, comet_task_memory_manager_obj)?); + let memory_pool_type = memory_pool_type.try_to_string(env)?; + let mut local_dirs_vec = vec![]; + for i in 0..local_dirs.len(env)? { + let dir = local_dirs.get_element(env, i)?; + let dir = unsafe { JString::from_raw(&*env, dir.into_raw()) }; + local_dirs_vec.push(dir.try_to_string(env)?); + } + let id = next_id()?; + let inputs = PlanInputs { + id, + settings, + serialized_plan, + partition_count: partitions as usize, + batch_size: batch_size as usize, + off_heap_mode: off_heap_mode != JNI_FALSE, + memory_pool_type, + memory_limit, + task_attempt_id, + task_cpus: task_cpus as usize, + local_dirs: local_dirs_vec, + metrics_update_interval: 0, + start, + jvm: PlanJvmRefs::default(), + }; + let context = create_execution_context(inputs, |config| { + create_memory_pool(config, task_memory_manager, task_attempt_id) + })?; + let plan = local_plan(&context, partitions)?; + if plan.schema().fields().len() != columns as usize { + return Err(CometError::Internal( + "Local output schema width mismatch".into(), + )); + } + let query = LocalQuery::new(plan, context.session_ctx.task_ctx()).in_partition_order(); + let query = QueryHandle::start(query, &get_runtime())?; + register( + id, + Entry { + query, + columns: columns as usize, + _context: Mutex::new(context), + }, + )?; + Ok(id) + }) +} + +const MAX_PARTITIONS: jint = 1024; + +/// One graph for the whole query. Each partition's operators are created by the planner for +/// that partition, as a Spark task would create them, so per-partition state such as a range +/// slice or a partition-dependent expression is correct. Partition `i` of the graph is the +/// single output partition of the operators created for partition `i`. +fn local_plan(context: &ExecutionContext, partitions: jint) -> CometResult> { + let mut plans = Vec::with_capacity(partitions as usize); + for partition in 0..partitions { + let (scans, shuffle_scans, root) = context.planner(partition).create_plan( + &context.spark_plan, + &mut vec![], + partitions as usize, + )?; + if !scans.is_empty() + || !shuffle_scans.is_empty() + || root.native_plan.output_partitioning().partition_count() != 1 + { + return Err(CometError::Internal( + "Local query plans must have no JVM input and one output partition".into(), + )); + } + plans.push(Arc::clone(&root.native_plan)); + } + Ok(UnionExec::try_new(plans)?) +} + +fn next_id() -> CometResult { + NEXT_ID + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| id.checked_add(1)) + .map_err(|_| CometError::Internal("Local query IDs exhausted".into())) +} + +fn register(id: i64, entry: Entry) -> CometResult<()> { + let mut registry = QUERIES.lock(); + if registry.len() >= MAX_LIVE_QUERIES { + return Err(CometError::Internal("Too many live local queries".into())); + } + registry.insert(id, Arc::new(entry)); + Ok(()) +} + +fn next(env: &mut Env, id: i64, arrays: JLongArray, schemas: JLongArray) -> CometResult { + let entry = QUERIES + .lock() + .get(&id) + .cloned() + .ok_or_else(|| CometError::Internal("Local query is closed or unknown".into()))?; + if arrays.len(env)? != entry.columns || schemas.len(env)? != entry.columns { + return Err(CometError::Internal( + "Local output column count mismatch".into(), + )); + } + // The caller supplies writable Arrow C structs. Reject null pointers before exporting. + for addresses in [&arrays, &schemas] { + let elements = unsafe { addresses.get_elements(env, ReleaseMode::NoCopyBack)? }; + if elements.contains(&0) { + return Err(CometError::Internal("Null local output pointer".into())); + } + } + match entry.query.poll(Duration::from_millis(50))? { + QueryPoll::Pending => Ok(-2), + QueryPoll::Finished => Ok(-1), + QueryPoll::Batch(batch) => prepare_output(env, arrays, schemas, batch, false), + } +} + +/// -2 means pending, -1 means EOF; nonnegative values are exported row counts. +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_local_NativeLocal_nextBatch( + e: EnvUnowned, + _: JObject, + id: jlong, + arrays: JLongArray, + schemas: JLongArray, +) -> jlong { + try_unwrap_or_throw(&e, |env| { + let result = next(env, id, arrays, schemas); + if result.is_err() || matches!(result, Ok(-1)) { + close(id); + } + result + }) +} + +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_local_NativeLocal_close( + e: EnvUnowned, + _: JObject, + id: jlong, +) { + try_unwrap_or_throw(&e, |_| { + close(id); + Ok(()) + }) +} + +#[no_mangle] +pub extern "system" fn Java_org_apache_comet_local_NativeLocal_activeQueries( + e: EnvUnowned, + _: JObject, +) -> jlong { + try_unwrap_or_throw(&e, |_| Ok(QUERIES.lock().len() as jlong)) +} diff --git a/native/local/Cargo.toml b/native/local/Cargo.toml new file mode 100644 index 00000000000..74c6cc4ab6b --- /dev/null +++ b/native/local/Cargo.toml @@ -0,0 +1,36 @@ +# 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] +name = "datafusion-comet-local" +description = "Query-scoped native execution for Comet local mode" +version = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } +authors = { workspace = true } +license = { workspace = true } +edition = { workspace = true } +rust-version = { workspace = true } +publish = false + +[dependencies] +datafusion = { workspace = true } +futures = { workspace = true } +tokio = { version = "1", features = ["rt", "sync"] } + +[dev-dependencies] +tokio = { version = "1", features = ["macros", "rt-multi-thread", "time"] } diff --git a/native/local/src/handle.rs b/native/local/src/handle.rs new file mode 100644 index 00000000000..b5f1f3f56a6 --- /dev/null +++ b/native/local/src/handle.rs @@ -0,0 +1,226 @@ +// 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. + +//! Bounded handoff to blocking callers, with cancellation independent of the reader lock. + +use std::panic::AssertUnwindSafe; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{mpsc, Arc, Mutex}; +use std::task::Poll; +use std::time::Duration; + +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::{DataFusionError, Result}; +use futures::future::poll_fn; +use futures::{FutureExt, StreamExt}; +use tokio::runtime::Handle; +use tokio::sync::{OwnedSemaphorePermit, Semaphore}; +use tokio::task::JoinHandle; + +use crate::{lock, LocalQuery, LocalQueryStream}; + +struct Delivery { + batch: Result, + _permit: OwnedSemaphorePermit, +} + +pub enum QueryPoll { + Batch(RecordBatch), + Pending, + Finished, +} + +/// Holds at most one queued batch, plus the batch owned by the external consumer. +/// Cancelling or dropping the handle drops the query's stream on the calling thread and +/// unread batches with it. The runtime must remain alive until DataFusion finishes aborting +/// the tasks its operators spawn. +pub struct QueryHandle { + receiver: Mutex>, + /// The query's stream. The producer locks it only while polling it, so `cancel` can take it + /// whenever the producer is waiting. + stream: Arc>>, + producer: JoinHandle<()>, + cancelled: AtomicBool, + completed: Arc, +} + +impl QueryHandle { + pub fn start(query: LocalQuery, runtime: &Handle) -> Result { + let stream = { + let _guard = runtime.enter(); + query.execute()? + }; + let (sender, receiver) = mpsc::channel(); + let completed = Arc::new(AtomicBool::new(false)); + let worker_completed = Arc::clone(&completed); + let stream = Arc::new(Mutex::new(Some(stream))); + let polled = Arc::clone(&stream); + let producer = runtime.spawn(async move { + let capacity = Arc::new(Semaphore::new(1)); + let result = AssertUnwindSafe(async { + loop { + let permit = Arc::clone(&capacity) + .acquire_owned() + .await + .map_err(|e| DataFusionError::Execution(e.to_string()))?; + let next = poll_fn(|cx| match lock(&polled).as_mut() { + Some(stream) => stream.poll_next_unpin(cx).map(Some), + // Cancelled: the stream was taken, so this is not completion. + None => Poll::Ready(None), + }) + .await; + let Some(next) = next else { + return Ok(false); + }; + match next { + Some(batch) => { + let failed = batch.is_err(); + if sender + .send(Delivery { + batch, + _permit: permit, + }) + .is_err() + || failed + { + break; + } + } + None => break, + } + } + Ok::<_, DataFusionError>(true) + }) + .catch_unwind() + .await; + let error = match result { + Ok(Ok(true)) => None, + Ok(Ok(false)) => return, + Ok(Err(e)) => Some(e), + Err(_) => Some(DataFusionError::Execution( + "Local query worker panicked".into(), + )), + }; + if let Some(error) = error { + if let Ok(permit) = capacity.acquire_owned().await { + let _ = sender.send(Delivery { + batch: Err(error), + _permit: permit, + }); + } + } + // Publish before dropping sender. Runtime shutdown/abort skips this store; + // channel disconnection alone must never turn a truncated result into EOF. + worker_completed.store(true, Ordering::Release); + }); + Ok(Self { + receiver: Mutex::new(receiver), + stream, + producer, + cancelled: AtomicBool::new(false), + completed, + }) + } + + /// A bounded wait lets a JVM caller check Spark task interruption between polls. + /// Only one reader may call this at a time. Cancellation never needs this lock. + pub fn poll(&self, wait: Duration) -> Result { + if self.cancelled.load(Ordering::Acquire) { + return Err(DataFusionError::Execution("Local query cancelled".into())); + } + let receiver = self.receiver.try_lock().map_err(|_| { + DataFusionError::Execution("Concurrent local query result readers".into()) + })?; + let delivery = receiver.recv_timeout(wait); + if self.cancelled.load(Ordering::Acquire) { + return Err(DataFusionError::Execution("Local query cancelled".into())); + } + match delivery { + Ok(delivery) => delivery.batch.map(QueryPoll::Batch), + Err(mpsc::RecvTimeoutError::Timeout) => Ok(QueryPoll::Pending), + Err(mpsc::RecvTimeoutError::Disconnected) => { + if self.completed.load(Ordering::Acquire) { + Ok(QueryPoll::Finished) + } else { + Err(DataFusionError::Execution( + "Local query worker stopped without completing".into(), + )) + } + } + } + } + + /// Stops the query and drops its stream on the calling thread, rather than leaving that to + /// the cancelled producer: a cancellation only runs once a Tokio worker is free, and every + /// worker can be waiting, for instance in Spark's `acquireMemory` for memory this stream + /// holds. Taking the stream waits only for a poll the producer is already in. + pub fn cancel(&self) { + self.cancelled.store(true, Ordering::Release); + self.producer.abort(); + let stream = lock(&self.stream).take(); + // A panic while dropping must not unwind into the JNI caller. + let _ = std::panic::catch_unwind(AssertUnwindSafe(|| drop(stream))); + } +} + +impl Drop for QueryHandle { + fn drop(&mut self) { + self.cancel(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::arrow::datatypes::Schema; + use datafusion::execution::TaskContext; + use datafusion::physical_plan::test::exec::BlockingExec; + + #[test] + fn cancel_interrupts_a_blocked_reader_without_its_mutex() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .build() + .unwrap(); + let plan = Arc::new(BlockingExec::new(Arc::new(Schema::empty()), 2)); + let query = Arc::new( + QueryHandle::start( + LocalQuery::new(plan, Arc::new(TaskContext::default())), + runtime.handle(), + ) + .unwrap(), + ); + let reader = Arc::clone(&query); + let (started_tx, started_rx) = mpsc::channel(); + let (sender, receiver) = mpsc::channel(); + let thread = std::thread::spawn(move || { + // Hold exactly the receiver lock used by poll. A handshake avoids racing + // poll's try_lock with a test thread inspecting that same lock. + let incoming = reader.receiver.lock().unwrap(); + started_tx.send(()).unwrap(); + let result = incoming.recv_timeout(Duration::from_secs(60)); + sender + .send(matches!(result, Err(mpsc::RecvTimeoutError::Disconnected))) + .unwrap(); + }); + started_rx.recv_timeout(Duration::from_secs(5)).unwrap(); + assert!(query.poll(Duration::ZERO).is_err()); + query.cancel(); + assert!(receiver.recv_timeout(Duration::from_secs(5)).unwrap()); + thread.join().unwrap(); + } +} diff --git a/native/local/src/lib.rs b/native/local/src/lib.rs new file mode 100644 index 00000000000..360a331ccb0 --- /dev/null +++ b/native/local/src/lib.rs @@ -0,0 +1,297 @@ +// 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. + +//! Query-scoped execution, independent of Spark tasks and JNI. +//! +//! A caller supplies a fresh, fully planned graph for each execution and receives one +//! result stream. The caller owns the Tokio runtime and must keep it running until +//! cleanup finishes. Spark admission and JNI live outside this crate. + +pub mod handle; + +use std::collections::VecDeque; +use std::panic::AssertUnwindSafe; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex, MutexGuard}; +use std::task::{Context, Poll}; + +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::{DataFusionError, Result}; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + execute_stream, ExecutionPlan, ExecutionPlanProperties, RecordBatchStream, + SendableRecordBatchStream, +}; +use futures::future::poll_fn; +use futures::task::AtomicWaker; +use futures::{FutureExt, Stream, StreamExt}; +use tokio::task::JoinHandle; + +/// Batches each root partition may run ahead of the one being delivered, when results are +/// delivered in partition order. +const PARTITION_BUFFER: usize = 2; + +/// One execution's graph and resources. Consumed when execution starts; not a plan cache. +pub struct LocalQuery { + plan: Arc, + context: Arc, + in_partition_order: bool, +} + +impl LocalQuery { + /// Takes ownership of a fresh graph and its query-scoped context. + /// + /// The planner must satisfy operator distribution and ordering requirements. It must not + /// reuse stateful operator instances from another query or attempt. + pub fn new(plan: Arc, context: Arc) -> Self { + Self { + plan, + context, + in_partition_order: false, + } + } + + /// Delivers the root partitions one after another, in partition order, as Spark returns + /// the partitions of a result. All partitions still run concurrently, each up to a few + /// batches ahead. Only for roots whose partitions do not share an exchange: a partition + /// that waits to be delivered would otherwise leave its exchange buffering for the others. + pub fn in_partition_order(mut self) -> Self { + self.in_partition_order = true; + self + } + + /// Starts all root partitions and exposes one streaming result. + /// + /// By default the partitions are coalesced as they produce batches, which does not + /// preserve their order. A single-partition root preserves its ordering either way. No + /// result-wide collection is performed here. + pub fn execute(self) -> Result { + tokio::runtime::Handle::try_current().map_err(|_| { + DataFusionError::Execution("LocalQuery requires an active Tokio runtime".into()) + })?; + let schema = self.plan.schema(); + let stream = if self.in_partition_order { + in_partition_order(&self.plan, &self.context)? + } else { + execute_stream(Arc::clone(&self.plan), Arc::clone(&self.context))? + }; + Ok(LocalQueryStream { + stream: Some(stream), + query: Some(self), + schema, + }) + } +} + +fn in_partition_order( + plan: &Arc, + context: &Arc, +) -> Result { + let failure = Arc::new(Failure::default()); + let runtime = tokio::runtime::Handle::current(); + let partitions = (0..plan.output_partitioning().partition_count()) + .map(|partition| { + let stream = plan.execute(partition, Arc::clone(context))?; + Ok(Partition::spawn(&runtime, stream, Arc::clone(&failure))) + }) + .collect::>>()?; + Ok(Box::pin(RecordBatchStreamAdapter::new( + plan.schema(), + PartitionOrder { + partitions, + failure, + }, + ))) +} + +/// One root partition running on its own Tokio task, a few batches ahead of delivery. +struct Partition { + batches: tokio::sync::mpsc::Receiver, + /// Set once the task has sent the stream's last batch. The channel closes whenever the + /// task ends, and a runtime that shuts down cancels its tasks, so only this tells that + /// every batch arrived. + ended: Arc, + /// The partition's stream. The task locks it only while polling it, so dropping the + /// partition can take it whenever the task is waiting. + stream: Arc>>, + task: JoinHandle<()>, +} + +impl Partition { + fn spawn( + runtime: &tokio::runtime::Handle, + stream: SendableRecordBatchStream, + failure: Arc, + ) -> Self { + let (sender, batches) = tokio::sync::mpsc::channel(PARTITION_BUFFER); + let ended = Arc::new(AtomicBool::new(false)); + let stream = Arc::new(Mutex::new(Some(stream))); + let (task_ended, polled) = (Arc::clone(&ended), Arc::clone(&stream)); + let task = runtime.spawn(async move { + let result = AssertUnwindSafe(async { + // Reserve before polling, so at most `PARTITION_BUFFER` batches are held. + while let Ok(permit) = sender.reserve().await { + let next = poll_fn(|cx| match lock(&polled).as_mut() { + Some(stream) => stream.poll_next_unpin(cx), + None => Poll::Ready(None), + }) + .await; + match next { + Some(Ok(batch)) => permit.send(batch), + Some(Err(error)) => return failure.record(error), + None => { + // A taken stream also ends here, but then nobody reads the flag. + task_ended.store(true, Ordering::Release); + return; + } + } + } + }) + .catch_unwind() + .await; + if result.is_err() { + failure.record(DataFusionError::Execution( + "Local query partition panicked".into(), + )); + } + }); + Self { + batches, + ended, + stream, + task, + } + } +} + +impl Drop for Partition { + /// Drops the partition's stream on the calling thread, since cancelling the task only takes + /// effect once a Tokio worker is free. See `QueryHandle::cancel`. + fn drop(&mut self) { + self.task.abort(); + let stream = lock(&self.stream).take(); + let _ = std::panic::catch_unwind(AssertUnwindSafe(|| drop(stream))); + } +} + +/// The first error of any partition, reported as soon as it happens rather than when the +/// failed partition's turn comes. +#[derive(Default)] +struct Failure { + error: Mutex>, + waker: AtomicWaker, +} + +impl Failure { + /// Keeps the first error and wakes the consumer. + fn record(&self, error: DataFusionError) { + lock(&self.error).get_or_insert(error); + self.waker.wake(); + } + + fn take(&self) -> Option { + lock(&self.error).take() + } +} + +struct PartitionOrder { + partitions: VecDeque, + failure: Arc, +} + +impl Stream for PartitionOrder { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + self.failure.waker.register(cx.waker()); + loop { + if let Some(error) = self.failure.take() { + self.partitions.clear(); + return Poll::Ready(Some(Err(error))); + } + let Some(partition) = self.partitions.front_mut() else { + return Poll::Ready(None); + }; + match partition.batches.poll_recv(cx) { + Poll::Ready(Some(batch)) => return Poll::Ready(Some(Ok(batch))), + Poll::Ready(None) if partition.ended.load(Ordering::Acquire) => { + self.partitions.pop_front(); + } + Poll::Ready(None) => { + // A failed partition records its error before its channel closes. + let error = self.failure.take().unwrap_or_else(|| { + DataFusionError::Execution( + "Local query partition stopped without completing".into(), + ) + }); + self.partitions.clear(); + return Poll::Ready(Some(Err(error))); + } + Poll::Pending => return Poll::Pending, + } + } + } +} + +/// Locks `mutex`, ignoring poisoning: the guarded values stay consistent across a panic. +pub(crate) fn lock(mutex: &Mutex) -> MutexGuard<'_, T> { + mutex.lock().unwrap_or_else(|e| e.into_inner()) +} + +/// Owns the graph until EOF, the first error, cancellation, or drop. +/// +/// Cancellation requests teardown by dropping DataFusion streams. Spawned tasks finish +/// aborting asynchronously on the caller's runtime; this is not a join barrier. +pub struct LocalQueryStream { + // Drop streams before releasing our graph/context references. + stream: Option, + query: Option, + schema: SchemaRef, +} + +impl LocalQueryStream { + /// Stops result delivery and requests cancellation of the whole native execution. + /// Calling this again, or polling afterwards, is harmless. + pub fn cancel(&mut self) { + self.stream = None; + self.query = None; + } +} + +impl Stream for LocalQueryStream { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let Some(stream) = self.stream.as_mut() else { + return Poll::Ready(None); + }; + let result = stream.as_mut().poll_next(cx); + if matches!(result, Poll::Ready(None | Some(Err(_)))) { + self.cancel(); + } + result + } +} + +impl RecordBatchStream for LocalQueryStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} diff --git a/native/local/tests/bridge.rs b/native/local/tests/bridge.rs new file mode 100644 index 00000000000..7d953aff8b9 --- /dev/null +++ b/native/local/tests/bridge.rs @@ -0,0 +1,303 @@ +// 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. + +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use datafusion::arrow::array::Int64Array; +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::test::exec::{BlockingExec, ErrorExec, MockExec, PanicExec}; +use datafusion::physical_plan::union::UnionExec; +use datafusion::physical_plan::ExecutionPlan; +use datafusion_comet_local::handle::{QueryHandle, QueryPoll}; +use datafusion_comet_local::LocalQuery; +use tokio::runtime::Builder; + +fn runtime() -> tokio::runtime::Runtime { + Builder::new_multi_thread() + .worker_threads(1) + .enable_all() + .build() + .unwrap() +} + +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])) +} + +/// One partition per `(start, end)` range, each split into batches of `rows` values. +fn ranges(ranges: &[(i64, i64)], rows: i64) -> Arc { + let partitions = ranges + .iter() + .map(|&(first, last)| { + let batches = (first..last) + .step_by(rows as usize) + .map(|start| { + let end = (start + rows).min(last); + Ok(RecordBatch::try_new( + schema(), + vec![Arc::new(Int64Array::from_iter_values(start..end))], + ) + .unwrap()) + }) + .collect(); + Arc::new(MockExec::new(batches, schema())) as Arc + }) + .collect(); + UnionExec::try_new(partitions).unwrap() +} + +fn values(batch: &RecordBatch) -> &[i64] { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() +} + +#[test] +fn partition_order_is_preserved_through_the_handoff() { + let runtime = runtime(); + let plan = ranges(&[(0, 10), (10, 10), (10, 25), (25, 26)], 3); + let refs = Arc::downgrade(&plan); + let query = QueryHandle::start( + LocalQuery::new(plan, Arc::new(TaskContext::default())).in_partition_order(), + runtime.handle(), + ) + .unwrap(); + let mut actual = vec![]; + let deadline = Instant::now() + Duration::from_secs(5); + loop { + assert!(Instant::now() < deadline); + match query.poll(Duration::from_millis(20)).unwrap() { + QueryPoll::Batch(batch) => actual.extend_from_slice(values(&batch)), + QueryPoll::Pending => (), + QueryPoll::Finished => break, + } + } + assert_eq!(actual, (0..26).collect::>()); + drop(query); + while refs.strong_count() > 0 { + assert!(Instant::now() < deadline); + std::thread::yield_now(); + } +} + +#[test] +fn pending_is_distinct_from_eof_and_drop_releases_graph() { + let runtime = runtime(); + let plan = BlockingExec::new(schema(), 4); + let refs = plan.refs(); + let query = QueryHandle::start( + LocalQuery::new(Arc::new(plan), Arc::new(TaskContext::default())), + runtime.handle(), + ) + .unwrap(); + assert!(matches!( + query.poll(Duration::from_millis(20)).unwrap(), + QueryPoll::Pending + )); + assert!( + refs.strong_count() > 1, + "pending streams should have started" + ); + drop(query); + let deadline = Instant::now() + Duration::from_secs(5); + while refs.strong_count() > 0 { + assert!(Instant::now() < deadline); + std::thread::yield_now(); + } +} + +#[test] +fn worker_panic_is_an_error_not_successful_eof() { + let runtime = runtime(); + let plan = Arc::new(PanicExec::new(schema(), 2)); + let query = QueryHandle::start( + LocalQuery::new(plan, Arc::new(TaskContext::default())), + runtime.handle(), + ) + .unwrap(); + let deadline = Instant::now() + Duration::from_secs(5); + loop { + assert!(Instant::now() < deadline); + match query.poll(Duration::from_millis(20)) { + Err(error) => { + assert!(error.to_string().contains("panic")); + break; + } + Ok(QueryPoll::Pending | QueryPoll::Batch(_)) => (), + Ok(QueryPoll::Finished) => panic!("worker failure was lost"), + } + } +} + +#[test] +fn startup_failure_is_returned() { + let runtime = runtime(); + let plan = Arc::new(ErrorExec::new()); + assert!(QueryHandle::start( + LocalQuery::new(plan, Arc::new(TaskContext::default())), + runtime.handle() + ) + .is_err()); +} + +#[test] +fn runtime_shutdown_is_an_error_not_successful_eof() { + let runtime = runtime(); + let query = QueryHandle::start( + LocalQuery::new( + Arc::new(BlockingExec::new(schema(), 4)), + Arc::new(TaskContext::default()), + ), + runtime.handle(), + ) + .unwrap(); + assert!(matches!( + query.poll(Duration::from_millis(20)).unwrap(), + QueryPoll::Pending + )); + drop(runtime); + match query.poll(Duration::from_millis(50)) { + Err(error) => assert!(error.to_string().contains("without completing")), + Ok(_) => panic!("runtime shutdown must not look like successful EOF"), + } +} + +#[test] +fn backpressured_query_does_not_starve_other_queries_and_cancel_releases_graph() { + let runtime = runtime(); + for _ in 0..12 { + let stalled = ranges(&[(0, 100_000), (100_000, 200_000)], 5); + let refs = Arc::downgrade(&stalled); + let stalled = QueryHandle::start( + LocalQuery::new(stalled, Arc::new(TaskContext::default())).in_partition_order(), + runtime.handle(), + ) + .unwrap(); + // Leave its one-batch result channel unread while a second query makes progress. + let healthy = QueryHandle::start( + LocalQuery::new( + ranges(&[(100, 150), (150, 200), (200, 237)], 3), + Arc::new(TaskContext::default()), + ) + .in_partition_order(), + runtime.handle(), + ) + .unwrap(); + let deadline = Instant::now() + Duration::from_secs(5); + let mut next = 100; + loop { + assert!( + Instant::now() < deadline, + "a stalled consumer starved another query" + ); + match healthy.poll(Duration::from_millis(20)).unwrap() { + QueryPoll::Batch(batch) => { + for value in values(&batch) { + assert_eq!(*value, next); + next += 1; + } + } + QueryPoll::Pending => (), + QueryPoll::Finished => break, + } + } + assert_eq!(next, 237); + stalled.cancel(); + assert!(stalled.poll(Duration::ZERO).is_err()); + drop(stalled); + drop(healthy); + while refs.strong_count() != 0 { + assert!(Instant::now() < deadline, "cancelled graph remained live"); + std::thread::yield_now(); + } + } +} + +#[test] +fn shutdown_after_partial_output_never_silently_truncates_the_result() { + let runtime = runtime(); + let query = QueryHandle::start( + LocalQuery::new(ranges(&[(0, 100_000)], 5), Arc::new(TaskContext::default())), + runtime.handle(), + ) + .unwrap(); + let deadline = Instant::now() + Duration::from_secs(5); + let first = loop { + assert!(Instant::now() < deadline); + match query.poll(Duration::from_millis(20)).unwrap() { + QueryPoll::Batch(batch) => break batch, + QueryPoll::Pending => (), + QueryPoll::Finished => panic!("large fixture ended"), + } + }; + drop(runtime); + assert_eq!(values(&first)[0], 0); + let mut queued = 0; + loop { + assert!(Instant::now() < deadline); + match query.poll(Duration::from_millis(20)) { + Ok(QueryPoll::Batch(_)) => { + queued += 1; + assert!(queued <= 1, "handoff exceeded its one-batch bound"); + } + Ok(QueryPoll::Pending) => (), + Ok(QueryPoll::Finished) => panic!("partial output was reported as complete"), + Err(error) => { + assert!(error.to_string().contains("without completing")); + break; + } + } + } +} + +#[test] +fn cancel_releases_the_graph_while_every_worker_is_busy() { + let runtime = runtime(); + let plan = BlockingExec::new(schema(), 3); + let refs = plan.refs(); + let query = QueryHandle::start( + LocalQuery::new(Arc::new(plan), Arc::new(TaskContext::default())).in_partition_order(), + runtime.handle(), + ) + .unwrap(); + assert!(matches!( + query.poll(Duration::from_millis(20)).unwrap(), + QueryPoll::Pending + )); + // Occupy the runtime's only worker, as Spark's acquireMemory can, so that no cancelled + // task gets to run. + let (release, blocked) = std::sync::mpsc::channel::<()>(); + let (started, running) = std::sync::mpsc::channel::<()>(); + runtime.spawn(async move { + started.send(()).unwrap(); + let _ = blocked.recv(); + }); + running.recv_timeout(Duration::from_secs(5)).unwrap(); + query.cancel(); + assert_eq!( + refs.strong_count(), + 0, + "cancel must drop the graph on the calling thread" + ); + release.send(()).unwrap(); +} diff --git a/native/local/tests/execution.rs b/native/local/tests/execution.rs new file mode 100644 index 00000000000..2fc13831157 --- /dev/null +++ b/native/local/tests/execution.rs @@ -0,0 +1,525 @@ +// 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. + +use std::collections::{BTreeMap, BTreeSet}; +use std::fmt; +use std::sync::{Arc, Mutex, Weak}; +use std::time::Duration; + +use datafusion::arrow::array::Int64Array; +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{DataFusionError, Result}; +use datafusion::execution::config::SessionConfig; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::{expressions::Column, PhysicalExpr}; +use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion::physical_plan::empty::EmptyExec; +use datafusion::physical_plan::limit::GlobalLimitExec; +use datafusion::physical_plan::repartition::RepartitionExec; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::test::exec::{BlockingExec, ErrorExec, MockExec}; +use datafusion::physical_plan::union::UnionExec; +use datafusion::physical_plan::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, + PlanProperties, RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, +}; +use datafusion_comet_local::LocalQuery; +use futures::{StreamExt, TryStreamExt}; +use tokio::time::timeout; + +const DEADLINE: Duration = Duration::from_secs(10); +const INPUTS: usize = 4; +const ROWS_PER_INPUT: usize = 512; + +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int64, false), + Field::new("id", DataType::Int64, false), + ])) +} + +fn context() -> Arc { + Arc::new(TaskContext::default().with_session_config(SessionConfig::new().with_batch_size(4))) +} + +fn batch(ids: impl Iterator) -> RecordBatch { + let ids: Vec<_> = ids.collect(); + RecordBatch::try_new( + schema(), + vec![ + Arc::new(Int64Array::from_iter_values(ids.iter().map(|id| id % 17))), + Arc::new(Int64Array::from(ids)), + ], + ) + .unwrap() +} + +fn input() -> Arc { + let inputs = (0..INPUTS) + .map(|partition| { + let batches = (0..ROWS_PER_INPUT) + .step_by(4) + .map(|offset| { + let start = (partition * ROWS_PER_INPUT + offset) as i64; + Ok(batch(start..start + 4)) + }) + .collect(); + Arc::new(MockExec::new(batches, schema())) as Arc + }) + .collect(); + UnionExec::try_new(inputs).unwrap() +} + +fn repartition( + input: Arc, + partitioning: Partitioning, +) -> Arc { + Arc::new(RepartitionExec::try_new(input, partitioning).unwrap()) +} + +#[derive(Debug, Default)] +struct Observations { + executions: Mutex>, + key_partitions: Mutex>>, +} + +/// Records actual partition execution and routing without changing the data or scheduling. +#[derive(Debug)] +struct ObserveExec { + input: Arc, + observations: Arc, +} + +impl DisplayAs for ObserveExec { + fn fmt_as(&self, _: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "ObserveExec") + } +} + +impl ExecutionPlan for ObserveExec { + fn name(&self) -> &'static str { + "ObserveExec" + } + + fn properties(&self) -> &Arc { + self.input.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + + fn apply_expressions( + &self, + _: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc, + children: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + assert_eq!(children.len(), 1); + Ok(Arc::new(Self { + input: Arc::clone(&children[0]), + observations: Arc::clone(&self.observations), + })) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + partition: usize, + ctx: Arc, + ) -> Result { + self.observations.executions.lock().unwrap().push(partition); + let observations = Arc::clone(&self.observations); + let stream = self.input.execute(partition, ctx)?.map_ok(move |batch| { + let keys = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + let mut locations = observations.key_partitions.lock().unwrap(); + for key in keys.values() { + locations.entry(*key).or_default().insert(partition); + } + batch + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream, + ))) + } +} + +fn observe(input: Arc) -> (Arc, Arc) { + let observations = Arc::new(Observations::default()); + ( + Arc::new(ObserveExec { + input, + observations: Arc::clone(&observations), + }), + observations, + ) +} + +fn assert_executed_once(observations: &Observations, partitions: usize) { + let mut executed = observations.executions.lock().unwrap().clone(); + executed.sort_unstable(); + assert_eq!(executed, (0..partitions).collect::>()); +} + +async fn wait_released(refs: &Weak) { + timeout(DEADLINE, async { + while refs.strong_count() != 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("execution retained a graph, context, or input stream after teardown"); +} + +async fn wait_started(observations: &Observations, partitions: usize) { + timeout(DEADLINE, async { + while observations.executions.lock().unwrap().len() < partitions { + tokio::task::yield_now().await; + } + }) + .await + .expect("input partitions did not start"); + assert_executed_once(observations, partitions); +} + +async fn check_exchange(hash: bool, nested: bool) { + let (input, source_observations) = observe(input()); + let partitioning = if hash { + Partitioning::Hash(vec![Arc::new(Column::new("key", 0))], 7) + } else { + Partitioning::RoundRobinBatch(7) + }; + let (mut plan, exchange_observations) = observe(repartition(input, partitioning)); + if nested { + plan = repartition(plan, Partitioning::RoundRobinBatch(11)); + } + let plan_ref = Arc::downgrade(&plan); + let ctx = context(); + let ctx_ref = Arc::downgrade(&ctx); + let pool = Arc::clone(ctx.memory_pool()); + let mut stream = LocalQuery::new(plan, ctx).execute().unwrap(); + let batches: Vec = timeout(DEADLINE, stream.by_ref().try_collect()) + .await + .expect("repartition stalled") + .unwrap(); + let mut ids: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + ids.sort_unstable(); + assert_eq!( + ids, + (0..(INPUTS * ROWS_PER_INPUT) as i64).collect::>() + ); + assert_executed_once(&source_observations, INPUTS); + assert_executed_once(&exchange_observations, 7); + if hash { + let locations = exchange_observations.key_partitions.lock().unwrap(); + assert_eq!(locations.len(), 17); + assert!(locations.values().all(|partitions| partitions.len() == 1)); + } + // EOF releases resources even if the JVM would retain its exhausted iterator. + assert!(stream.next().await.is_none()); + assert_eq!(stream.schema(), schema()); + wait_released(&plan_ref).await; + wait_released(&ctx_ref).await; + assert_eq!(pool.reserved(), 0); +} + +#[tokio::test(flavor = "current_thread")] +async fn round_robin_more_partitions_than_workers() { + check_exchange(false, false).await; +} + +#[tokio::test(flavor = "current_thread")] +async fn hash_routes_all_input_partitions_through_one_exchange() { + check_exchange(true, false).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn nested_exchanges_more_partitions_than_workers() { + check_exchange(true, true).await; +} + +#[tokio::test] +async fn single_partition_preserves_order() { + let expected = batch(0..12); + let plan = Arc::new(MockExec::new(vec![Ok(expected.clone())], schema())); + let batches: Vec = LocalQuery::new(plan, context()) + .execute() + .unwrap() + .try_collect() + .await + .unwrap(); + assert_eq!(batches, vec![expected]); +} + +#[tokio::test] +async fn empty_result_finishes_with_schema() { + for partitions in [0, 1, 7] { + let plan = Arc::new(EmptyExec::new(schema()).with_partitions(partitions)); + let mut stream = LocalQuery::new(plan, context()).execute().unwrap(); + assert!(timeout(DEADLINE, stream.next()).await.unwrap().is_none()); + assert_eq!(stream.schema(), schema()); + } +} + +#[tokio::test] +async fn cancel_pending_exchange_releases_inputs() { + let input = BlockingExec::new(schema(), 4); + let refs = input.refs(); + let (input, observations) = observe(Arc::new(input)); + let plan = repartition(input, Partitioning::RoundRobinBatch(7)); + let mut stream = LocalQuery::new(plan, context()).execute().unwrap(); + // Prove cancellation tears down running streams, not just an unstarted graph. + wait_started(&observations, 4).await; + // Drive the execution to Pending, then cancel without waiting for a batch. + assert!(futures::poll!(stream.next()).is_pending()); + stream.cancel(); + stream.cancel(); + assert!(stream.next().await.is_none()); + wait_released(&refs).await; +} + +fn partially_blocked_input(fail: bool) -> (Arc, Weak<()>, Arc) { + let blocking = BlockingExec::new(schema(), 3); + let refs = blocking.refs(); + let (blocking, observations) = observe(Arc::new(blocking)); + let value = if fail { + Err(DataFusionError::Execution("injected source failure".into())) + } else { + Ok(batch(0..4)) + }; + let input = UnionExec::try_new(vec![ + Arc::new(MockExec::new(vec![value], schema()).with_unknown_statistics()), + blocking, + ]) + .unwrap(); + ( + repartition(input, Partitioning::RoundRobinBatch(7)), + refs, + observations, + ) +} + +#[tokio::test] +async fn drop_after_first_batch_aborts_sibling_partitions() { + let (plan, refs, observations) = partially_blocked_input(false); + let mut stream = LocalQuery::new(plan, context()).execute().unwrap(); + wait_started(&observations, 3).await; + timeout(DEADLINE, stream.next()) + .await + .unwrap() + .unwrap() + .unwrap(); + drop(stream); + wait_released(&refs).await; +} + +#[tokio::test] +async fn first_error_cancels_siblings_and_is_terminal() { + let (plan, refs, observations) = partially_blocked_input(true); + let ctx = context(); + let ctx_ref = Arc::downgrade(&ctx); + let pool = Arc::clone(ctx.memory_pool()); + let mut stream = LocalQuery::new(plan, ctx).execute().unwrap(); + wait_started(&observations, 3).await; + let error = timeout(DEADLINE, stream.next()) + .await + .unwrap() + .unwrap() + .unwrap_err(); + assert!(error.to_string().contains("injected source failure")); + assert!(stream.next().await.is_none()); + wait_released(&refs).await; + wait_released(&ctx_ref).await; + assert_eq!(pool.reserved(), 0); +} + +#[tokio::test] +async fn global_limit_stops_exchange_without_draining_inputs() { + let (input, refs, observations) = partially_blocked_input(false); + let plan = Arc::new(GlobalLimitExec::new( + Arc::new(CoalescePartitionsExec::new(input)), + 0, + Some(1), + )); + let mut stream = LocalQuery::new(plan, context()).execute().unwrap(); + wait_started(&observations, 3).await; + let batches: Vec = timeout(DEADLINE, stream.by_ref().try_collect()) + .await + .unwrap() + .unwrap(); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); + wait_released(&refs).await; +} + +#[tokio::test] +async fn cancelling_one_query_does_not_cancel_another() { + let blocking_input = BlockingExec::new(schema(), 4); + let refs = blocking_input.refs(); + let (blocking_input, observations) = observe(Arc::new(blocking_input)); + let mut blocked = LocalQuery::new( + repartition(blocking_input, Partitioning::RoundRobinBatch(7)), + context(), + ) + .execute() + .unwrap(); + wait_started(&observations, 4).await; + let other = LocalQuery::new(input(), context()).execute().unwrap(); + blocked.cancel(); + let batches: Vec = timeout(DEADLINE, other.try_collect()) + .await + .unwrap() + .unwrap(); + assert_eq!( + batches.iter().map(RecordBatch::num_rows).sum::(), + INPUTS * ROWS_PER_INPUT + ); + wait_released(&refs).await; +} + +#[tokio::test] +async fn startup_error_releases_query() { + let plan = Arc::new(ErrorExec::new()); + let refs = Arc::downgrade(&plan); + let error = LocalQuery::new(plan, context()).execute().err().unwrap(); + assert!(error.to_string().contains("ErrorExec")); + wait_released(&refs).await; +} + +#[test] +fn missing_runtime_is_an_error_not_a_panic() { + let plan = Arc::new(EmptyExec::new(schema())); + let refs = Arc::downgrade(&plan); + let error = LocalQuery::new(plan, context()).execute().err().unwrap(); + assert!(error.to_string().contains("active Tokio runtime")); + assert_eq!(refs.strong_count(), 0); +} + +fn ordered_input(partitions: usize, batches: usize) -> Arc { + let plans = (0..partitions) + .map(|partition| { + let base = (partition * batches * 4) as i64; + let values = (0..batches) + .map(|i| Ok(batch(base + i as i64 * 4..base + i as i64 * 4 + 4))) + .collect(); + Arc::new(MockExec::new(values, schema())) as Arc + }) + .collect(); + UnionExec::try_new(plans).unwrap() +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn partition_order_delivers_partitions_in_turn() { + let plan = ordered_input(7, 9); + let batches: Vec = timeout( + DEADLINE, + LocalQuery::new(plan, context()) + .in_partition_order() + .execute() + .unwrap() + .try_collect(), + ) + .await + .unwrap() + .unwrap(); + let values: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + assert_eq!(values, (0..7 * 9 * 4).collect::>()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn partition_order_reports_a_later_failure_without_waiting_for_earlier_partitions() { + let blocking = BlockingExec::new(schema(), 1); + let refs = blocking.refs(); + let failing = MockExec::new( + vec![Err(DataFusionError::Execution( + "injected late failure".into(), + ))], + schema(), + ); + let plan = UnionExec::try_new(vec![Arc::new(blocking), Arc::new(failing)]).unwrap(); + let mut stream = LocalQuery::new(plan, context()) + .in_partition_order() + .execute() + .unwrap(); + let error = timeout(DEADLINE, stream.next()) + .await + .unwrap() + .unwrap() + .unwrap_err(); + assert!(error.to_string().contains("injected late failure")); + assert!(stream.next().await.is_none()); + wait_released(&refs).await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn partition_order_cancel_releases_running_partitions() { + let blocking = BlockingExec::new(schema(), 3); + let refs = blocking.refs(); + let (blocking, observations) = observe(Arc::new(blocking)); + let mut stream = LocalQuery::new(blocking, context()) + .in_partition_order() + .execute() + .unwrap(); + wait_started(&observations, 3).await; + assert!(futures::poll!(stream.next()).is_pending()); + stream.cancel(); + assert!(stream.next().await.is_none()); + wait_released(&refs).await; +} diff --git a/pom.xml b/pom.xml index 084c086e447..6cad4cb4edb 100644 --- a/pom.xml +++ b/pom.xml @@ -938,6 +938,7 @@ under the License. -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\.local\.NativeLocal\..*:s -Wconf:cat=unused-params&src=.*/src/[a-z]+/spark-[^/]+/.*:s -Wconf:cat=unused-params&site=org\.apache\.comet\.rules\.CometScanContrib\.tryTransformV[12]\..*:s @@ -983,6 +984,7 @@ under the License. -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\.local\.NativeLocal\..*:s -Wconf:cat=unused-params&src=.*/src/[a-z]+/spark-[^/]+/.*:s -Wconf:cat=unused-params&site=org\.apache\.comet\.rules\.CometScanContrib\.tryTransformV[12]\..*:s @@ -1199,6 +1201,7 @@ under the License. src/main/scala/**/*.scala src/test/scala/**/*.scala + spark-local/src/**/*.scala src/test/celeborn-reflection-compatibility/**/*.scala src/main/spark-*/**/*.scala diff --git a/spark-local/src/main/scala/org/apache/comet/local/CometLocalExec.scala b/spark-local/src/main/scala/org/apache/comet/local/CometLocalExec.scala new file mode 100644 index 00000000000..726a4827f9e --- /dev/null +++ b/spark-local/src/main/scala/org/apache/comet/local/CometLocalExec.scala @@ -0,0 +1,88 @@ +/* + * 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.local + +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, SortOrder} +import org.apache.spark.sql.catalyst.plans.QueryPlan +import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, UnknownPartitioning} +import org.apache.spark.sql.execution.{ColumnarToRowExec, LeafExecNode, SparkPlan} +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.vectorized.ColumnarBatch + +/** What the result task needs to create the native graph: no plan objects. */ +private[local] case class LocalQuerySpec(plan: Array[Byte], partitions: Int, columns: Int) + +/** + * A whole query run as one native graph by the single Spark task of its result. `block` is the + * Comet native block the graph is created from, kept for explain only. + */ +case class CometLocalExec private[local] ( + override val output: Seq[Attribute], + spec: LocalQuerySpec, + ordering: Seq[SortOrder], + block: SparkPlan) + extends LeafExecNode { + override def supportsColumnar: Boolean = true + + // The result is one partition, but its rows are not clustered by anything. A query over a + // cached local result must still exchange to get its required distribution. + override def outputPartitioning: Partitioning = UnknownPartitioning(1) + + // Partitions are delivered in order, so the block's ordering holds across the result. + override def outputOrdering: Seq[SortOrder] = ordering + + override def innerChildren: Seq[QueryPlan[_]] = Seq(block) + + override def stringArgs: Iterator[Any] = Iterator(output, s"partitions=${spec.partitions}") + + override lazy val metrics: Map[String, SQLMetric] = Map( + "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"), + "numOutputBatches" -> SQLMetrics.createMetric(sparkContext, "number of output batches")) + + override protected def doExecute(): RDD[InternalRow] = ColumnarToRowExec(this).execute() + + override protected def doExecuteColumnar(): RDD[ColumnarBatch] = { + val rows = longMetric("numOutputRows") + val batches = longMetric("numOutputBatches") + new LocalQueryRDD(sparkContext, spec).mapPartitions { input => + input.map { batch => + rows += batch.numRows().toLong + batches += 1L + batch + } + } + } +} + +private[local] class LocalQueryRDD(sc: SparkContext, spec: LocalQuerySpec) + extends RDD[ColumnarBatch](sc, Nil) { + override protected def getPartitions: Array[Partition] = Array(new Partition { + override def index: Int = 0 + }) + + override def compute(split: Partition, context: TaskContext): Iterator[ColumnarBatch] = { + // A retried result task creates an entirely fresh native graph. No partition of an + // existing graph is replayed. Creation stays lazy, so explain alone allocates nothing. + new LocalQueryIterator(spec, context) + } +} diff --git a/spark-local/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala b/spark-local/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala new file mode 100644 index 00000000000..61b96d27baa --- /dev/null +++ b/spark-local/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala @@ -0,0 +1,122 @@ +/* + * 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.local + +import java.util.concurrent.ConcurrentHashMap +import java.util.concurrent.atomic.AtomicLong + +import scala.collection.mutable.ArrayBuffer + +import org.apache.spark.SparkException +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder, UnsafeRow} +import org.apache.spark.sql.catalyst.plans.physical.Partitioning +import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} + +/** + * Row result boundary of a local query. Spark's collect encodes and compresses every row of a + * result partition in its task, then decodes them on the driver. A local query has one result + * partition, so that work would run on a single thread. Admission requires an in-process local + * master, so collect and take still run one Spark task (keeping cancellation, job groups and SQL + * metrics) but hand copied rows to the driver in this JVM. + */ +case class CometLocalResultExec(child: SparkPlan) extends UnaryExecNode { + override def output: Seq[Attribute] = child.output + override def outputPartitioning: Partitioning = child.outputPartitioning + override def outputOrdering: Seq[SortOrder] = child.outputOrdering + + override protected def doExecute(): RDD[InternalRow] = child.execute() + + override def executeCollect(): Array[InternalRow] = + if (sparkContext.isLocal) LocalResultHandoff.collect(execute(), -1, maxResultSize) + else super.executeCollect() + + override def executeTake(n: Int): Array[InternalRow] = + if (n <= 0) Array.empty + else if (sparkContext.isLocal) LocalResultHandoff.collect(execute(), n, maxResultSize) + else super.executeTake(n) + + // Spark enforces this on serialized task results; this boundary enforces it while copying. + private def maxResultSize: Long = + sparkContext.getConf.getSizeAsBytes("spark.driver.maxResultSize", "1g") + + override protected def withNewChildInternal(newChild: SparkPlan): CometLocalResultExec = + copy(child = newChild) +} + +/** Driver-owned result slots, keyed by numeric IDs that the task closure carries. */ +private[local] object LocalResultHandoff { + private final class Slot(val maxResultSize: Long) { + @volatile var rows: Array[InternalRow] = _ + @volatile var exceededBytes: Long = -1L + } + + private val slots = new ConcurrentHashMap[Long, Slot]() + private val nextId = new AtomicLong() + + /** Number of live slots; zero whenever no collect or take is running. */ + def active: Int = slots.size() + + /** + * Runs the single result partition of `rdd` and returns at most `take` rows (all if negative). + * Stops reading, which closes the native query, once the copied rows exceed `maxResultSize` + * bytes (unlimited if not positive). Unlike Spark's compressed size, this counts uncompressed + * UnsafeRow bytes, so it is the more conservative of the two. + */ + def collect(rdd: RDD[InternalRow], take: Int, maxResultSize: Long): Array[InternalRow] = { + require(rdd.getNumPartitions == 1, "Local results have exactly one partition") + val id = nextId.incrementAndGet() + val slot = new Slot(maxResultSize) + slots.put(id, slot) + try { + rdd.sparkContext.runJob(rdd, (rows: Iterator[InternalRow]) => fill(id, rows, take), Seq(0)) + if (slot.exceededBytes >= 0) { + throw new SparkException( + s"Total size of local query results (at least ${slot.exceededBytes} bytes, " + + s"uncompressed) is bigger than spark.driver.maxResultSize ($maxResultSize bytes)") + } + Option(slot.rows).getOrElse(throw new IllegalStateException("Local result not delivered")) + } finally { + slots.remove(id) + () + } + } + + // Runs in the result task. A retried attempt rebuilds the result and publishes it only once + // complete. A slot the driver already removed means nobody will read the rows. + private def fill(id: Long, rows: Iterator[InternalRow], take: Int): Unit = { + val slot = slots.get(id) + if (slot == null) return + val buffer = new ArrayBuffer[InternalRow]() + var bytes = 0L + while ((take < 0 || buffer.length < take) && rows.hasNext) { + // ColumnarToRow reuses one UnsafeRow; Spark's collect makes the same assumption. + val row = rows.next().asInstanceOf[UnsafeRow].copy() + bytes += row.getSizeInBytes + if (slot.maxResultSize > 0 && bytes > slot.maxResultSize) { + slot.exceededBytes = bytes + return + } + buffer += row + } + slot.rows = buffer.toArray + } +} diff --git a/spark-local/src/main/scala/org/apache/comet/local/CometLocalRule.scala b/spark-local/src/main/scala/org/apache/comet/local/CometLocalRule.scala new file mode 100644 index 00000000000..1c600815e44 --- /dev/null +++ b/spark-local/src/main/scala/org/apache/comet/local/CometLocalRule.scala @@ -0,0 +1,135 @@ +/* + * 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.local + +import scala.jdk.CollectionConverters._ + +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.expressions.SubqueryExpression +import org.apache.spark.sql.comet.{CometNativeExec, CometRangeExec} +import org.apache.spark.sql.execution.{CollectLimitExec, QueryExecution, SparkPlan} +import org.apache.spark.sql.internal.SQLConf + +import com.google.protobuf.Descriptors.FieldDescriptor.JavaType +import com.google.protobuf.Message + +import org.apache.comet.CometConf +import org.apache.comet.serde.ExprOuterClass.Expr +import org.apache.comet.shims.ShimCometStreaming + +/** + * Admission of whole queries to local execution. A query is admitted only as a whole, from the + * plan Comet's ordinary conversion produces: one native block with no JVM input, whose operators + * then run for all partitions in one native graph. Everything else keeps ordinary execution. + */ +private[comet] object CometLocalRule { + + /** At most this many partitions run in one graph. */ + private val MaxPartitions = 1024 + + /** Expressions that call back into the JVM, which a local query does not do. */ + private val JvmCallbacks = Set(Expr.ExprStructCase.SUBQUERY, Expr.ExprStructCase.JVM_SCALAR_UDF) + + /** An already admitted local query, with or without its result boundary or limit wrapper. */ + def isLocal(plan: SparkPlan): Boolean = plan match { + case _: CometLocalExec | _: CometLocalResultExec => true + case limit: CollectLimitExec => isLocal(limit.child) + case _ => false + } + + /** + * `plan` as a local query if it is admitted, and otherwise as `convert`, Comet's ordinary + * conversion, produces it. A root `CollectLimitExec` (from `take`, `head` or `show`) is kept: + * it reads the result through `CometLocalResultExec.executeTake`. + */ + def apply( + session: SparkSession, + plan: SparkPlan, + convert: SparkPlan => SparkPlan): SparkPlan = { + if (!CometConf.COMET_EXEC_LOCAL_ENABLED + .get(SQLConf.get) || !environmentAdmits(session, plan)) { + return convert(plan) + } + val limit = plan match { + case limit: CollectLimitExec if limit.offset == 0 => Some(limit) + case _ => None + } + convert(limit.map(_.child).getOrElse(plan)) match { + case block: CometNativeExec if admits(block) => + val partitions = block.collectLeaves().head.outputPartitioning.numPartitions + val spec = LocalQuerySpec(block.serializedPlanOpt.plan.get, partitions, block.output.size) + val local = CometLocalResultExec( + CometLocalExec(block.output, spec, block.outputOrdering, block)) + limit.map(_.copy(child = local)).getOrElse(local) + case converted if limit.isEmpty => converted + case _ => convert(plan) + } + } + + private def environmentAdmits(session: SparkSession, plan: SparkPlan): Boolean = { + session.sparkContext.isLocal && + !ShimCometStreaming.isStreamingPlan(plan) && + // As in CometRule's plan-only reporting, this identifies subquery preparation. + !Thread + .currentThread() + .getStackTrace + .exists(frame => + frame.getClassName == QueryExecution.getClass.getName && + frame.getMethodName == "prepareExecutedPlan") + } + + /** + * A single native block that a Spark task could run for each partition on its own: its root + * carries the serialized plan, every other operator is native and part of the same block, its + * only leaf is a range, and its expressions are deterministic and do not call back into the + * JVM. Running all partitions in one graph then produces what the Spark tasks would. + */ + private def admits(block: CometNativeExec): Boolean = { + block.serializedPlanOpt.isDefined && + block.output.nonEmpty && + block.collectLeaves().map(_.getClass) == Seq(classOf[CometRangeExec]) && + block.collectFirst { + case node: CometNativeExec if (node ne block) && node.serializedPlanOpt.isDefined => node + case node if !node.isInstanceOf[CometNativeExec] => node + }.isEmpty && + !block.exists(_.expressions.exists(e => + !e.deterministic || e.exists(_.isInstanceOf[SubqueryExpression]))) && { + val partitions = block.collectLeaves().head.outputPartitioning.numPartitions + partitions >= 1 && partitions <= MaxPartitions + } && + withoutJvmCallbacks(block.nativeOp) + } + + private def withoutJvmCallbacks(message: Message): Boolean = { + val accepted = message match { + case expression: Expr => !JvmCallbacks.contains(expression.getExprStructCase) + case _ => true + } + accepted && message.getAllFields.asScala.forall { case (field, value) => + if (field.getJavaType != JavaType.MESSAGE) { + true + } else if (field.isRepeated) { + value.asInstanceOf[java.util.List[Message]].asScala.forall(withoutJvmCallbacks) + } else { + withoutJvmCallbacks(value.asInstanceOf[Message]) + } + } + } +} diff --git a/spark-local/src/main/scala/org/apache/comet/local/LocalQueryIterator.scala b/spark-local/src/main/scala/org/apache/comet/local/LocalQueryIterator.scala new file mode 100644 index 00000000000..074218a809a --- /dev/null +++ b/spark-local/src/main/scala/org/apache/comet/local/LocalQueryIterator.scala @@ -0,0 +1,143 @@ +/* + * 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.local + +import org.apache.spark.{SparkEnv, TaskContext, TaskKilledException} +import org.apache.spark.sql.comet.CometExec +import org.apache.spark.sql.vectorized.ColumnarBatch + +import org.apache.comet.{CometConf, CometExecIterator, SparkErrorConverter} +import org.apache.comet.exceptions.CometQueryExecutionException +import org.apache.comet.vector.NativeUtil + +/** Task-thread-owned Arrow batches; native close can also interrupt a concurrent result pull. */ +private[local] class LocalQueryIterator(spec: LocalQuerySpec, context: TaskContext) + extends Iterator[ColumnarBatch] + with AutoCloseable { + private val native = new NativeLocal + private val util = new NativeUtil + private var id = 0L + private var closed = false + private var pending: ColumnarBatch = _ + private var previous: ColumnarBatch = _ + + // Shared by every native plan of the task, and released when the task completes. + private val taskMemory = CometExecIterator.taskMemory(context, CometExec.newIterId) + + // Install cleanup before allocating the native execution. + context.addTaskCompletionListener[Unit](_ => close()) + try { + checkCancellation() + // The memory pool, spill directories and session settings of any Comet plan in this task. + val memory = CometExecIterator.getMemoryConfig(SparkEnv.get.conf) + id = native.createQuery( + spec.plan, + CometExecIterator.serializeCometSQLConfs(), + spec.partitions, + spec.columns, + taskMemory.manager, + SparkEnv.get.blockManager.getLocalDiskDirs, + CometConf.COMET_BATCH_SIZE.get(), + memory.offHeapMode, + memory.memoryPoolType, + memory.memoryLimit, + context.taskAttemptId(), + context.cpus().toLong) + } catch { + case failure: Throwable => + throw executionFailure(failure) + } + + private def checkCancellation(): Unit = { + if (context.isInterrupted()) throw new TaskKilledException("Local query cancelled") + } + + private def closeAfterFailure(failure: Throwable): Unit = { + try close() + catch { case cleanup: Throwable => failure.addSuppressed(cleanup) } + } + + private def executionFailure(failure: Throwable): Throwable = { + closeAfterFailure(failure) + failure match { + case e: CometQueryExecutionException => SparkErrorConverter.convertToSparkException(e) + case other => other + } + } + + override def hasNext: Boolean = { + if (closed) return false + if (pending != null) return true + try { + if (previous != null) { + previous.close() + previous = null + } + pending = util + .getNextBatch( + spec.columns, + (arrays, schemas) => { + var rows = -2L + while (rows == -2L) { + checkCancellation() + rows = native.nextBatch(id, arrays, schemas) + } + checkCancellation() + rows + }) + .orNull + if (pending == null) close() + pending != null + } catch { + case failure: Throwable => + throw executionFailure(failure) + } + } + + override def next(): ColumnarBatch = { + if (!hasNext) throw new NoSuchElementException("Local query exhausted") + previous = pending + pending = null + previous + } + + override def close(): Unit = { + if (!closed) { + closed = true + var failure: Throwable = null + def attempt(cleanup: => Unit): Unit = { + try cleanup + catch { + case caught: Throwable => + if (failure == null) failure = caught else failure.addSuppressed(caught) + } + } + // Cancel native production before releasing imported Arrow buffers. + if (id != 0L) attempt(native.close(id)) + id = 0L + if (pending != null) attempt(pending.close()) + pending = null + if (previous != null) attempt(previous.close()) + previous = null + attempt(util.close()) + if (failure != null) throw failure + } + } +} diff --git a/spark-local/src/main/scala/org/apache/comet/local/NativeLocal.scala b/spark-local/src/main/scala/org/apache/comet/local/NativeLocal.scala new file mode 100644 index 00000000000..6444cc97046 --- /dev/null +++ b/spark-local/src/main/scala/org/apache/comet/local/NativeLocal.scala @@ -0,0 +1,58 @@ +/* + * 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.local + +import org.apache.spark.CometTaskMemoryManager + +import org.apache.comet.NativeBase + +/** + * Numeric handles, never native pointers. Every successful create needs close, including errors. + */ +private[local] class NativeLocal extends NativeBase { + + /** + * Creates and starts a local query over `partitions` partitions of a serialized Comet plan + * without JVM inputs, with the memory pool and session that `Native.createPlan` creates for the + * calling Spark task. Returns its handle. + */ + // scalastyle:off argcount + @native def createQuery( + plan: Array[Byte], + configs: Array[Byte], + partitions: Int, + columns: Int, + taskMemoryManager: CometTaskMemoryManager, + localDirs: Array[String], + batchSize: Int, + offHeapMode: Boolean, + memoryPoolType: String, + memoryLimit: Long, + taskAttemptId: Long, + taskCpus: Long): Long + // scalastyle:on argcount + + /** Exports the next batch; returns its row count, -1 at EOF, or -2 while none is ready. */ + @native def nextBatch(id: Long, arrays: Array[Long], schemas: Array[Long]): Long + + @native def close(id: Long): Unit + + @native def activeQueries(): Long +} diff --git a/spark-local/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala b/spark-local/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala new file mode 100644 index 00000000000..8fa75fc026d --- /dev/null +++ b/spark-local/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala @@ -0,0 +1,410 @@ +/* + * 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.local + +import java.util.concurrent.{Callable, Executors, TimeUnit} + +import scala.jdk.CollectionConverters._ +import scala.reflect.ClassTag + +import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach} +import org.scalatest.funsuite.AnyFunSuite + +import org.apache.spark.{SparkEnv, SparkException, TaskContext} +import org.apache.spark.sql.{DataFrame, Row, SparkSession} +import org.apache.spark.sql.execution.columnar.InMemoryRelation +import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike + +import org.apache.comet.{CometArrowImportAllocator, CometConf, CometExecIterator, CometSparkSessionExtensions, NativeBase} +import org.apache.comet.vector.NativeUtil + +class CometLocalExecutionSuite + extends AnyFunSuite + with BeforeAndAfterAll + with BeforeAndAfterEach { + private var spark: SparkSession = _ + private val enabled = CometConf.COMET_EXEC_LOCAL_ENABLED.key + + override protected def beforeAll(): Unit = { + super.beforeAll() + spark = SparkSession + .builder() + .master("local[2]") + .appName("CometLocalExecutionSuite") + .config("spark.ui.enabled", "false") + .config("spark.sql.session.timeZone", "UTC") + .config("spark.sql.extensions", classOf[CometSparkSessionExtensions].getName) + // AQE stays at Spark's default: it does not wrap the exchange-free plans admitted here. + .config("spark.memory.offHeap.enabled", "true") + .config("spark.memory.offHeap.size", "256m") + .config("spark.comet.enabled", "true") + .config("spark.comet.exec.enabled", "true") + .config("spark.comet.shuffle.enabled", "false") + .config(CometConf.COMET_EXEC_RANGE_ENABLED.key, "true") + .config("spark.comet.batchSize", "17") + .config(enabled, "true") + .getOrCreate() + spark.sparkContext.setLogLevel("WARN") + assert(NativeBase.isLoaded) + } + + override protected def afterEach(): Unit = { + try { + assert(new NativeLocal().activeQueries() == 0) + assert(CometArrowImportAllocator.getAllocatedMemory == 0) + assert(LocalResultHandoff.active == 0) + } finally super.afterEach() + } + + override protected def afterAll(): Unit = { + try { + if (spark != null) spark.stop() + SparkSession.clearActiveSession() + SparkSession.clearDefaultSession() + } finally super.afterAll() + } + + private def withConf[T](key: String, value: String)(f: => T): T = { + val previous = spark.conf.getOption(key) + spark.conf.set(key, value) + try f + finally + previous match { + case Some(v) => spark.conf.set(key, v) + case None => spark.conf.unset(key) + } + } + + private def localNodes(df: DataFrame): Seq[CometLocalExec] = + df.queryExecution.executedPlan.collect { case p: CometLocalExec => p } + + /** Spark's own answer, with Comet disabled. */ + private def expected(query: => DataFrame): Seq[Row] = + withConf("spark.comet.enabled", "false")(query.collect().toSeq) + + /** Checks that `query` runs locally and returns Spark's rows in Spark's order. */ + private def checkLocal(query: => DataFrame): Unit = { + val actual = query + assert(localNodes(actual).size == 1, actual.queryExecution.executedPlan.toString) + assert(actual.collect().toSeq == expected(query)) + } + + /** Runs `f` in a Spark task, as the local query iterator does. */ + private def inTask[T: ClassTag](f: TaskContext => T): T = + spark.sparkContext + .runJob( + spark.sparkContext.parallelize(Seq(0), 1), + { _: Iterator[Int] => + f(TaskContext.get()) + }) + .head + + test("opt-in is required and admitted ranges match Spark in order") { + assert(!CometConf.COMET_EXEC_LOCAL_ENABLED.defaultValue.get) + withConf(enabled, "false") { + assert(localNodes(spark.range(0, 100, 1, 7).toDF()).isEmpty) + } + checkLocal(spark.range(0, 4099, 1, 7).selectExpr("id AS x", "id * 2 AS y")) + checkLocal(spark.range(30, -10, -3, 7).toDF()) + checkLocal(spark.range(0, 100, 7, 3).filter("id % 3 = 1").selectExpr("id + 1 AS v")) + checkLocal(spark.range(0, 1000, 1, 1).toDF()) + } + + test("ranges near the long boundaries match Spark, and overflowing ones stay on Spark") { + checkLocal(spark.range(Long.MinValue, Long.MaxValue, Long.MaxValue, 7).toDF()) + checkLocal(spark.range(Long.MaxValue, Long.MinValue, Long.MinValue, 3).toDF()) + // Spark's generated code overflows for these, so Comet keeps them on Spark. + val stepped = () => spark.range(Long.MinValue, Long.MaxValue, 1L << 62, 1).toDF() + assert(localNodes(stepped()).isEmpty, stepped().queryExecution.executedPlan.toString) + assert(stepped().collect().toSeq == expected(stepped())) + val wide = () => spark.range(Long.MinValue, Long.MaxValue, 1, 1).toDF() + assert(localNodes(wide()).isEmpty) + assert(wide().take(3).toSeq == withConf("spark.comet.enabled", "false")(wide().take(3).toSeq)) + } + + test("empty ranges are not admitted and return no rows") { + val query = spark.range(5, 5, 1, 7).toDF() + assert(localNodes(query).isEmpty) + assert(query.collect().isEmpty) + } + + test("partition-dependent, nondeterministic and subquery expressions stay on Comet") { + for (query <- Seq( + () => spark.range(0, 50, 1, 7).selectExpr("id", "spark_partition_id() AS p"), + () => spark.range(0, 50, 1, 7).selectExpr("id", "monotonically_increasing_id() AS m"), + () => spark.range(0, 50, 1, 7).selectExpr("id", "rand(3) AS r"))) { + assert(localNodes(query()).isEmpty, query().queryExecution.executedPlan.toString) + assert(query().collect().toSeq == expected(query())) + } + val subquery = spark.sql("SELECT (SELECT max(id) FROM range(5)) + id AS v FROM range(3)") + assert(subquery.queryExecution.executedPlan.collectWithSubqueries { case p: CometLocalExec => + p + }.isEmpty) + assert(subquery.collect().toSeq == Seq(Row(4L), Row(5L), Row(6L))) + } + + test("admission does not depend on AQE for exchange-free plans") { + for (adaptive <- Seq("true", "false")) { + withConf("spark.sql.adaptive.enabled", adaptive) { + checkLocal(spark.range(0, 500, 1, 5).selectExpr("id * 3 AS v")) + } + } + } + + test("ordinary Comet planning applies when Comet execution is off") { + withConf("spark.comet.exec.enabled", "false") { + assert(localNodes(spark.range(0, 10, 1, 2).toDF()).isEmpty) + } + withConf(CometConf.COMET_EXEC_RANGE_ENABLED.key, "false") { + assert(localNodes(spark.range(0, 10, 1, 2).toDF()).isEmpty) + } + } + + test("limits read partitions in order and take, head and show close early") { + val query = spark.range(0, 1000, 1, 7).toDF() + // Unordered limits return the first rows of the first partitions, as Spark does. + for (_ <- 0 until 3) { + assert( + query.limit(10).collect().toSeq == expected(spark.range(0, 1000, 1, 7).toDF().limit(10))) + assert(query.take(5).map(_.getLong(0)).toSeq == (0L until 5L)) + } + assert(query.head(0).isEmpty) + val unbounded = spark.range(0, Long.MaxValue, 1, 7).toDF() + assert(localNodes(unbounded).size == 1) + assert(unbounded.take(3).map(_.getLong(0)).toSeq == Seq(0L, 1L, 2L)) + unbounded.show(2) + } + + test("a query over a cached local result keeps its required distribution") { + withConf("spark.sql.adaptive.enabled", "false") { + val cached = spark.range(0, 1000, 1, 7).selectExpr("id", "id % 5 AS k") + assert(localNodes(cached).size == 1) + cached.cache() + try { + val counts = cached.groupBy("k").count() + // The cache holds the local result, a single partition with no clustering. + val cachedPlans = counts.queryExecution.withCachedData.collect { + case relation: InMemoryRelation => relation.cacheBuilder.cachedPlan + } + assert(cachedPlans.flatMap(_.collect { case p: CometLocalExec => p }).size == 1) + val plan = counts.queryExecution.executedPlan + assert(plan.collect { case e: ShuffleExchangeLike => e }.nonEmpty, plan.toString) + assert( + counts.collect().map(r => r.getLong(0) -> r.getLong(1)).toMap == + (0L until 5L).map(_ -> 200L).toMap) + } finally cached.unpersist(blocking = true) + } + } + + test("repeated actions build fresh executions and iterator results") { + val query = spark.range(0, 300, 3, 7).selectExpr("id AS value") + val rows = (0L until 300L by 3L).map(Row(_)) + for (_ <- 0 until 3) { + assert(query.collect().toSeq == rows) + assert(new NativeLocal().activeQueries() == 0) + } + assert(query.toLocalIterator().asScala.toSeq == rows) + } + + test("collect and take hand local results to the driver within a Spark job") { + val query = spark.range(0, 200, 1, 5).selectExpr("id", "id % 3 AS k") + assert(query.queryExecution.executedPlan.isInstanceOf[CometLocalResultExec]) + assert( + query.collect().toSeq == expected( + spark.range(0, 200, 1, 5).selectExpr("id", "id % 3 AS k"))) + assert(query.take(4).map(_.getLong(0)).toSeq == (0L until 4L)) + } + + test("job group cancellation interrupts a local collect") { + val query = spark.range(0, Long.MaxValue, 1, 7).toDF() + assert(query.queryExecution.executedPlan.isInstanceOf[CometLocalResultExec]) + val group = "comet-local-collect-cancel" + val pool = Executors.newSingleThreadExecutor() + try { + val future = pool.submit(new Callable[Int] { + override def call(): Int = { + spark.sparkContext.setJobGroup(group, "cancel local collect", interruptOnCancel = true) + try query.collect().length + finally spark.sparkContext.clearJobGroup() + } + }) + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(20) + while (new NativeLocal().activeQueries() == 0 && System.nanoTime() < deadline) { + Thread.sleep(10) + } + assert(new NativeLocal().activeQueries() == 1) + spark.sparkContext.cancelJobGroup(group) + intercept[java.util.concurrent.ExecutionException] { future.get(20, TimeUnit.SECONDS) } + while (new NativeLocal().activeQueries() != 0 && System.nanoTime() < deadline) { + Thread.sleep(10) + } + } finally pool.shutdownNow() + } + + test("local result size limit stops native production and fails the action") { + val query = spark.range(0, Long.MaxValue, 1, 7).selectExpr("id", "id AS other") + val plan = query.queryExecution.executedPlan + assert(plan.isInstanceOf[CometLocalResultExec], plan.toString) + // Each two-long UnsafeRow is 24 bytes, so 1,000 bytes fit 41 rows and the 42nd exceeds it. + val failure = intercept[SparkException] { + LocalResultHandoff.collect(plan.execute(), -1, 1000) + } + assert(failure.getMessage.contains("spark.driver.maxResultSize"), failure.getMessage) + assert(failure.getMessage.contains("at least 1008 bytes"), failure.getMessage) + assert(new NativeLocal().activeQueries() == 0) + // A take within the limit still succeeds against the same bound. + assert(LocalResultHandoff.collect(plan.execute(), 41, 1000).length == 41) + // Non-positive limits mean unlimited, as in Spark. + assert(LocalResultHandoff.collect(plan.execute(), 100, 0).length == 100) + } + + test("Spark job cancellation releases a live local execution") { + val executor = Executors.newSingleThreadExecutor() + val query = spark.range(0, Long.MaxValue, 1, 7).toDF() + assert(localNodes(query).size == 1) + val future = executor.submit(new Callable[Unit] { + override def call(): Unit = { + spark.sparkContext.setJobGroup("comet-local-cancel", "local cancellation", true) + query.queryExecution.executedPlan.execute().foreachPartition { rows => + while (rows.hasNext) rows.next() + } + } + }) + try { + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(20) + while (new NativeLocal().activeQueries() == 0 && !future.isDone && + System.nanoTime() < deadline) Thread.sleep(10) + assert(new NativeLocal().activeQueries() > 0, "native execution did not start") + spark.sparkContext.cancelJobGroup("comet-local-cancel") + intercept[java.util.concurrent.ExecutionException] { future.get(20, TimeUnit.SECONDS) } + while (new NativeLocal().activeQueries() != 0 && System.nanoTime() < deadline) + Thread.sleep(10) + assert(new NativeLocal().activeQueries() == 0) + } finally { + spark.sparkContext.cancelJobGroup("comet-local-cancel") + executor.shutdownNow() + } + } + + test("native creation and output errors fail without leaking handles") { + val spec = localNodes(spark.range(10, 20, 1, 1).toDF()).head.spec + val (failures, live) = inTask(CometLocalExecutionSuite.creationAndOutputErrors(spec, _)) + assert(failures == Seq(true, true, true)) + assert(live == 0) + } + + test("imported Arrow batches outlive native close") { + val spec = localNodes(spark.range(10, 20, 1, 1).toDF()).head.spec + assert(inTask(CometLocalExecutionSuite.readAfterClose(spec, _)) == (10L until 20L)) + } + + test("explain shows the native block and metrics count delivered rows and batches") { + val query = spark.range(0, 100, 1, 4).selectExpr("id + 1 AS v") + val local = localNodes(query).head + assert(local.toString.contains("partitions=4")) + assert(query.queryExecution.executedPlan.treeString.contains("CometRange")) + query.collect() + assert(local.metrics("numOutputRows").value == 100) + assert(local.metrics("numOutputBatches").value >= 4) + } + + test("the native graph takes its batch size from the session when the query runs") { + for (batchSize <- Seq(10, 17)) { + withConf("spark.comet.batchSize", batchSize.toString) { + val query = spark.range(0, 100, 1, 1).toDF() + assert(query.collect().length == 100) + val batches = localNodes(query).head.metrics("numOutputBatches").value + assert(batches == (100 + batchSize - 1) / batchSize) + } + } + } +} + +/** Bodies of tests that run in a Spark task, kept out of the suite so it is not serialized. */ +object CometLocalExecutionSuite { + + /** Creates a native query for `spec` in the calling task, as the iterator does. */ + private def createQuery( + native: NativeLocal, + spec: LocalQuerySpec, + context: TaskContext, + partitions: Int, + columns: Int): Long = { + val memory = CometExecIterator.getMemoryConfig(SparkEnv.get.conf) + native.createQuery( + spec.plan, + CometExecIterator.serializeCometSQLConfs(), + partitions, + columns, + CometExecIterator.taskMemory(context, 0L).manager, + SparkEnv.get.blockManager.getLocalDiskDirs, + 17, + memory.offHeapMode, + memory.memoryPoolType, + memory.memoryLimit, + context.taskAttemptId(), + 1L) + } + + private def fails(f: => Any): Boolean = + try { + f + false + } catch { case _: Exception => true } + + /** Whether invalid creation and output pointers fail, and the live queries left after. */ + def creationAndOutputErrors( + spec: LocalQuerySpec, + context: TaskContext): (Seq[Boolean], Long) = { + val native = new NativeLocal + val noPartitions = fails(createQuery(native, spec, context, partitions = 0, columns = 1)) + val wrongWidth = fails(createQuery(native, spec, context, partitions = 1, columns = 2)) + val id = createQuery(native, spec, context, partitions = 1, columns = 1) + val nullOutput = fails(native.nextBatch(id, Array(0L), Array(0L))) + native.close(id) + native.close(id) + (Seq(noPartitions, wrongWidth, nullOutput), native.activeQueries()) + } + + /** The values of the first batch, read after its native query is closed. */ + def readAfterClose(spec: LocalQuerySpec, context: TaskContext): Seq[Long] = { + val native = new NativeLocal + val util = new NativeUtil + val id = createQuery(native, spec, context, partitions = 1, columns = 1) + try { + val batch = util + .getNextBatch( + 1, + (arrays, schemas) => { + var rows = -2L + while (rows == -2L) rows = native.nextBatch(id, arrays, schemas) + rows + }) + .get + try { + native.close(id) + (0 until batch.numRows()).map(batch.column(0).getLong) + } finally batch.close() + } finally { + native.close(id) + util.close() + } + } +} diff --git a/spark/pom.xml b/spark/pom.xml index 6b71b199eb8..94546e8a725 100644 --- a/spark/pom.xml +++ b/spark/pom.xml @@ -683,6 +683,21 @@ under the License. + + org.scalastyle + scalastyle-maven-plugin + + + check-local-source-module + compile + check + + ${project.basedir}/../spark-local/src/main/scala + ${project.build.directory}/scalastyle-local-output.xml + + + + io.github.git-commit-id git-commit-id-maven-plugin @@ -877,6 +892,7 @@ under the License. + ../spark-local/src/test/scala src/test/${shims.majorVerSrc} src/test/${shims.minorPlusVerSrc} src/test/${shims.minorVerSrc} @@ -891,6 +907,7 @@ under the License. + ../spark-local/src/main/scala src/main/${shims.majorVerSrc} src/main/${shims.minorPlusVerSrc} src/main/${shims.minorVerSrc} diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 9bd32816c64..21bab80553b 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -299,6 +299,16 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(true) + val COMET_EXEC_LOCAL_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.exec.local.enabled") + .category(CATEGORY_EXEC) + .internal() + .doc( + "Experimental. When the Spark application runs with a local master, run a whole query " + + "that Comet executes natively without JVM input as one native graph in a single " + + "Spark task. Other queries use ordinary Comet execution.") + .booleanConf + .createWithDefault(false) + val COMET_EXEC_PROJECT_ENABLED: ConfigEntry[Boolean] = createExecEnabledConfig("project", defaultValue = true) val COMET_EXEC_FILTER_ENABLED: ConfigEntry[Boolean] = 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 1a1e033bf60..5bd8560d884 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometRule.scala @@ -33,6 +33,7 @@ import org.apache.spark.sql.internal.SQLConf import org.apache.comet.{CometConf, ExtendedExplainInfo} import org.apache.comet.CometSparkSessionExtensions.isCometLoaded +import org.apache.comet.local.CometLocalRule import org.apache.comet.shims.ShimCometStreaming object CometRule { @@ -162,8 +163,14 @@ case class CometRule(session: SparkSession, queryStagePrep: Boolean = false) if (CometRule.planOnlyApplies(conf, plan)) { reportPlanOnlyCoverage(plan) plan - } else { + } else if (CometLocalRule.isLocal(plan)) { + plan + } else if (queryStagePrep) { convert(plan) + } else { + // Local execution admits only exchange-free plans, which AQE does not wrap, so this + // instance sees them whether or not AQE is enabled. + CometLocalRule(session, plan, convert) } } From 49dbdf859da3a8cc9022fb75c0633bb7da7c183e Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Mon, 12 Oct 2026 23:06:46 -0700 Subject: [PATCH 02/12] fix: keep Spark's ordering and limit semantics in local results Spark's outputOrdering holds within each partition, but a local result concatenates its partitions. Report the block's ordering only when it holds across them: with one partition, or when the partitions are ranges of a prefix of that ordering. Otherwise a cached result sorted within partitions claimed a global order and a window over it skipped its sort. Report a partition's error when its turn comes rather than at once. A limit satisfied by earlier partitions no longer fails on a later partition that Spark's executeTake would never evaluate. Replace the AtomicI64::fetch_update call, deprecated in current stable Rust, with fetch_add. --- .../contributor-guide/local-execution.md | 15 ++-- native/core/src/local.rs | 9 ++- native/local/src/lib.rs | 74 ++++++------------- native/local/tests/execution.rs | 40 ++++++++-- .../apache/comet/local/CometLocalExec.scala | 2 +- .../apache/comet/local/CometLocalRule.scala | 21 +++++- .../local/CometLocalExecutionSuite.scala | 59 +++++++++++++++ 7 files changed, 149 insertions(+), 71 deletions(-) diff --git a/docs/source/contributor-guide/local-execution.md b/docs/source/contributor-guide/local-execution.md index 87d9250fe82..f7c1cf39a02 100644 --- a/docs/source/contributor-guide/local-execution.md +++ b/docs/source/contributor-guide/local-execution.md @@ -60,13 +60,16 @@ single partition. That partition's task creates the native graph: Comet execution. - `native/local` runs all partitions concurrently on Comet's Tokio runtime and delivers them in partition order, each up to two batches ahead. The result therefore has Spark's row order, and - a limit returns the same rows as Spark's `executeTake`. The first error of any partition ends - the query at once. + a limit returns the same rows as Spark's `executeTake`. A partition's error is reported when its + turn comes, so a limit satisfied by earlier partitions does not fail on a later one, as with + `executeTake`, which never evaluates the partitions it does not need. `CometLocalExec` reports `UnknownPartitioning(1)`: the result is one partition, but its rows are not clustered by anything, so a query over a cached local result still gets the exchange -its distribution requires. It reports the block's ordering, which delivery in partition order -preserves. +its distribution requires. Spark's ordering of the block holds within each partition, so the +result reports it only when it holds across the concatenated partitions: with one partition, or +when the partitions are ranges of a prefix of that ordering, as a range's slices are. Sorting +within partitions, for example, reports no ordering. Explain shows the native block under `CometLocalExec`. Its native operators do not report SQL metrics; `CometLocalExec` counts the rows and batches it delivers. @@ -112,8 +115,8 @@ cd .. ./mvnw test -Dsuites=org.apache.comet.local.CometLocalExecutionSuite ``` -The native tests cover partition-order delivery, early failure, cancellation with busy workers, +The native tests cover partition-order delivery, errors reported in a partition's turn, cancellation with busy workers, the one-batch handoff and runtime shutdown. `CometLocalExecutionSuite` compares admitted queries with Spark's results and order, and covers ranges Spark would overflow, declined expressions, -AQE, limits, caching, the result size limit, job and task cancellation, and native errors. After +AQE, reported orderings, limits that stop before a failing partition, caching, the result size limit, job and task cancellation, and native errors. After each test it checks that no native query, imported Arrow memory or result slot is left. diff --git a/native/core/src/local.rs b/native/core/src/local.rs index 60c836f66c9..10d0ed1abec 100644 --- a/native/core/src/local.rs +++ b/native/core/src/local.rs @@ -171,9 +171,12 @@ fn local_plan(context: &ExecutionContext, partitions: jint) -> CometResult CometResult { - NEXT_ID - .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| id.checked_add(1)) - .map_err(|_| CometError::Internal("Local query IDs exhausted".into())) + // IDs are never reused. A counter that wrapped would reuse them, so fail instead. + let id = NEXT_ID.fetch_add(1, Ordering::Relaxed); + if id < 1 { + return Err(CometError::Internal("Local query IDs exhausted".into())); + } + Ok(id) } fn register(id: i64, entry: Entry) -> CometResult<()> { diff --git a/native/local/src/lib.rs b/native/local/src/lib.rs index 360a331ccb0..f843cb3ca51 100644 --- a/native/local/src/lib.rs +++ b/native/local/src/lib.rs @@ -40,7 +40,6 @@ use datafusion::physical_plan::{ SendableRecordBatchStream, }; use futures::future::poll_fn; -use futures::task::AtomicWaker; use futures::{FutureExt, Stream, StreamExt}; use tokio::task::JoinHandle; @@ -104,26 +103,22 @@ fn in_partition_order( plan: &Arc, context: &Arc, ) -> Result { - let failure = Arc::new(Failure::default()); let runtime = tokio::runtime::Handle::current(); let partitions = (0..plan.output_partitioning().partition_count()) .map(|partition| { let stream = plan.execute(partition, Arc::clone(context))?; - Ok(Partition::spawn(&runtime, stream, Arc::clone(&failure))) + Ok(Partition::spawn(&runtime, stream)) }) .collect::>>()?; Ok(Box::pin(RecordBatchStreamAdapter::new( plan.schema(), - PartitionOrder { - partitions, - failure, - }, + PartitionOrder { partitions }, ))) } /// One root partition running on its own Tokio task, a few batches ahead of delivery. struct Partition { - batches: tokio::sync::mpsc::Receiver, + batches: tokio::sync::mpsc::Receiver>, /// Set once the task has sent the stream's last batch. The channel closes whenever the /// task ends, and a runtime that shuts down cancels its tasks, so only this tells that /// every batch arrived. @@ -135,11 +130,7 @@ struct Partition { } impl Partition { - fn spawn( - runtime: &tokio::runtime::Handle, - stream: SendableRecordBatchStream, - failure: Arc, - ) -> Self { + fn spawn(runtime: &tokio::runtime::Handle, stream: SendableRecordBatchStream) -> Self { let (sender, batches) = tokio::sync::mpsc::channel(PARTITION_BUFFER); let ended = Arc::new(AtomicBool::new(false)); let stream = Arc::new(Mutex::new(Some(stream))); @@ -154,8 +145,8 @@ impl Partition { }) .await; match next { - Some(Ok(batch)) => permit.send(batch), - Some(Err(error)) => return failure.record(error), + Some(Ok(batch)) => permit.send(Ok(batch)), + Some(Err(error)) => return permit.send(Err(error)), None => { // A taken stream also ends here, but then nobody reads the flag. task_ended.store(true, Ordering::Release); @@ -167,9 +158,11 @@ impl Partition { .catch_unwind() .await; if result.is_err() { - failure.record(DataFusionError::Execution( - "Local query partition panicked".into(), - )); + let _ = sender + .send(Err(DataFusionError::Execution( + "Local query partition panicked".into(), + ))) + .await; } }); Self { @@ -191,58 +184,35 @@ impl Drop for Partition { } } -/// The first error of any partition, reported as soon as it happens rather than when the -/// failed partition's turn comes. -#[derive(Default)] -struct Failure { - error: Mutex>, - waker: AtomicWaker, -} - -impl Failure { - /// Keeps the first error and wakes the consumer. - fn record(&self, error: DataFusionError) { - lock(&self.error).get_or_insert(error); - self.waker.wake(); - } - - fn take(&self) -> Option { - lock(&self.error).take() - } -} - +/// Delivers the partitions in turn. A partition's error is reported when its turn comes, as +/// Spark's `executeTake` evaluates a later partition only when it needs that partition's rows: +/// a consumer that stops early never sees the failure of a partition it did not reach. struct PartitionOrder { partitions: VecDeque, - failure: Arc, } impl Stream for PartitionOrder { type Item = Result; fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - self.failure.waker.register(cx.waker()); loop { - if let Some(error) = self.failure.take() { - self.partitions.clear(); - return Poll::Ready(Some(Err(error))); - } let Some(partition) = self.partitions.front_mut() else { return Poll::Ready(None); }; match partition.batches.poll_recv(cx) { - Poll::Ready(Some(batch)) => return Poll::Ready(Some(Ok(batch))), + Poll::Ready(Some(Ok(batch))) => return Poll::Ready(Some(Ok(batch))), + Poll::Ready(Some(Err(error))) => { + self.partitions.clear(); + return Poll::Ready(Some(Err(error))); + } Poll::Ready(None) if partition.ended.load(Ordering::Acquire) => { self.partitions.pop_front(); } Poll::Ready(None) => { - // A failed partition records its error before its channel closes. - let error = self.failure.take().unwrap_or_else(|| { - DataFusionError::Execution( - "Local query partition stopped without completing".into(), - ) - }); self.partitions.clear(); - return Poll::Ready(Some(Err(error))); + return Poll::Ready(Some(Err(DataFusionError::Execution( + "Local query partition stopped without completing".into(), + )))); } Poll::Pending => return Poll::Pending, } diff --git a/native/local/tests/execution.rs b/native/local/tests/execution.rs index 2fc13831157..1d3c8a6bf20 100644 --- a/native/local/tests/execution.rs +++ b/native/local/tests/execution.rs @@ -483,21 +483,30 @@ async fn partition_order_delivers_partitions_in_turn() { assert_eq!(values, (0..7 * 9 * 4).collect::>()); } -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn partition_order_reports_a_later_failure_without_waiting_for_earlier_partitions() { - let blocking = BlockingExec::new(schema(), 1); - let refs = blocking.refs(); - let failing = MockExec::new( +fn failing() -> Arc { + Arc::new(MockExec::new( vec![Err(DataFusionError::Execution( "injected late failure".into(), ))], schema(), - ); - let plan = UnionExec::try_new(vec![Arc::new(blocking), Arc::new(failing)]).unwrap(); + )) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn partition_order_reports_a_later_failure_in_its_turn() { + let plan = UnionExec::try_new(vec![ordered_input(1, 3), failing()]).unwrap(); let mut stream = LocalQuery::new(plan, context()) .in_partition_order() .execute() .unwrap(); + // The earlier partition's batches come first, even though the later one failed at once. + for _ in 0..3 { + timeout(DEADLINE, stream.next()) + .await + .unwrap() + .unwrap() + .unwrap(); + } let error = timeout(DEADLINE, stream.next()) .await .unwrap() @@ -505,6 +514,23 @@ async fn partition_order_reports_a_later_failure_without_waiting_for_earlier_par .unwrap_err(); assert!(error.to_string().contains("injected late failure")); assert!(stream.next().await.is_none()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn partition_order_never_reports_the_failure_of_a_partition_not_reached() { + let blocking = BlockingExec::new(schema(), 1); + let refs = blocking.refs(); + let plan = UnionExec::try_new(vec![Arc::new(blocking), failing()]).unwrap(); + let mut stream = LocalQuery::new(plan, context()) + .in_partition_order() + .execute() + .unwrap(); + // A consumer that stops in the first partition, as a limit does, sees no error. + assert!(timeout(Duration::from_millis(200), stream.next()) + .await + .is_err()); + stream.cancel(); + assert!(stream.next().await.is_none()); wait_released(&refs).await; } diff --git a/spark-local/src/main/scala/org/apache/comet/local/CometLocalExec.scala b/spark-local/src/main/scala/org/apache/comet/local/CometLocalExec.scala index 726a4827f9e..8e623ac0d61 100644 --- a/spark-local/src/main/scala/org/apache/comet/local/CometLocalExec.scala +++ b/spark-local/src/main/scala/org/apache/comet/local/CometLocalExec.scala @@ -48,7 +48,7 @@ case class CometLocalExec private[local] ( // cached local result must still exchange to get its required distribution. override def outputPartitioning: Partitioning = UnknownPartitioning(1) - // Partitions are delivered in order, so the block's ordering holds across the result. + // The ordering of the whole result, which admission derives from the block's partitions. override def outputOrdering: Seq[SortOrder] = ordering override def innerChildren: Seq[QueryPlan[_]] = Seq(block) diff --git a/spark-local/src/main/scala/org/apache/comet/local/CometLocalRule.scala b/spark-local/src/main/scala/org/apache/comet/local/CometLocalRule.scala index 1c600815e44..69d49b21bef 100644 --- a/spark-local/src/main/scala/org/apache/comet/local/CometLocalRule.scala +++ b/spark-local/src/main/scala/org/apache/comet/local/CometLocalRule.scala @@ -22,7 +22,8 @@ package org.apache.comet.local import scala.jdk.CollectionConverters._ import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.expressions.SubqueryExpression +import org.apache.spark.sql.catalyst.expressions.{SortOrder, SubqueryExpression} +import org.apache.spark.sql.catalyst.plans.physical.RangePartitioning import org.apache.spark.sql.comet.{CometNativeExec, CometRangeExec} import org.apache.spark.sql.execution.{CollectLimitExec, QueryExecution, SparkPlan} import org.apache.spark.sql.internal.SQLConf @@ -76,13 +77,29 @@ private[comet] object CometLocalRule { val partitions = block.collectLeaves().head.outputPartitioning.numPartitions val spec = LocalQuerySpec(block.serializedPlanOpt.plan.get, partitions, block.output.size) val local = CometLocalResultExec( - CometLocalExec(block.output, spec, block.outputOrdering, block)) + CometLocalExec(block.output, spec, resultOrdering(block, partitions), block)) limit.map(_.copy(child = local)).getOrElse(local) case converted if limit.isEmpty => converted case _ => convert(plan) } } + /** + * The ordering of the whole result. Spark's `outputOrdering` holds within each partition, and + * the result concatenates the partitions in order, so it holds for the whole result only with + * one partition, or when the partitions are ranges of a prefix of that ordering, as a range's + * slices are. + */ + private def resultOrdering(block: CometNativeExec, partitions: Int): Seq[SortOrder] = { + val ordering = block.outputOrdering + block.outputPartitioning match { + case _ if partitions == 1 => ordering + case RangePartitioning(ranges, _) if SortOrder.orderingSatisfies(ordering, ranges) => + ordering + case _ => Nil + } + } + private def environmentAdmits(session: SparkSession, plan: SparkPlan): Boolean = { session.sparkContext.isLocal && !ShimCometStreaming.isStreamingPlan(plan) && diff --git a/spark-local/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala b/spark-local/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala index 8fa75fc026d..58ebc8f72e7 100644 --- a/spark-local/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala +++ b/spark-local/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala @@ -31,6 +31,8 @@ import org.apache.spark.{SparkEnv, SparkException, TaskContext} import org.apache.spark.sql.{DataFrame, Row, SparkSession} import org.apache.spark.sql.execution.columnar.InMemoryRelation import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike +import org.apache.spark.sql.expressions.Window +import org.apache.spark.sql.functions.{col, desc, row_number} import org.apache.comet.{CometArrowImportAllocator, CometConf, CometExecIterator, CometSparkSessionExtensions, NativeBase} import org.apache.comet.vector.NativeUtil @@ -213,6 +215,63 @@ class CometLocalExecutionSuite } } + test("only an ordering that holds across concatenated partitions is reported") { + // A range's slices are ranges of its ordering, so the result keeps it. + assert(localNodes(spark.range(0, 100, 1, 7).toDF()).head.outputOrdering.nonEmpty) + // Sorting within partitions orders each slice, not their concatenation. + val sorted = spark.range(0, 20, 1, 2).toDF().sortWithinPartitions(desc("id")) + assert(localNodes(sorted).size == 1) + assert(localNodes(sorted).head.outputOrdering.isEmpty) + assert( + localNodes( + spark + .range(0, 20, 1, 1) + .toDF() + .sortWithinPartitions(desc("id"))).head.outputOrdering.nonEmpty) + } + + test("a window over a cached partition-sorted result still sorts") { + withConf("spark.sql.adaptive.enabled", "false") { + val cached = spark.range(0, 20, 1, 2).toDF().sortWithinPartitions(desc("id")) + assert(localNodes(cached).size == 1) + cached.cache() + try { + cached.count() + val ranks = withConf("spark.comet.enabled", "false") { + cached + .select(col("id"), row_number().over(Window.orderBy(desc("id")))) + .collect() + .map(r => r.getLong(0) -> r.getInt(1)) + .toMap + } + assert(ranks == (0L until 20L).map(id => id -> (20 - id.toInt)).toMap) + } finally cached.unpersist(blocking = true) + } + } + + test("a limit does not fail on an error in a partition it does not reach") { + withConf("spark.sql.ansi.enabled", "true") { + def query(): DataFrame = + spark + .range(0, 200000, 1, 2) + .filter("id >= 99999") + .selectExpr("100L DIV (100000L - id) AS v") + .limit(1) + assert(localNodes(query()).size == 1, query().queryExecution.executedPlan.toString) + // Spark reads only the first partition; the second divides by zero at id 100000. + assert(query().collect().toSeq == Seq(Row(100L))) + assert(query().collect().toSeq == expected(query())) + // A full collect reaches the failing partition and fails as Spark does. + def all(): DataFrame = + spark + .range(0, 200000, 1, 2) + .filter("id >= 99999") + .selectExpr("100L DIV (100000L - id) AS v") + val failure = intercept[Exception](all().collect()) + assert(failure.getMessage.contains("DIVIDE_BY_ZERO"), failure.getMessage) + } + } + test("repeated actions build fresh executions and iterator results") { val query = spark.range(0, 300, 3, 7).selectExpr("id AS value") val rows = (0L until 300L by 3L).map(Row(_)) From 89f4adae9ee259fdce81b05b9cfc2301585f30ec Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Tue, 13 Oct 2026 20:13:02 -0700 Subject: [PATCH 03/12] fix: run iterator actions on local queries as ordinary Comet Spark's executeToIterator fetches one result partition at a time so that the driver holds at most one. A local query has a single partition holding the whole result, so toLocalIterator buffered all of it and could exceed spark.driver.maxResultSize where ordinary execution streams. CometLocalResultExec now runs the admitted native block as ordinary Comet execution for executeToIterator, one Spark task per partition. --- .../contributor-guide/local-execution.md | 10 ++++++++-- .../comet/local/CometLocalResultExec.scala | 14 +++++++++++++- .../comet/local/CometLocalExecutionSuite.scala | 18 ++++++++++++++++++ 3 files changed, 39 insertions(+), 3 deletions(-) diff --git a/docs/source/contributor-guide/local-execution.md b/docs/source/contributor-guide/local-execution.md index f7c1cf39a02..9bea4263cdd 100644 --- a/docs/source/contributor-guide/local-execution.md +++ b/docs/source/contributor-guide/local-execution.md @@ -89,7 +89,13 @@ them on the driver, all on one thread. While copying, the task counts uncompress bytes and stops reading, which closes the native query, once they exceed `spark.driver.maxResultSize`; the driver then fails the action. Spark counts compressed bytes, so this check is the stricter one. On a non-local master the node uses Spark's ordinary -collect. Other actions, such as `toLocalIterator`, keep Spark's behavior. +collect. + +An iterator fetches one result partition at a time so that it holds at most one, but a local +query has a single partition holding the whole result. `toLocalIterator`, which also serves the +Thrift server's incremental collect, therefore runs the native block as ordinary Comet execution, +one Spark task per partition. Consumers that read the result RDD partition by partition, such as +PySpark's `toLocalIterator`, still receive the whole result as one partition. ## Lifecycle and cancellation @@ -118,5 +124,5 @@ cd .. The native tests cover partition-order delivery, errors reported in a partition's turn, cancellation with busy workers, the one-batch handoff and runtime shutdown. `CometLocalExecutionSuite` compares admitted queries with Spark's results and order, and covers ranges Spark would overflow, declined expressions, -AQE, reported orderings, limits that stop before a failing partition, caching, the result size limit, job and task cancellation, and native errors. After +AQE, reported orderings, limits that stop before a failing partition, iterators, caching, the result size limit, job and task cancellation, and native errors. After each test it checks that no native query, imported Arrow memory or result slot is left. diff --git a/spark-local/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala b/spark-local/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala index 61b96d27baa..86c2b651d5f 100644 --- a/spark-local/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala +++ b/spark-local/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala @@ -29,7 +29,7 @@ import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder, UnsafeRow} import org.apache.spark.sql.catalyst.plans.physical.Partitioning -import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} +import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, UnaryExecNode} /** * Row result boundary of a local query. Spark's collect encodes and compresses every row of a @@ -54,6 +54,18 @@ case class CometLocalResultExec(child: SparkPlan) extends UnaryExecNode { else if (sparkContext.isLocal) LocalResultHandoff.collect(execute(), n, maxResultSize) else super.executeTake(n) + /** + * Runs the query as ordinary Comet execution instead. An iterator fetches one result partition + * at a time to bound what it holds, and a local query has a single partition holding the whole + * result, so the native block runs as Spark tasks, one per partition, as without local + * execution. + */ + override def executeToIterator(): Iterator[InternalRow] = + child + .collectFirst { case local: CometLocalExec => ColumnarToRowExec(local.block) } + .getOrElse(child) + .executeToIterator() + // Spark enforces this on serialized task results; this boundary enforces it while copying. private def maxResultSize: Long = sparkContext.getConf.getSizeAsBytes("spark.driver.maxResultSize", "1g") diff --git a/spark-local/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala b/spark-local/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala index 58ebc8f72e7..7a2face3292 100644 --- a/spark-local/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala +++ b/spark-local/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala @@ -282,6 +282,24 @@ class CometLocalExecutionSuite assert(query.toLocalIterator().asScala.toSeq == rows) } + test("an iterator fetches one partition at a time through ordinary execution") { + val query = spark.range(0, 400, 1, 8).selectExpr("id", "id * 2 AS v") + assert(query.queryExecution.executedPlan.isInstanceOf[CometLocalResultExec]) + val group = "comet-local-iterator" + spark.sparkContext.setJobGroup(group, "local iterator") + val rows = + try query.toLocalIterator().asScala.toList + finally spark.sparkContext.clearJobGroup() + assert(rows == expected(spark.range(0, 400, 1, 8).selectExpr("id", "id * 2 AS v"))) + // One job per partition, as Spark's executeToIterator runs, rather than one for the query. + val tracker = spark.sparkContext.statusTracker + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(20) + while (tracker.getJobIdsForGroup(group).length < 8 && System.nanoTime() < deadline) { + Thread.sleep(10) + } + assert(tracker.getJobIdsForGroup(group).length == 8) + } + test("collect and take hand local results to the driver within a Spark job") { val query = spark.range(0, 200, 1, 5).selectExpr("id", "id % 3 AS k") assert(query.queryExecution.executedPlan.isInstanceOf[CometLocalResultExec]) From b82f3878384044d6607c9c4fdf6cf09c77ddcbcd Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Mon, 19 Oct 2026 09:50:00 -0700 Subject: [PATCH 04/12] fix: run tails of local queries as ordinary Comet and move local sources Spark's executeTail scans partitions from the last one until it has enough rows, which PySpark's tail uses. A local query has a single partition, so a tail evaluated the whole query, including partitions Spark never reads. CometLocalResultExec now runs the admitted native block as ordinary Comet execution for executeTail, as it does for executeToIterator. Move the local execution sources from spark-local into the spark module, which they compile into anyway, removing the extra source roots, scalastyle execution, spotless include and CI change filters. --- dev/ci/check-ci-config.py | 4 ---- dev/ci/compute-changes.py | 12 ------------ pom.xml | 1 - spark/pom.xml | 17 ----------------- .../org/apache/comet/local/CometLocalExec.scala | 0 .../comet/local/CometLocalResultExec.scala | 17 +++++++++-------- .../org/apache/comet/local/CometLocalRule.scala | 0 .../apache/comet/local/LocalQueryIterator.scala | 0 .../org/apache/comet/local/NativeLocal.scala | 0 .../comet/local/CometLocalExecutionSuite.scala | 16 ++++++++++++++++ 10 files changed, 25 insertions(+), 42 deletions(-) rename {spark-local => spark}/src/main/scala/org/apache/comet/local/CometLocalExec.scala (100%) rename {spark-local => spark}/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala (90%) rename {spark-local => spark}/src/main/scala/org/apache/comet/local/CometLocalRule.scala (100%) rename {spark-local => spark}/src/main/scala/org/apache/comet/local/LocalQueryIterator.scala (100%) rename {spark-local => spark}/src/main/scala/org/apache/comet/local/NativeLocal.scala (100%) rename {spark-local => spark}/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala (96%) diff --git a/dev/ci/check-ci-config.py b/dev/ci/check-ci-config.py index d65f14ab983..3832fda92dc 100644 --- a/dev/ci/check-ci-config.py +++ b/dev/ci/check-ci-config.py @@ -127,10 +127,6 @@ MVN_JOBS = BUILD_JOBS | {"delta_gate", "pyarrow_udf"} ROUTING_CASES = [ - (["spark-local/src/main/scala/org/apache/comet/local/CometLocalRule.scala"], - BUILD_JOBS | {"delta_gate"}), - (["spark-local/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala"], - {"build_linux", "build_linux_full", "build_linux_all_profiles", "build_macos", "build_macos_full"}), # The Maven wrapper and its config feed every job that runs ./mvnw: the # Linux/macOS builds, setup-spark-builder, the Iceberg `mvnw install`, the # Delta gate's effective-pom check and the PyArrow suite's `mvnw install`. diff --git a/dev/ci/compute-changes.py b/dev/ci/compute-changes.py index 87b9b4ef730..4d62020bb83 100644 --- a/dev/ci/compute-changes.py +++ b/dev/ci/compute-changes.py @@ -43,7 +43,6 @@ "native/**", "common/**", "spark/**", - "spark-local/**", "spark-integration/**", "pom.xml", "**/pom.xml", @@ -81,7 +80,6 @@ "native/**", "common/**", "spark/**", - "spark-local/**", "spark-integration/**", "pom.xml", "**/pom.xml", @@ -118,7 +116,6 @@ "native/**", "common/src/main/**", "spark/src/main/**", - "spark-local/src/main/**", "contrib/delta/**", "pom.xml", "**/pom.xml", @@ -200,7 +197,6 @@ "common/src/main/**", "common/pom.xml", "spark/src/main/**", - "spark-local/src/main/**", "!spark/src/main/spark-3.5/**", "!spark/src/main/spark-4.0/**", "!spark/src/main/spark-4.1/**", @@ -229,7 +225,6 @@ "common/src/main/**", "common/pom.xml", "spark/src/main/**", - "spark-local/src/main/**", "!spark/src/main/spark-3.4/**", "!spark/src/main/spark-4.0/**", "!spark/src/main/spark-4.1/**", @@ -258,7 +253,6 @@ "common/src/main/**", "common/pom.xml", "spark/src/main/**", - "spark-local/src/main/**", "!spark/src/main/spark-3.4/**", "!spark/src/main/spark-3.5/**", "!spark/src/main/spark-3.x/**", @@ -287,7 +281,6 @@ "common/src/main/**", "common/pom.xml", "spark/src/main/**", - "spark-local/src/main/**", "!spark/src/main/spark-3.4/**", "!spark/src/main/spark-3.5/**", "!spark/src/main/spark-3.x/**", @@ -316,7 +309,6 @@ "common/src/main/**", "common/pom.xml", "spark/src/main/**", - "spark-local/src/main/**", "!spark/src/main/spark-3.4/**", "!spark/src/main/spark-3.5/**", "!spark/src/main/spark-3.x/**", @@ -350,7 +342,6 @@ "common/src/main/**", "common/pom.xml", "spark/src/main/**", - "spark-local/src/main/**", "!spark/src/main/scala/org/apache/comet/GenerateDocs.scala", "spark/pom.xml", "dev/diffs/iceberg/**", @@ -376,7 +367,6 @@ "common/src/main/**", "common/pom.xml", "spark/src/main/**", - "spark-local/src/main/**", "!spark/src/main/scala/org/apache/comet/GenerateDocs.scala", "spark/pom.xml", "dev/diffs/iceberg/**", @@ -402,7 +392,6 @@ "common/src/main/**", "common/pom.xml", "spark/src/main/**", - "spark-local/src/main/**", "!spark/src/main/scala/org/apache/comet/GenerateDocs.scala", "spark/pom.xml", "dev/diffs/iceberg/**", @@ -428,7 +417,6 @@ "common/src/main/**", "common/pom.xml", "spark/src/main/**", - "spark-local/src/main/**", "!spark/src/main/scala/org/apache/comet/GenerateDocs.scala", "spark/pom.xml", "dev/diffs/iceberg/**", diff --git a/pom.xml b/pom.xml index 6cad4cb4edb..72b33fed924 100644 --- a/pom.xml +++ b/pom.xml @@ -1201,7 +1201,6 @@ under the License. src/main/scala/**/*.scala src/test/scala/**/*.scala - spark-local/src/**/*.scala src/test/celeborn-reflection-compatibility/**/*.scala src/main/spark-*/**/*.scala diff --git a/spark/pom.xml b/spark/pom.xml index 94546e8a725..6b71b199eb8 100644 --- a/spark/pom.xml +++ b/spark/pom.xml @@ -683,21 +683,6 @@ under the License. - - org.scalastyle - scalastyle-maven-plugin - - - check-local-source-module - compile - check - - ${project.basedir}/../spark-local/src/main/scala - ${project.build.directory}/scalastyle-local-output.xml - - - - io.github.git-commit-id git-commit-id-maven-plugin @@ -892,7 +877,6 @@ under the License. - ../spark-local/src/test/scala src/test/${shims.majorVerSrc} src/test/${shims.minorPlusVerSrc} src/test/${shims.minorVerSrc} @@ -907,7 +891,6 @@ under the License. - ../spark-local/src/main/scala src/main/${shims.majorVerSrc} src/main/${shims.minorPlusVerSrc} src/main/${shims.minorVerSrc} diff --git a/spark-local/src/main/scala/org/apache/comet/local/CometLocalExec.scala b/spark/src/main/scala/org/apache/comet/local/CometLocalExec.scala similarity index 100% rename from spark-local/src/main/scala/org/apache/comet/local/CometLocalExec.scala rename to spark/src/main/scala/org/apache/comet/local/CometLocalExec.scala diff --git a/spark-local/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala b/spark/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala similarity index 90% rename from spark-local/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala rename to spark/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala index 86c2b651d5f..8abd68ff2fe 100644 --- a/spark-local/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala +++ b/spark/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala @@ -54,17 +54,18 @@ case class CometLocalResultExec(child: SparkPlan) extends UnaryExecNode { else if (sparkContext.isLocal) LocalResultHandoff.collect(execute(), n, maxResultSize) else super.executeTake(n) - /** - * Runs the query as ordinary Comet execution instead. An iterator fetches one result partition - * at a time to bound what it holds, and a local query has a single partition holding the whole - * result, so the native block runs as Spark tasks, one per partition, as without local - * execution. - */ - override def executeToIterator(): Iterator[InternalRow] = + // An iterator fetches one result partition at a time to bound what it holds, and a tail scans + // partitions from the last one until it has enough rows. A local query has a single partition + // holding the whole result, so both run the native block as ordinary Comet execution instead, + // one Spark task per partition, as without local execution. + override def executeToIterator(): Iterator[InternalRow] = ordinary.executeToIterator() + + override def executeTail(n: Int): Array[InternalRow] = ordinary.executeTail(n) + + private def ordinary: SparkPlan = child .collectFirst { case local: CometLocalExec => ColumnarToRowExec(local.block) } .getOrElse(child) - .executeToIterator() // Spark enforces this on serialized task results; this boundary enforces it while copying. private def maxResultSize: Long = diff --git a/spark-local/src/main/scala/org/apache/comet/local/CometLocalRule.scala b/spark/src/main/scala/org/apache/comet/local/CometLocalRule.scala similarity index 100% rename from spark-local/src/main/scala/org/apache/comet/local/CometLocalRule.scala rename to spark/src/main/scala/org/apache/comet/local/CometLocalRule.scala diff --git a/spark-local/src/main/scala/org/apache/comet/local/LocalQueryIterator.scala b/spark/src/main/scala/org/apache/comet/local/LocalQueryIterator.scala similarity index 100% rename from spark-local/src/main/scala/org/apache/comet/local/LocalQueryIterator.scala rename to spark/src/main/scala/org/apache/comet/local/LocalQueryIterator.scala diff --git a/spark-local/src/main/scala/org/apache/comet/local/NativeLocal.scala b/spark/src/main/scala/org/apache/comet/local/NativeLocal.scala similarity index 100% rename from spark-local/src/main/scala/org/apache/comet/local/NativeLocal.scala rename to spark/src/main/scala/org/apache/comet/local/NativeLocal.scala diff --git a/spark-local/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala b/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala similarity index 96% rename from spark-local/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala rename to spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala index 7a2face3292..a8dd5366ee3 100644 --- a/spark-local/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala @@ -300,6 +300,22 @@ class CometLocalExecutionSuite assert(tracker.getJobIdsForGroup(group).length == 8) } + test("a tail scans partitions from the last one through ordinary execution") { + withConf("spark.sql.ansi.enabled", "true") { + // As PySpark's tail does. The first partition divides by zero at id 0, but Spark reads only + // the last partition, which has enough rows. + def query(): DataFrame = spark.range(0, 20, 1, 2).selectExpr("100L DIV id AS v") + val plan = query().queryExecution.executedPlan + assert(plan.isInstanceOf[CometLocalResultExec], plan.toString) + assert(plan.executeTail(1).map(_.getLong(0)).toSeq == Seq(5L)) + val sparkTail = withConf("spark.comet.enabled", "false") { + query().queryExecution.executedPlan.executeTail(3).map(_.getLong(0)).toSeq + } + assert(plan.executeTail(3).map(_.getLong(0)).toSeq == sparkTail) + assert(plan.executeTail(0).isEmpty) + } + } + test("collect and take hand local results to the driver within a Spark job") { val query = spark.range(0, 200, 1, 5).selectExpr("id", "id % 3 AS k") assert(query.queryExecution.executedPlan.isInstanceOf[CometLocalResultExec]) From 61f954082349ac6fc213ab9991a0f65fc39986c6 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Mon, 19 Oct 2026 10:29:44 -0700 Subject: [PATCH 05/12] fix: bound how many local query partitions run at once Delivering partitions in order started every partition at once, and each partition's operators register their memory consumers when it starts. With many partitions the task's fair pool split into shares too small for a sort's up-front merge reservation, so a query that ordinary Comet runs failed with "Not enough memory to continue external sort". Run at most spark.executor.cores partitions at once, which local[N] resolves to N, and start the next when the first is delivered, as Spark runs at most that many tasks. Also parenthesize a typed lambda parameter in the suite, which the syntactic scalafix check cannot parse. --- .../contributor-guide/local-execution.md | 11 ++- native/core/src/execution/jni_api.rs | 2 +- native/core/src/local.rs | 7 +- native/local/src/lib.rs | 87 +++++++++++++------ native/local/tests/bridge.rs | 10 ++- native/local/tests/execution.rs | 41 ++++++++- .../local/CometLocalExecutionSuite.scala | 15 +++- 7 files changed, 131 insertions(+), 42 deletions(-) diff --git a/docs/source/contributor-guide/local-execution.md b/docs/source/contributor-guide/local-execution.md index 9bea4263cdd..c6c9ab94a89 100644 --- a/docs/source/contributor-guide/local-execution.md +++ b/docs/source/contributor-guide/local-execution.md @@ -58,8 +58,11 @@ single partition. That partition's task creates the native graph: Spark task would create them. Partition `i` of the graph is the output of the operators created for partition `i`, so per-partition state such as a range slice is the same as with ordinary Comet execution. -- `native/local` runs all partitions concurrently on Comet's Tokio runtime and delivers them in - partition order, each up to two batches ahead. The result therefore has Spark's row order, and +- `native/local` runs the partitions on Comet's Tokio runtime and delivers them in partition + order. At most `spark.executor.cores` partitions (N for `local[N]`) run at once, each up to two + batches ahead, and the next starts when the first is delivered, as Spark runs at most that + many tasks. A partition's operators register their memory consumers when it starts, so this + also bounds how many share the task's memory pool. The result therefore has Spark's row order, and a limit returns the same rows as Spark's `executeTake`. A partition's error is reported when its turn comes, so a limit satisfied by earlier partitions does not fail on a later one, as with `executeTake`, which never evaluates the partitions it does not need. @@ -121,8 +124,8 @@ cd .. ./mvnw test -Dsuites=org.apache.comet.local.CometLocalExecutionSuite ``` -The native tests cover partition-order delivery, errors reported in a partition's turn, cancellation with busy workers, +The native tests cover partition-order delivery, the bound on running partitions, errors reported in a partition's turn, cancellation with busy workers, the one-batch handoff and runtime shutdown. `CometLocalExecutionSuite` compares admitted queries with Spark's results and order, and covers ranges Spark would overflow, declined expressions, -AQE, reported orderings, limits that stop before a failing partition, iterators, caching, the result size limit, job and task cancellation, and native errors. After +AQE, reported orderings, limits that stop before a failing partition, iterators, many partitions with sorts in a small memory pool, caching, the result size limit, job and task cancellation, and native errors. After each test it checks that no native query, imported Arrow memory or result slot is left. diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 11eaa2bcb1a..94abb636e8e 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -720,7 +720,7 @@ impl ExecutionContext { /// What `createPlan` takes from the Spark configs. pub(crate) struct PlanSettings { - spark_config: HashMap, + pub(crate) spark_config: HashMap, debug_native: bool, strict_timestamp_labels: bool, explain_native: bool, diff --git a/native/core/src/local.rs b/native/core/src/local.rs index 10d0ed1abec..784a392b84d 100644 --- a/native/core/src/local.rs +++ b/native/core/src/local.rs @@ -41,6 +41,7 @@ use crate::execution::jni_api::{ PlanInputs, PlanJvmRefs, }; use crate::execution::memory_pools::create_memory_pool; +use crate::execution::spark_config::{SparkConfig, SPARK_EXECUTOR_CORES}; struct Entry { query: QueryHandle, @@ -92,6 +93,9 @@ pub extern "system" fn Java_org_apache_comet_local_NativeLocal_createQuery( } let start = Instant::now(); let settings = prepare_plan_settings(&env.convert_byte_array(serialized_spark_configs)?)?; + // Run as many partitions at once as Spark would run tasks, which also bounds how many + // operators share the memory pool. `local[N]` resolves to N here. + let max_running = settings.spark_config.get_usize(SPARK_EXECUTOR_CORES, 1); JVMClasses::init(env); let serialized_plan = env.convert_byte_array(serialized_query)?; let task_memory_manager = @@ -129,7 +133,8 @@ pub extern "system" fn Java_org_apache_comet_local_NativeLocal_createQuery( "Local output schema width mismatch".into(), )); } - let query = LocalQuery::new(plan, context.session_ctx.task_ctx()).in_partition_order(); + let query = + LocalQuery::new(plan, context.session_ctx.task_ctx()).in_partition_order(max_running); let query = QueryHandle::start(query, &get_runtime())?; register( id, diff --git a/native/local/src/lib.rs b/native/local/src/lib.rs index f843cb3ca51..8d465b73155 100644 --- a/native/local/src/lib.rs +++ b/native/local/src/lib.rs @@ -51,7 +51,8 @@ const PARTITION_BUFFER: usize = 2; pub struct LocalQuery { plan: Arc, context: Arc, - in_partition_order: bool, + /// Delivers the partitions in order, with at most this many running at once. + in_partition_order: Option, } impl LocalQuery { @@ -63,16 +64,19 @@ impl LocalQuery { Self { plan, context, - in_partition_order: false, + in_partition_order: None, } } /// Delivers the root partitions one after another, in partition order, as Spark returns - /// the partitions of a result. All partitions still run concurrently, each up to a few - /// batches ahead. Only for roots whose partitions do not share an exchange: a partition - /// that waits to be delivered would otherwise leave its exchange buffering for the others. - pub fn in_partition_order(mut self) -> Self { - self.in_partition_order = true; + /// the partitions of a result. Up to `max_running` partitions run at once, each up to a few + /// batches ahead, and the next one starts when the first is delivered, as Spark runs at most + /// as many tasks as it has cores. A partition's operators register their memory consumers + /// when it starts, so this also bounds how many share the memory pool. Only for roots whose + /// partitions do not share an exchange: a partition that waits to be delivered would + /// otherwise leave its exchange buffering for the others. + pub fn in_partition_order(mut self, max_running: usize) -> Self { + self.in_partition_order = Some(max_running.max(1)); self } @@ -86,8 +90,8 @@ impl LocalQuery { DataFusionError::Execution("LocalQuery requires an active Tokio runtime".into()) })?; let schema = self.plan.schema(); - let stream = if self.in_partition_order { - in_partition_order(&self.plan, &self.context)? + let stream = if let Some(max_running) = self.in_partition_order { + in_partition_order(&self.plan, &self.context, max_running)? } else { execute_stream(Arc::clone(&self.plan), Arc::clone(&self.context))? }; @@ -102,17 +106,21 @@ impl LocalQuery { fn in_partition_order( plan: &Arc, context: &Arc, + max_running: usize, ) -> Result { - let runtime = tokio::runtime::Handle::current(); - let partitions = (0..plan.output_partitioning().partition_count()) - .map(|partition| { - let stream = plan.execute(partition, Arc::clone(context))?; - Ok(Partition::spawn(&runtime, stream)) - }) - .collect::>>()?; + let mut order = PartitionOrder { + plan: Arc::clone(plan), + context: Arc::clone(context), + runtime: tokio::runtime::Handle::current(), + next: 0, + max_running, + running: VecDeque::new(), + }; + // Start the first partitions now, so that a failure to start them fails here. + order.start_partitions()?; Ok(Box::pin(RecordBatchStreamAdapter::new( plan.schema(), - PartitionOrder { partitions }, + order, ))) } @@ -188,7 +196,33 @@ impl Drop for Partition { /// Spark's `executeTake` evaluates a later partition only when it needs that partition's rows: /// a consumer that stops early never sees the failure of a partition it did not reach. struct PartitionOrder { - partitions: VecDeque, + plan: Arc, + context: Arc, + runtime: tokio::runtime::Handle, + /// The first partition not started yet. + next: usize, + max_running: usize, + /// Started partitions, the first of which is being delivered. + running: VecDeque, +} + +impl PartitionOrder { + fn start_partitions(&mut self) -> Result<()> { + let partitions = self.plan.output_partitioning().partition_count(); + while self.running.len() < self.max_running && self.next < partitions { + let stream = self.plan.execute(self.next, Arc::clone(&self.context))?; + self.running + .push_back(Partition::spawn(&self.runtime, stream)); + self.next += 1; + } + Ok(()) + } + + fn fail(&mut self, error: DataFusionError) -> Poll>> { + self.running.clear(); + self.next = usize::MAX; + Poll::Ready(Some(Err(error))) + } } impl Stream for PartitionOrder { @@ -196,23 +230,22 @@ impl Stream for PartitionOrder { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { loop { - let Some(partition) = self.partitions.front_mut() else { + if let Err(error) = self.start_partitions() { + return self.fail(error); + } + let Some(partition) = self.running.front_mut() else { return Poll::Ready(None); }; match partition.batches.poll_recv(cx) { Poll::Ready(Some(Ok(batch))) => return Poll::Ready(Some(Ok(batch))), - Poll::Ready(Some(Err(error))) => { - self.partitions.clear(); - return Poll::Ready(Some(Err(error))); - } + Poll::Ready(Some(Err(error))) => return self.fail(error), Poll::Ready(None) if partition.ended.load(Ordering::Acquire) => { - self.partitions.pop_front(); + self.running.pop_front(); } Poll::Ready(None) => { - self.partitions.clear(); - return Poll::Ready(Some(Err(DataFusionError::Execution( + return self.fail(DataFusionError::Execution( "Local query partition stopped without completing".into(), - )))); + )) } Poll::Pending => return Poll::Pending, } diff --git a/native/local/tests/bridge.rs b/native/local/tests/bridge.rs index 7d953aff8b9..c4402f139b7 100644 --- a/native/local/tests/bridge.rs +++ b/native/local/tests/bridge.rs @@ -78,7 +78,7 @@ fn partition_order_is_preserved_through_the_handoff() { let plan = ranges(&[(0, 10), (10, 10), (10, 25), (25, 26)], 3); let refs = Arc::downgrade(&plan); let query = QueryHandle::start( - LocalQuery::new(plan, Arc::new(TaskContext::default())).in_partition_order(), + LocalQuery::new(plan, Arc::new(TaskContext::default())).in_partition_order(usize::MAX), runtime.handle(), ) .unwrap(); @@ -189,7 +189,8 @@ fn backpressured_query_does_not_starve_other_queries_and_cancel_releases_graph() let stalled = ranges(&[(0, 100_000), (100_000, 200_000)], 5); let refs = Arc::downgrade(&stalled); let stalled = QueryHandle::start( - LocalQuery::new(stalled, Arc::new(TaskContext::default())).in_partition_order(), + LocalQuery::new(stalled, Arc::new(TaskContext::default())) + .in_partition_order(usize::MAX), runtime.handle(), ) .unwrap(); @@ -199,7 +200,7 @@ fn backpressured_query_does_not_starve_other_queries_and_cancel_releases_graph() ranges(&[(100, 150), (150, 200), (200, 237)], 3), Arc::new(TaskContext::default()), ) - .in_partition_order(), + .in_partition_order(usize::MAX), runtime.handle(), ) .unwrap(); @@ -276,7 +277,8 @@ fn cancel_releases_the_graph_while_every_worker_is_busy() { let plan = BlockingExec::new(schema(), 3); let refs = plan.refs(); let query = QueryHandle::start( - LocalQuery::new(Arc::new(plan), Arc::new(TaskContext::default())).in_partition_order(), + LocalQuery::new(Arc::new(plan), Arc::new(TaskContext::default())) + .in_partition_order(usize::MAX), runtime.handle(), ) .unwrap(); diff --git a/native/local/tests/execution.rs b/native/local/tests/execution.rs index 1d3c8a6bf20..8fb0eb8aa92 100644 --- a/native/local/tests/execution.rs +++ b/native/local/tests/execution.rs @@ -460,7 +460,7 @@ async fn partition_order_delivers_partitions_in_turn() { let batches: Vec = timeout( DEADLINE, LocalQuery::new(plan, context()) - .in_partition_order() + .in_partition_order(usize::MAX) .execute() .unwrap() .try_collect(), @@ -496,7 +496,7 @@ fn failing() -> Arc { async fn partition_order_reports_a_later_failure_in_its_turn() { let plan = UnionExec::try_new(vec![ordered_input(1, 3), failing()]).unwrap(); let mut stream = LocalQuery::new(plan, context()) - .in_partition_order() + .in_partition_order(usize::MAX) .execute() .unwrap(); // The earlier partition's batches come first, even though the later one failed at once. @@ -522,7 +522,7 @@ async fn partition_order_never_reports_the_failure_of_a_partition_not_reached() let refs = blocking.refs(); let plan = UnionExec::try_new(vec![Arc::new(blocking), failing()]).unwrap(); let mut stream = LocalQuery::new(plan, context()) - .in_partition_order() + .in_partition_order(usize::MAX) .execute() .unwrap(); // A consumer that stops in the first partition, as a limit does, sees no error. @@ -540,7 +540,7 @@ async fn partition_order_cancel_releases_running_partitions() { let refs = blocking.refs(); let (blocking, observations) = observe(Arc::new(blocking)); let mut stream = LocalQuery::new(blocking, context()) - .in_partition_order() + .in_partition_order(usize::MAX) .execute() .unwrap(); wait_started(&observations, 3).await; @@ -549,3 +549,36 @@ async fn partition_order_cancel_releases_running_partitions() { assert!(stream.next().await.is_none()); wait_released(&refs).await; } + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn partition_order_runs_at_most_the_given_number_of_partitions() { + let (plan, observations) = observe(ordered_input(6, 3)); + let executed = || observations.executions.lock().unwrap().clone(); + let mut stream = LocalQuery::new(plan, context()) + .in_partition_order(2) + .execute() + .unwrap(); + // The first two partitions start with the query; the rest wait for their turn. + assert_eq!(executed(), vec![0, 1]); + for _ in 0..3 { + timeout(DEADLINE, stream.next()) + .await + .unwrap() + .unwrap() + .unwrap(); + } + assert_eq!(executed(), vec![0, 1]); + // Finishing the first partition starts the third. + timeout(DEADLINE, stream.next()) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(executed(), vec![0, 1, 2]); + let rest: Vec = timeout(DEADLINE, stream.try_collect()) + .await + .unwrap() + .unwrap(); + assert_eq!(rest.len(), 6 * 3 - 4); + assert_executed_once(&observations, 6); +} diff --git a/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala b/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala index a8dd5366ee3..1ecd7b77e0b 100644 --- a/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala @@ -113,7 +113,7 @@ class CometLocalExecutionSuite spark.sparkContext .runJob( spark.sparkContext.parallelize(Seq(0), 1), - { _: Iterator[Int] => + { (_: Iterator[Int]) => f(TaskContext.get()) }) .head @@ -272,6 +272,19 @@ class CometLocalExecutionSuite } } + test("many partitions with sorts fit the memory pool") { + // Each sort registers its consumers when its partition starts and reserves memory for its + // merge up front. Running only as many partitions at once as Spark runs tasks keeps their + // fair shares of the 256 MiB pool large enough, as with ordinary execution. + def query(): DataFrame = + spark + .range(0, 32000, 1, 32) + .selectExpr("id", "sha2(cast(id as string), 256) AS k") + .sortWithinPartitions("k") + assert(localNodes(query()).size == 1, query().queryExecution.executedPlan.toString) + assert(query().collect().toSeq == expected(query())) + } + test("repeated actions build fresh executions and iterator results") { val query = spark.range(0, 300, 3, 7).selectExpr("id AS value") val rows = (0L until 300L by 3L).map(Row(_)) From b456e6d0368562966c4bd06c5dc1d30a4740c2a7 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Mon, 19 Oct 2026 22:45:54 -0700 Subject: [PATCH 06/12] fix: run takes of local queries as ordinary Comet Spark's executeTake requests partitions in growing batches and runs each batch to completion, so it can evaluate partitions beyond the one that supplies the last row. Reading the local result in order stopped at that row instead, so a take could succeed where Spark fails on a later partition of the same batch. CometLocalResultExec now runs the admitted native block as ordinary Comet execution for executeTake, as it does for executeToIterator and executeTail, so take, head, show and limit(n).collect() evaluate the partitions Spark evaluates. The result handoff now serves only collect. --- .../contributor-guide/local-execution.md | 31 +++++++-------- .../comet/local/CometLocalResultExec.scala | 38 +++++++++---------- .../local/CometLocalExecutionSuite.scala | 32 +++++++++++++--- 3 files changed, 61 insertions(+), 40 deletions(-) diff --git a/docs/source/contributor-guide/local-execution.md b/docs/source/contributor-guide/local-execution.md index c6c9ab94a89..198b0bbec00 100644 --- a/docs/source/contributor-guide/local-execution.md +++ b/docs/source/contributor-guide/local-execution.md @@ -42,7 +42,7 @@ admitted when: - it has between 1 and 1024 partitions. A root `CollectLimitExec`, which Spark plans for `take`, `head` and `show`, is kept above the -local query. Such plans have no exchange, so AQE does not wrap them and they are admitted +local query, and reads it as described under result delivery. Such plans have no exchange, so AQE does not wrap them and they are admitted whether or not AQE is enabled. ## Execution @@ -62,10 +62,8 @@ single partition. That partition's task creates the native graph: order. At most `spark.executor.cores` partitions (N for `local[N]`) run at once, each up to two batches ahead, and the next starts when the first is delivered, as Spark runs at most that many tasks. A partition's operators register their memory consumers when it starts, so this - also bounds how many share the task's memory pool. The result therefore has Spark's row order, and - a limit returns the same rows as Spark's `executeTake`. A partition's error is reported when its - turn comes, so a limit satisfied by earlier partitions does not fail on a later one, as with - `executeTake`, which never evaluates the partitions it does not need. + also bounds how many share the task's memory pool. The result therefore has Spark's row order. + A partition's error is reported when its turn comes. `CometLocalExec` reports `UnknownPartitioning(1)`: the result is one partition, but its rows are not clustered by anything, so a query over a cached local result still gets the exchange @@ -84,21 +82,24 @@ Arrow arrays use `prepare_output` and `NativeUtil`; an imported batch stays vali native query is closed, until the JVM closes it. A pull returns "pending" after 50 ms, so that the task can check whether it was killed. -`CometLocalResultExec` sits above Spark's columnar-to-row transition. For `collect` and `take` -it still runs the single task as a Spark job, which keeps cancellation, job groups and SQL -metrics, but the task copies its rows into a driver-side slot instead of returning an encoded -task result. Spark would otherwise encode and compress every row in that one task and decode +`CometLocalResultExec` sits above Spark's columnar-to-row transition. For `collect` it still +runs the single task as a Spark job, which keeps cancellation, job groups and SQL metrics, but +the task copies its rows into a driver-side slot instead of returning an encoded task result. Spark would otherwise encode and compress every row in that one task and decode them on the driver, all on one thread. While copying, the task counts uncompressed `UnsafeRow` bytes and stops reading, which closes the native query, once they exceed `spark.driver.maxResultSize`; the driver then fails the action. Spark counts compressed bytes, so this check is the stricter one. On a non-local master the node uses Spark's ordinary collect. -An iterator fetches one result partition at a time so that it holds at most one, but a local -query has a single partition holding the whole result. `toLocalIterator`, which also serves the -Thrift server's incremental collect, therefore runs the native block as ordinary Comet execution, -one Spark task per partition. Consumers that read the result RDD partition by partition, such as -PySpark's `toLocalIterator`, still receive the whole result as one partition. +Spark's other actions read a result partition by partition: a take requests partitions in +growing batches until it has enough rows, an iterator fetches one partition at a time so that it +holds at most one, and a tail scans partitions from the last one. A local query has a single +partition holding the whole result, so `take`, `head`, `show`, `limit(n).collect()`, +`toLocalIterator` (which also serves the Thrift server's incremental collect) and PySpark's +`tail` run the native block as ordinary Comet execution, one Spark task per partition. The +partitions they evaluate, and so the rows they return and the errors they raise, are Spark's. +Consumers that read the result RDD partition by partition, such as PySpark's `toLocalIterator`, +still receive the whole result as one partition. ## Lifecycle and cancellation @@ -127,5 +128,5 @@ cd .. The native tests cover partition-order delivery, the bound on running partitions, errors reported in a partition's turn, cancellation with busy workers, the one-batch handoff and runtime shutdown. `CometLocalExecutionSuite` compares admitted queries with Spark's results and order, and covers ranges Spark would overflow, declined expressions, -AQE, reported orderings, limits that stop before a failing partition, iterators, many partitions with sorts in a small memory pool, caching, the result size limit, job and task cancellation, and native errors. After +AQE, reported orderings, takes and limits, iterators, tails, many partitions with sorts in a small memory pool, caching, the result size limit, job and task cancellation, and native errors. After each test it checks that no native query, imported Arrow memory or result slot is left. diff --git a/spark/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala b/spark/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala index 8abd68ff2fe..491f8319e82 100644 --- a/spark/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala +++ b/spark/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala @@ -35,8 +35,8 @@ import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, UnaryExecNo * Row result boundary of a local query. Spark's collect encodes and compresses every row of a * result partition in its task, then decodes them on the driver. A local query has one result * partition, so that work would run on a single thread. Admission requires an in-process local - * master, so collect and take still run one Spark task (keeping cancellation, job groups and SQL - * metrics) but hand copied rows to the driver in this JVM. + * master, so collect still runs one Spark task (keeping cancellation, job groups and SQL metrics) + * but hands copied rows to the driver in this JVM. */ case class CometLocalResultExec(child: SparkPlan) extends UnaryExecNode { override def output: Seq[Attribute] = child.output @@ -46,18 +46,16 @@ case class CometLocalResultExec(child: SparkPlan) extends UnaryExecNode { override protected def doExecute(): RDD[InternalRow] = child.execute() override def executeCollect(): Array[InternalRow] = - if (sparkContext.isLocal) LocalResultHandoff.collect(execute(), -1, maxResultSize) + if (sparkContext.isLocal) LocalResultHandoff.collect(execute(), maxResultSize) else super.executeCollect() - override def executeTake(n: Int): Array[InternalRow] = - if (n <= 0) Array.empty - else if (sparkContext.isLocal) LocalResultHandoff.collect(execute(), n, maxResultSize) - else super.executeTake(n) + // A take requests partitions in growing batches until it has enough rows, an iterator fetches + // one result partition at a time to bound what it holds, and a tail scans partitions from the + // last one. A local query has a single partition holding the whole result, so these run the + // native block as ordinary Comet execution instead, one Spark task per partition, as without + // local execution. Which partitions they evaluate, and so which errors they raise, match Spark. + override def executeTake(n: Int): Array[InternalRow] = ordinary.executeTake(n) - // An iterator fetches one result partition at a time to bound what it holds, and a tail scans - // partitions from the last one until it has enough rows. A local query has a single partition - // holding the whole result, so both run the native block as ordinary Comet execution instead, - // one Spark task per partition, as without local execution. override def executeToIterator(): Iterator[InternalRow] = ordinary.executeToIterator() override def executeTail(n: Int): Array[InternalRow] = ordinary.executeTail(n) @@ -85,22 +83,22 @@ private[local] object LocalResultHandoff { private val slots = new ConcurrentHashMap[Long, Slot]() private val nextId = new AtomicLong() - /** Number of live slots; zero whenever no collect or take is running. */ + /** Number of live slots; zero whenever no collect is running. */ def active: Int = slots.size() /** - * Runs the single result partition of `rdd` and returns at most `take` rows (all if negative). - * Stops reading, which closes the native query, once the copied rows exceed `maxResultSize` - * bytes (unlimited if not positive). Unlike Spark's compressed size, this counts uncompressed - * UnsafeRow bytes, so it is the more conservative of the two. + * Runs the single result partition of `rdd` and returns its rows. Stops reading, which closes + * the native query, once the copied rows exceed `maxResultSize` bytes (unlimited if not + * positive). Unlike Spark's compressed size, this counts uncompressed UnsafeRow bytes, so it is + * the more conservative of the two. */ - def collect(rdd: RDD[InternalRow], take: Int, maxResultSize: Long): Array[InternalRow] = { + def collect(rdd: RDD[InternalRow], maxResultSize: Long): Array[InternalRow] = { require(rdd.getNumPartitions == 1, "Local results have exactly one partition") val id = nextId.incrementAndGet() val slot = new Slot(maxResultSize) slots.put(id, slot) try { - rdd.sparkContext.runJob(rdd, (rows: Iterator[InternalRow]) => fill(id, rows, take), Seq(0)) + rdd.sparkContext.runJob(rdd, (rows: Iterator[InternalRow]) => fill(id, rows), Seq(0)) if (slot.exceededBytes >= 0) { throw new SparkException( s"Total size of local query results (at least ${slot.exceededBytes} bytes, " + @@ -115,12 +113,12 @@ private[local] object LocalResultHandoff { // Runs in the result task. A retried attempt rebuilds the result and publishes it only once // complete. A slot the driver already removed means nobody will read the rows. - private def fill(id: Long, rows: Iterator[InternalRow], take: Int): Unit = { + private def fill(id: Long, rows: Iterator[InternalRow]): Unit = { val slot = slots.get(id) if (slot == null) return val buffer = new ArrayBuffer[InternalRow]() var bytes = 0L - while ((take < 0 || buffer.length < take) && rows.hasNext) { + while (rows.hasNext) { // ColumnarToRow reuses one UnsafeRow; Spark's collect makes the same assumption. val row = rows.next().asInstanceOf[UnsafeRow].copy() bytes += row.getSizeInBytes diff --git a/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala b/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala index 1ecd7b77e0b..4a1c440acb9 100644 --- a/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala @@ -329,7 +329,7 @@ class CometLocalExecutionSuite } } - test("collect and take hand local results to the driver within a Spark job") { + test("collect hands local results to the driver, and take runs as ordinary execution") { val query = spark.range(0, 200, 1, 5).selectExpr("id", "id % 3 AS k") assert(query.queryExecution.executedPlan.isInstanceOf[CometLocalResultExec]) assert( @@ -338,6 +338,26 @@ class CometLocalExecutionSuite assert(query.take(4).map(_.getLong(0)).toSeq == (0L until 4L)) } + test("a take evaluates the partitions Spark requests together") { + withConf("spark.sql.ansi.enabled", "true") { + // After the first partition supplies one row, Spark requests the second and third + // together, and the third divides by zero at id 4, although the second has the missing row. + def query(): DataFrame = + spark + .range(0, 6, 1, 3) + .filter("id % 2 = 0") + .selectExpr("100L DIV (4L - id) AS v") + .limit(2) + assert(localNodes(query()).size == 1, query().queryExecution.executedPlan.toString) + val failure = intercept[Exception](query().collect()) + assert(failure.getMessage.contains("DIVIDE_BY_ZERO"), failure.getMessage) + val sparkFailure = withConf("spark.comet.enabled", "false") { + intercept[Exception](query().collect()) + } + assert(sparkFailure.getMessage.contains("DIVIDE_BY_ZERO"), sparkFailure.getMessage) + } + } + test("job group cancellation interrupts a local collect") { val query = spark.range(0, Long.MaxValue, 1, 7).toDF() assert(query.queryExecution.executedPlan.isInstanceOf[CometLocalResultExec]) @@ -370,15 +390,17 @@ class CometLocalExecutionSuite assert(plan.isInstanceOf[CometLocalResultExec], plan.toString) // Each two-long UnsafeRow is 24 bytes, so 1,000 bytes fit 41 rows and the 42nd exceeds it. val failure = intercept[SparkException] { - LocalResultHandoff.collect(plan.execute(), -1, 1000) + LocalResultHandoff.collect(plan.execute(), 1000) } assert(failure.getMessage.contains("spark.driver.maxResultSize"), failure.getMessage) assert(failure.getMessage.contains("at least 1008 bytes"), failure.getMessage) assert(new NativeLocal().activeQueries() == 0) - // A take within the limit still succeeds against the same bound. - assert(LocalResultHandoff.collect(plan.execute(), 41, 1000).length == 41) + // A result within the limit still succeeds against the same bound. + def bounded(rows: Long) = + spark.range(0, rows, 1, 7).selectExpr("id", "id AS other").queryExecution.executedPlan + assert(LocalResultHandoff.collect(bounded(41).execute(), 1000).length == 41) // Non-positive limits mean unlimited, as in Spark. - assert(LocalResultHandoff.collect(plan.execute(), 100, 0).length == 100) + assert(LocalResultHandoff.collect(bounded(100).execute(), 0).length == 100) } test("Spark job cancellation releases a live local execution") { From 2bc36af1d54291eb639043bdc6b28b0ac06c2000 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Wed, 21 Oct 2026 09:12:17 -0700 Subject: [PATCH 07/12] fix: keep memory_pools private to avoid a conflict with upstream Re-export create_memory_pool from jni_api for local execution instead of making the memory_pools module crate-visible. --- native/core/src/execution/jni_api.rs | 4 +++- native/core/src/execution/mod.rs | 2 +- native/core/src/local.rs | 5 ++--- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index 94abb636e8e..7c11f8ff422 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -117,8 +117,10 @@ use tokio::runtime::{Handle, Runtime}; use tokio::sync::mpsc; use tokio::task::JoinHandle; +// Re-exported for local execution, whose memory pool is the one `createPlan` creates. +pub(crate) use crate::execution::memory_pools::create_memory_pool; use crate::execution::memory_pools::{ - create_memory_pool, overcommit, parse_memory_pool_config, MemoryPoolConfig, PlanMemoryPool, + overcommit, parse_memory_pool_config, MemoryPoolConfig, PlanMemoryPool, }; use crate::execution::operators::{ScanExec, ShuffleScanExec}; use crate::execution::shuffle::{ diff --git a/native/core/src/execution/mod.rs b/native/core/src/execution/mod.rs index bdea0f9825c..9a2e54783b4 100644 --- a/native/core/src/execution/mod.rs +++ b/native/core/src/execution/mod.rs @@ -26,7 +26,7 @@ pub mod operators; pub(crate) mod planner; pub mod serde; pub use datafusion_comet_shuffle as shuffle; -pub(crate) mod memory_pools; +mod memory_pools; pub(crate) mod sort; pub(crate) mod spark_aggregate_state; pub(crate) mod spark_config; diff --git a/native/core/src/local.rs b/native/core/src/local.rs index 784a392b84d..a99eb0bc445 100644 --- a/native/core/src/local.rs +++ b/native/core/src/local.rs @@ -37,10 +37,9 @@ use parking_lot::Mutex; use crate::errors::{try_unwrap_or_throw, CometError, CometResult}; use crate::execution::jni_api::{ - create_execution_context, get_runtime, prepare_output, prepare_plan_settings, ExecutionContext, - PlanInputs, PlanJvmRefs, + create_execution_context, create_memory_pool, get_runtime, prepare_output, + prepare_plan_settings, ExecutionContext, PlanInputs, PlanJvmRefs, }; -use crate::execution::memory_pools::create_memory_pool; use crate::execution::spark_config::{SparkConfig, SPARK_EXECUTOR_CORES}; struct Entry { From b55a395f9db95b9c270dd0038c539241669266cc Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Fri, 23 Oct 2026 10:00:43 -0700 Subject: [PATCH 08/12] fix: let running local partitions read ahead up to maxResultSize Partitions waiting for their turn stopped after two batches, so for a streaming block about one partition made progress at a time. They now run ahead while what they hold totals less than spark.driver.maxResultSize, and hold at most two batches each beyond that. --- .../contributor-guide/local-execution.md | 16 ++- native/core/src/local.rs | 5 +- native/local/src/lib.rs | 121 +++++++++++++----- native/local/tests/bridge.rs | 9 +- native/local/tests/execution.rs | 58 ++++++++- .../comet/local/LocalQueryIterator.scala | 5 +- .../org/apache/comet/local/NativeLocal.scala | 6 +- .../local/CometLocalExecutionSuite.scala | 3 +- 8 files changed, 170 insertions(+), 53 deletions(-) diff --git a/docs/source/contributor-guide/local-execution.md b/docs/source/contributor-guide/local-execution.md index 198b0bbec00..acaf2f4c421 100644 --- a/docs/source/contributor-guide/local-execution.md +++ b/docs/source/contributor-guide/local-execution.md @@ -59,11 +59,15 @@ single partition. That partition's task creates the native graph: for partition `i`, so per-partition state such as a range slice is the same as with ordinary Comet execution. - `native/local` runs the partitions on Comet's Tokio runtime and delivers them in partition - order. At most `spark.executor.cores` partitions (N for `local[N]`) run at once, each up to two - batches ahead, and the next starts when the first is delivered, as Spark runs at most that - many tasks. A partition's operators register their memory consumers when it starts, so this - also bounds how many share the task's memory pool. The result therefore has Spark's row order. - A partition's error is reported when its turn comes. + order. At most `spark.executor.cores` partitions (N for `local[N]`) run at once, and the next + starts when the first is delivered, as Spark runs at most that many tasks. A partition's + operators register their memory consumers when it starts, so this also bounds how many share + the task's memory pool. The result therefore has Spark's row order. A partition's error is + reported when its turn comes. +- The running partitions do not wait for their turn to produce batches, so they all make + progress, as Spark's tasks do. What they hold ahead of delivery is bounded by + `spark.driver.maxResultSize`, the most a collect may bring to the driver; past it, each holds + at most two batches until the earlier partitions are delivered. `CometLocalExec` reports `UnknownPartitioning(1)`: the result is one partition, but its rows are not clustered by anything, so a query over a cached local result still gets the exchange @@ -125,7 +129,7 @@ cd .. ./mvnw test -Dsuites=org.apache.comet.local.CometLocalExecutionSuite ``` -The native tests cover partition-order delivery, the bound on running partitions, errors reported in a partition's turn, cancellation with busy workers, +The native tests cover partition-order delivery, the bound on running partitions, read-ahead, errors reported in a partition's turn, cancellation with busy workers, the one-batch handoff and runtime shutdown. `CometLocalExecutionSuite` compares admitted queries with Spark's results and order, and covers ranges Spark would overflow, declined expressions, AQE, reported orderings, takes and limits, iterators, tails, many partitions with sorts in a small memory pool, caching, the result size limit, job and task cancellation, and native errors. After diff --git a/native/core/src/local.rs b/native/core/src/local.rs index a99eb0bc445..2cb74b80988 100644 --- a/native/core/src/local.rs +++ b/native/core/src/local.rs @@ -83,6 +83,7 @@ pub extern "system" fn Java_org_apache_comet_local_NativeLocal_createQuery( memory_limit: jlong, task_attempt_id: jlong, task_cpus: jlong, + read_ahead_bytes: jlong, ) -> jlong { try_unwrap_or_throw(&e, |env| { if !(1..=MAX_PARTITIONS).contains(&partitions) || columns < 1 { @@ -132,8 +133,8 @@ pub extern "system" fn Java_org_apache_comet_local_NativeLocal_createQuery( "Local output schema width mismatch".into(), )); } - let query = - LocalQuery::new(plan, context.session_ctx.task_ctx()).in_partition_order(max_running); + let query = LocalQuery::new(plan, context.session_ctx.task_ctx()) + .in_partition_order(max_running, read_ahead_bytes.max(0) as usize); let query = QueryHandle::start(query, &get_runtime())?; register( id, diff --git a/native/local/src/lib.rs b/native/local/src/lib.rs index 8d465b73155..3efc81a114c 100644 --- a/native/local/src/lib.rs +++ b/native/local/src/lib.rs @@ -26,7 +26,7 @@ pub mod handle; use std::collections::VecDeque; use std::panic::AssertUnwindSafe; use std::pin::Pin; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, MutexGuard}; use std::task::{Context, Poll}; @@ -41,18 +41,20 @@ use datafusion::physical_plan::{ }; use futures::future::poll_fn; use futures::{FutureExt, Stream, StreamExt}; +use tokio::sync::Notify; use tokio::task::JoinHandle; -/// Batches each root partition may run ahead of the one being delivered, when results are -/// delivered in partition order. +/// Batches each root partition may always hold ahead of delivery, when results are delivered in +/// partition order, whatever the read-ahead limit. const PARTITION_BUFFER: usize = 2; /// One execution's graph and resources. Consumed when execution starts; not a plan cache. pub struct LocalQuery { plan: Arc, context: Arc, - /// Delivers the partitions in order, with at most this many running at once. - in_partition_order: Option, + /// Delivers the partitions in order, with at most this many running at once and this many + /// bytes read ahead. + in_partition_order: Option<(usize, usize)>, } impl LocalQuery { @@ -69,14 +71,17 @@ impl LocalQuery { } /// Delivers the root partitions one after another, in partition order, as Spark returns - /// the partitions of a result. Up to `max_running` partitions run at once, each up to a few - /// batches ahead, and the next one starts when the first is delivered, as Spark runs at most - /// as many tasks as it has cores. A partition's operators register their memory consumers - /// when it starts, so this also bounds how many share the memory pool. Only for roots whose - /// partitions do not share an exchange: a partition that waits to be delivered would - /// otherwise leave its exchange buffering for the others. - pub fn in_partition_order(mut self, max_running: usize) -> Self { - self.in_partition_order = Some(max_running.max(1)); + /// the partitions of a result. Up to `max_running` partitions run at once, and the next one + /// starts when the first is delivered, as Spark runs at most as many tasks as it has cores. + /// A partition's operators register their memory consumers when it starts, so this also + /// bounds how many share the memory pool. + /// + /// The running partitions do not wait for their turn: each runs ahead of delivery while the + /// batches they all hold total less than `read_ahead` bytes, and up to a few batches beyond + /// that. Only for roots whose partitions do not share an exchange: a partition that waits to + /// be delivered would otherwise leave its exchange buffering for the others. + pub fn in_partition_order(mut self, max_running: usize, read_ahead: usize) -> Self { + self.in_partition_order = Some((max_running.max(1), read_ahead)); self } @@ -90,8 +95,8 @@ impl LocalQuery { DataFusionError::Execution("LocalQuery requires an active Tokio runtime".into()) })?; let schema = self.plan.schema(); - let stream = if let Some(max_running) = self.in_partition_order { - in_partition_order(&self.plan, &self.context, max_running)? + let stream = if let Some((max_running, read_ahead)) = self.in_partition_order { + in_partition_order(&self.plan, &self.context, max_running, read_ahead)? } else { execute_stream(Arc::clone(&self.plan), Arc::clone(&self.context))? }; @@ -107,6 +112,7 @@ fn in_partition_order( plan: &Arc, context: &Arc, max_running: usize, + read_ahead: usize, ) -> Result { let mut order = PartitionOrder { plan: Arc::clone(plan), @@ -114,6 +120,11 @@ fn in_partition_order( runtime: tokio::runtime::Handle::current(), next: 0, max_running, + read_ahead: Arc::new(ReadAhead { + limit: read_ahead, + held: AtomicUsize::new(0), + delivered: Notify::new(), + }), running: VecDeque::new(), }; // Start the first partitions now, so that a failure to start them fails here. @@ -124,9 +135,22 @@ fn in_partition_order( ))) } -/// One root partition running on its own Tokio task, a few batches ahead of delivery. +/// What the running partitions hold ahead of delivery, shared by them. +struct ReadAhead { + /// Bytes the partitions may hold before each waits at `PARTITION_BUFFER` batches. + limit: usize, + /// Bytes of the batches sent and not yet delivered. + held: AtomicUsize, + /// Notified whenever a batch is delivered. + delivered: Notify, +} + +/// One root partition running on its own Tokio task, ahead of delivery. struct Partition { - batches: tokio::sync::mpsc::Receiver>, + /// Each batch with its size in bytes. + batches: tokio::sync::mpsc::UnboundedReceiver>, + /// Batches sent and not yet delivered. + queued: Arc, /// Set once the task has sent the stream's last batch. The channel closes whenever the /// task ends, and a runtime that shuts down cancels its tasks, so only this tells that /// every batch arrived. @@ -138,23 +162,51 @@ struct Partition { } impl Partition { - fn spawn(runtime: &tokio::runtime::Handle, stream: SendableRecordBatchStream) -> Self { - let (sender, batches) = tokio::sync::mpsc::channel(PARTITION_BUFFER); + fn spawn( + runtime: &tokio::runtime::Handle, + stream: SendableRecordBatchStream, + read_ahead: &Arc, + ) -> Self { + let (sender, batches) = tokio::sync::mpsc::unbounded_channel(); + let queued = Arc::new(AtomicUsize::new(0)); let ended = Arc::new(AtomicBool::new(false)); let stream = Arc::new(Mutex::new(Some(stream))); - let (task_ended, polled) = (Arc::clone(&ended), Arc::clone(&stream)); + let (task_queued, task_ended, polled) = + (Arc::clone(&queued), Arc::clone(&ended), Arc::clone(&stream)); + let read_ahead = Arc::clone(read_ahead); let task = runtime.spawn(async move { let result = AssertUnwindSafe(async { - // Reserve before polling, so at most `PARTITION_BUFFER` batches are held. - while let Ok(permit) = sender.reserve().await { + loop { + // Wait before polling while this partition holds `PARTITION_BUFFER` batches + // and the partitions together hold the read-ahead limit. The first partition + // is being delivered, so it never waits for long. + loop { + let delivered = read_ahead.delivered.notified(); + if task_queued.load(Ordering::Acquire) < PARTITION_BUFFER + || read_ahead.held.load(Ordering::Acquire) < read_ahead.limit + { + break; + } + delivered.await; + } let next = poll_fn(|cx| match lock(&polled).as_mut() { Some(stream) => stream.poll_next_unpin(cx), None => Poll::Ready(None), }) .await; match next { - Some(Ok(batch)) => permit.send(Ok(batch)), - Some(Err(error)) => return permit.send(Err(error)), + Some(Ok(batch)) => { + let size = batch.get_array_memory_size(); + task_queued.fetch_add(1, Ordering::AcqRel); + read_ahead.held.fetch_add(size, Ordering::AcqRel); + if sender.send(Ok((batch, size))).is_err() { + return; + } + } + Some(Err(error)) => { + let _ = sender.send(Err(error)); + return; + } None => { // A taken stream also ends here, but then nobody reads the flag. task_ended.store(true, Ordering::Release); @@ -166,15 +218,14 @@ impl Partition { .catch_unwind() .await; if result.is_err() { - let _ = sender - .send(Err(DataFusionError::Execution( - "Local query partition panicked".into(), - ))) - .await; + let _ = sender.send(Err(DataFusionError::Execution( + "Local query partition panicked".into(), + ))); } }); Self { batches, + queued, ended, stream, task, @@ -202,6 +253,7 @@ struct PartitionOrder { /// The first partition not started yet. next: usize, max_running: usize, + read_ahead: Arc, /// Started partitions, the first of which is being delivered. running: VecDeque, } @@ -211,8 +263,8 @@ impl PartitionOrder { let partitions = self.plan.output_partitioning().partition_count(); while self.running.len() < self.max_running && self.next < partitions { let stream = self.plan.execute(self.next, Arc::clone(&self.context))?; - self.running - .push_back(Partition::spawn(&self.runtime, stream)); + let partition = Partition::spawn(&self.runtime, stream, &self.read_ahead); + self.running.push_back(partition); self.next += 1; } Ok(()) @@ -237,7 +289,12 @@ impl Stream for PartitionOrder { return Poll::Ready(None); }; match partition.batches.poll_recv(cx) { - Poll::Ready(Some(Ok(batch))) => return Poll::Ready(Some(Ok(batch))), + Poll::Ready(Some(Ok((batch, size)))) => { + partition.queued.fetch_sub(1, Ordering::AcqRel); + self.read_ahead.held.fetch_sub(size, Ordering::AcqRel); + self.read_ahead.delivered.notify_waiters(); + return Poll::Ready(Some(Ok(batch))); + } Poll::Ready(Some(Err(error))) => return self.fail(error), Poll::Ready(None) if partition.ended.load(Ordering::Acquire) => { self.running.pop_front(); diff --git a/native/local/tests/bridge.rs b/native/local/tests/bridge.rs index c4402f139b7..2d10a30ff6b 100644 --- a/native/local/tests/bridge.rs +++ b/native/local/tests/bridge.rs @@ -78,7 +78,8 @@ fn partition_order_is_preserved_through_the_handoff() { let plan = ranges(&[(0, 10), (10, 10), (10, 25), (25, 26)], 3); let refs = Arc::downgrade(&plan); let query = QueryHandle::start( - LocalQuery::new(plan, Arc::new(TaskContext::default())).in_partition_order(usize::MAX), + LocalQuery::new(plan, Arc::new(TaskContext::default())) + .in_partition_order(usize::MAX, usize::MAX), runtime.handle(), ) .unwrap(); @@ -190,7 +191,7 @@ fn backpressured_query_does_not_starve_other_queries_and_cancel_releases_graph() let refs = Arc::downgrade(&stalled); let stalled = QueryHandle::start( LocalQuery::new(stalled, Arc::new(TaskContext::default())) - .in_partition_order(usize::MAX), + .in_partition_order(usize::MAX, usize::MAX), runtime.handle(), ) .unwrap(); @@ -200,7 +201,7 @@ fn backpressured_query_does_not_starve_other_queries_and_cancel_releases_graph() ranges(&[(100, 150), (150, 200), (200, 237)], 3), Arc::new(TaskContext::default()), ) - .in_partition_order(usize::MAX), + .in_partition_order(usize::MAX, usize::MAX), runtime.handle(), ) .unwrap(); @@ -278,7 +279,7 @@ fn cancel_releases_the_graph_while_every_worker_is_busy() { let refs = plan.refs(); let query = QueryHandle::start( LocalQuery::new(Arc::new(plan), Arc::new(TaskContext::default())) - .in_partition_order(usize::MAX), + .in_partition_order(usize::MAX, usize::MAX), runtime.handle(), ) .unwrap(); diff --git a/native/local/tests/execution.rs b/native/local/tests/execution.rs index 8fb0eb8aa92..de45504f8c7 100644 --- a/native/local/tests/execution.rs +++ b/native/local/tests/execution.rs @@ -96,6 +96,7 @@ fn repartition( #[derive(Debug, Default)] struct Observations { executions: Mutex>, + batches: Mutex>, key_partitions: Mutex>>, } @@ -167,6 +168,12 @@ impl ExecutionPlan for ObserveExec { .as_any() .downcast_ref::() .unwrap(); + *observations + .batches + .lock() + .unwrap() + .entry(partition) + .or_default() += 1; let mut locations = observations.key_partitions.lock().unwrap(); for key in keys.values() { locations.entry(*key).or_default().insert(partition); @@ -460,7 +467,7 @@ async fn partition_order_delivers_partitions_in_turn() { let batches: Vec = timeout( DEADLINE, LocalQuery::new(plan, context()) - .in_partition_order(usize::MAX) + .in_partition_order(usize::MAX, usize::MAX) .execute() .unwrap() .try_collect(), @@ -496,7 +503,7 @@ fn failing() -> Arc { async fn partition_order_reports_a_later_failure_in_its_turn() { let plan = UnionExec::try_new(vec![ordered_input(1, 3), failing()]).unwrap(); let mut stream = LocalQuery::new(plan, context()) - .in_partition_order(usize::MAX) + .in_partition_order(usize::MAX, usize::MAX) .execute() .unwrap(); // The earlier partition's batches come first, even though the later one failed at once. @@ -522,7 +529,7 @@ async fn partition_order_never_reports_the_failure_of_a_partition_not_reached() let refs = blocking.refs(); let plan = UnionExec::try_new(vec![Arc::new(blocking), failing()]).unwrap(); let mut stream = LocalQuery::new(plan, context()) - .in_partition_order(usize::MAX) + .in_partition_order(usize::MAX, usize::MAX) .execute() .unwrap(); // A consumer that stops in the first partition, as a limit does, sees no error. @@ -540,7 +547,7 @@ async fn partition_order_cancel_releases_running_partitions() { let refs = blocking.refs(); let (blocking, observations) = observe(Arc::new(blocking)); let mut stream = LocalQuery::new(blocking, context()) - .in_partition_order(usize::MAX) + .in_partition_order(usize::MAX, usize::MAX) .execute() .unwrap(); wait_started(&observations, 3).await; @@ -555,7 +562,7 @@ async fn partition_order_runs_at_most_the_given_number_of_partitions() { let (plan, observations) = observe(ordered_input(6, 3)); let executed = || observations.executions.lock().unwrap().clone(); let mut stream = LocalQuery::new(plan, context()) - .in_partition_order(2) + .in_partition_order(2, usize::MAX) .execute() .unwrap(); // The first two partitions start with the query; the rest wait for their turn. @@ -582,3 +589,44 @@ async fn partition_order_runs_at_most_the_given_number_of_partitions() { assert_eq!(rest.len(), 6 * 3 - 4); assert_executed_once(&observations, 6); } + +/// A first partition that never delivers, and a second one of ten batches. +fn behind_a_blocked_partition() -> (Arc, Arc) { + let blocking = Arc::new(BlockingExec::new(schema(), 1)); + observe(UnionExec::try_new(vec![blocking, ordered_input(1, 10)]).unwrap()) +} + +fn produced(observations: &Observations, partition: usize) -> usize { + let batches = observations.batches.lock().unwrap(); + batches.get(&partition).copied().unwrap_or(0) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn partition_order_runs_later_partitions_ahead_within_the_limit() { + let (plan, observations) = behind_a_blocked_partition(); + let mut stream = LocalQuery::new(plan, context()) + .in_partition_order(2, usize::MAX) + .execute() + .unwrap(); + // The second partition runs to its end while the first has delivered nothing. + timeout(DEADLINE, async { + while produced(&observations, 1) < 10 { + tokio::task::yield_now().await; + } + }) + .await + .expect("the second partition did not run ahead"); + stream.cancel(); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn partition_order_holds_a_few_batches_beyond_the_limit() { + let (plan, observations) = behind_a_blocked_partition(); + let mut stream = LocalQuery::new(plan, context()) + .in_partition_order(2, 0) + .execute() + .unwrap(); + tokio::time::sleep(Duration::from_millis(200)).await; + assert_eq!(produced(&observations, 1), 2); + stream.cancel(); +} diff --git a/spark/src/main/scala/org/apache/comet/local/LocalQueryIterator.scala b/spark/src/main/scala/org/apache/comet/local/LocalQueryIterator.scala index 074218a809a..502f7e5b99d 100644 --- a/spark/src/main/scala/org/apache/comet/local/LocalQueryIterator.scala +++ b/spark/src/main/scala/org/apache/comet/local/LocalQueryIterator.scala @@ -47,6 +47,8 @@ private[local] class LocalQueryIterator(spec: LocalQuerySpec, context: TaskConte checkCancellation() // The memory pool, spill directories and session settings of any Comet plan in this task. val memory = CometExecIterator.getMemoryConfig(SparkEnv.get.conf) + // Partitions may run ahead holding as much of the result as the driver may collect. + val maxResultSize = SparkEnv.get.conf.getSizeAsBytes("spark.driver.maxResultSize", "1g") id = native.createQuery( spec.plan, CometExecIterator.serializeCometSQLConfs(), @@ -59,7 +61,8 @@ private[local] class LocalQueryIterator(spec: LocalQuerySpec, context: TaskConte memory.memoryPoolType, memory.memoryLimit, context.taskAttemptId(), - context.cpus().toLong) + context.cpus().toLong, + if (maxResultSize > 0) maxResultSize else Long.MaxValue) } catch { case failure: Throwable => throw executionFailure(failure) diff --git a/spark/src/main/scala/org/apache/comet/local/NativeLocal.scala b/spark/src/main/scala/org/apache/comet/local/NativeLocal.scala index 6444cc97046..a7985a58646 100644 --- a/spark/src/main/scala/org/apache/comet/local/NativeLocal.scala +++ b/spark/src/main/scala/org/apache/comet/local/NativeLocal.scala @@ -31,7 +31,8 @@ private[local] class NativeLocal extends NativeBase { /** * Creates and starts a local query over `partitions` partitions of a serialized Comet plan * without JVM inputs, with the memory pool and session that `Native.createPlan` creates for the - * calling Spark task. Returns its handle. + * calling Spark task. Its running partitions read ahead of delivery up to `readAheadBytes`. + * Returns its handle. */ // scalastyle:off argcount @native def createQuery( @@ -46,7 +47,8 @@ private[local] class NativeLocal extends NativeBase { memoryPoolType: String, memoryLimit: Long, taskAttemptId: Long, - taskCpus: Long): Long + taskCpus: Long, + readAheadBytes: Long): Long // scalastyle:on argcount /** Exports the next batch; returns its row count, -1 at EOF, or -2 while none is ready. */ diff --git a/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala b/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala index 4a1c440acb9..e0e9efd83a8 100644 --- a/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala @@ -488,7 +488,8 @@ object CometLocalExecutionSuite { memory.memoryPoolType, memory.memoryLimit, context.taskAttemptId(), - 1L) + 1L, + Long.MaxValue) } private def fails(f: => Any): Boolean = From 227e263ec90fcc2b8745d9266884a4123b74a041 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Thu, 29 Oct 2026 15:48:35 -0700 Subject: [PATCH 09/12] fix: run everything but collect of a local query as ordinary execution CometLocalResultExec now executes as the ordinary native block, with the block's partitioning and ordering, so a cache, df.rdd, foreachPartition, localCheckpoint and PySpark's collect keep the block's partitions. Only collect runs the whole query as one graph. Document that read-ahead batches are not reserved in the memory pool and how non-spillable consumers share the task's pool, and describe the local result path in the FFI guide. --- docs/source/contributor-guide/ffi.md | 10 +++- .../contributor-guide/local-execution.md | 48 +++++++++++-------- .../contributor-guide/memory_management.md | 3 ++ .../apache/comet/local/CometLocalExec.scala | 3 +- .../comet/local/CometLocalResultExec.scala | 35 ++++++-------- .../local/CometLocalExecutionSuite.scala | 27 +++++++---- 6 files changed, 75 insertions(+), 51 deletions(-) diff --git a/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 395b5c45909..bf6fce49c2d 100644 --- a/docs/source/contributor-guide/ffi.md +++ b/docs/source/contributor-guide/ffi.md @@ -26,7 +26,8 @@ Comet transfers Arrow data across the JVM/native boundary in two directions: 1. **JVM → Native**: Native code pulls batches from the JVM over the [Arrow C Stream Interface](https://arrow.apache.org/docs/format/CStreamInterface.html). The JVM exports each per-partition iterator once as an `ArrowArrayStream`, and native pulls every batch through a single C callback. -2. **Native → JVM**: JVM pulls batches from native code using `CometExecIterator`, via the +2. **Native → JVM**: JVM pulls batches from native code using `CometExecIterator` (or `LocalQueryIterator` for + [local execution](local-execution.md)), via the [Arrow C Data Interface](https://arrow.apache.org/docs/format/CDataInterface.html) (one `ArrowArray`/`ArrowSchema` pair per column of each batch). @@ -229,6 +230,11 @@ When the JVM needs results from native execution: `SchemaImporter`. Arrow Java's own `Data.importField` creates a new `SchemaImporter` for each field, and each one numbers dictionaries from 0, so two dictionary-encoded columns would collide in the shared `CDataDictionaryProvider`. +A [local execution](local-execution.md) query takes the same path through `LocalQueryIterator`, which calls +`NativeLocal.nextBatch` in place of `Native.executePlan`. `nextBatch` exports the next batch of the query with the same +`prepare_output`, and so with `move_to_spark`, and returns `-2` while none is ready yet, so that the task can check +whether it was killed between pulls. + ### Array Offsets Arrow Java's C Data import ignores `ArrowArray.offset` at every level @@ -262,6 +268,8 @@ Native allocates the data, and the JVM references it without copying: - `CometExecIterator` closes the batch it returned when the consumer next calls `hasNext` or `next`, and on `close()`. A batch is therefore valid only until the consumer asks for the next one, and a consumer that keeps the data longer has to copy it. +- `LocalQueryIterator` closes its batches the same way. A batch it returned stays valid after the native query is + closed, until the JVM closes the batch. By the time the JVM receives a batch, native has usually stopped reserving it in the memory pool, but it stays resident until the JVM closes it. See [Crossing the FFI boundary](memory_management.md#crossing-the-ffi-boundary). diff --git a/docs/source/contributor-guide/local-execution.md b/docs/source/contributor-guide/local-execution.md index acaf2f4c421..ba5da78da70 100644 --- a/docs/source/contributor-guide/local-execution.md +++ b/docs/source/contributor-guide/local-execution.md @@ -20,8 +20,9 @@ under the License. # Local Execution Local execution is an experimental mode for Spark applications that run with a local master. -Instead of running a query as Spark tasks that each create and run a native plan, it runs the -whole query as one DataFusion graph inside a single Spark task, for every partition at once. +Instead of running a query as Spark tasks that each create and run a native plan, a `collect` +runs the whole query as one DataFusion graph inside a single Spark task, for every partition at +once. Other actions run the query as ordinary Comet execution; see result delivery. The internal `spark.comet.exec.local.enabled` option, false by default, turns it on. A query is admitted only as a whole. Every other query uses ordinary Comet execution. @@ -67,14 +68,22 @@ single partition. That partition's task creates the native graph: - The running partitions do not wait for their turn to produce batches, so they all make progress, as Spark's tasks do. What they hold ahead of delivery is bounded by `spark.driver.maxResultSize`, the most a collect may bring to the driver; past it, each holds - at most two batches until the earlier partitions are delivered. - -`CometLocalExec` reports `UnknownPartitioning(1)`: the result is one partition, but its rows -are not clustered by anything, so a query over a cached local result still gets the exchange -its distribution requires. Spark's ordering of the block holds within each partition, so the -result reports it only when it holds across the concatenated partitions: with one partition, or -when the partitions are ranges of a prefix of that ordering, as a range's slices are. Sorting -within partitions, for example, reports no ordering. + at most two batches until the earlier partitions are delivered. These batches are not reserved + in the memory pool: a reservation could not be given back when an operator later needs the + memory, so it would fail queries that ordinary execution runs. Spark's own collect also holds + its task results outside the execution pool, on the driver. +- All the partitions' memory consumers share the one task's pool, where ordinary execution gives + each task its own share. Consumers that can spill each get a share of the task's pool, as they + would of a task's. Consumers that cannot spill, such as a sort's merge, are bounded together by + the pool, as the tasks' consumers are together, but no longer each by one task's share: one + partition can take more than a task could, leaving less for the others. + +`CometLocalExec` reports `UnknownPartitioning(1)`: its RDD is one partition, but its rows are +not clustered by anything. Spark's ordering of the block holds within each partition, so +`CometLocalExec` reports it only when it holds across the concatenated partitions: with one +partition, or when the partitions are ranges of a prefix of that ordering, as a range's slices +are. Sorting within partitions, for example, reports no ordering. `CometLocalResultExec` reports +the block's own partitioning and ordering, which hold for the partitions it returns. Explain shows the native block under `CometLocalExec`. Its native operators do not report SQL metrics; `CometLocalExec` counts the rows and batches it delivers. @@ -95,15 +104,14 @@ bytes and stops reading, which closes the native query, once they exceed so this check is the stricter one. On a non-local master the node uses Spark's ordinary collect. -Spark's other actions read a result partition by partition: a take requests partitions in -growing batches until it has enough rows, an iterator fetches one partition at a time so that it -holds at most one, and a tail scans partitions from the last one. A local query has a single -partition holding the whole result, so `take`, `head`, `show`, `limit(n).collect()`, -`toLocalIterator` (which also serves the Thrift server's incremental collect) and PySpark's -`tail` run the native block as ordinary Comet execution, one Spark task per partition. The -partitions they evaluate, and so the rows they return and the errors they raise, are Spark's. -Consumers that read the result RDD partition by partition, such as PySpark's `toLocalIterator`, -still receive the whole result as one partition. +Every other action reads the result as an RDD of the block's partitions: a take requests +partitions in growing batches until it has enough rows, an iterator fetches one partition at a +time so that it holds at most one, a tail scans partitions from the last one, and a cache, +`df.rdd`, `foreachPartition`, `localCheckpoint` or PySpark's collect and `toLocalIterator` keep +the partitions for later reads. So `CometLocalResultExec` executes as the ordinary native block, +one Spark task per partition, for everything but `collect`. The partitions these actions +evaluate, and so the rows they return and the errors they raise, are Spark's, and a cached result +is scanned as many partitions as without local execution. ## Lifecycle and cancellation @@ -132,5 +140,5 @@ cd .. The native tests cover partition-order delivery, the bound on running partitions, read-ahead, errors reported in a partition's turn, cancellation with busy workers, the one-batch handoff and runtime shutdown. `CometLocalExecutionSuite` compares admitted queries with Spark's results and order, and covers ranges Spark would overflow, declined expressions, -AQE, reported orderings, takes and limits, iterators, tails, many partitions with sorts in a small memory pool, caching, the result size limit, job and task cancellation, and native errors. After +AQE, reported orderings, takes and limits, iterators, tails, many partitions with sorts in a small memory pool, caching and other RDD consumers, the result size limit, job and task cancellation, and native errors. After each test it checks that no native query, imported Arrow memory or result slot is left. diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 13bc2f3def1..ed42b91b9fd 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -598,6 +598,9 @@ diverge for several structural reasons: Imported JVM batches are reserved only while a reserving operator holds them, and exported native batches have usually been released by the time the JVM receives them yet stay resident until the JVM closes them (see [Crossing the FFI boundary](#crossing-the-ffi-boundary)). +- **Local execution's read-ahead.** The partitions of a local query hold the batches they produce + ahead of delivery, up to `spark.driver.maxResultSize` in all, without reserving them (see + [Local Execution](local-execution.md)). - **Overcommit.** Unlike everything above, this is not accidental. When Spark grants less than a `grow` asked for, Comet holds the difference anyway and deliberately does not tell Spark about it (see [The unified pools](#the-unified-pools)). `reserved()` includes it, but Spark's memory diff --git a/spark/src/main/scala/org/apache/comet/local/CometLocalExec.scala b/spark/src/main/scala/org/apache/comet/local/CometLocalExec.scala index 8e623ac0d61..a9cb7a4f648 100644 --- a/spark/src/main/scala/org/apache/comet/local/CometLocalExec.scala +++ b/spark/src/main/scala/org/apache/comet/local/CometLocalExec.scala @@ -44,8 +44,7 @@ case class CometLocalExec private[local] ( extends LeafExecNode { override def supportsColumnar: Boolean = true - // The result is one partition, but its rows are not clustered by anything. A query over a - // cached local result must still exchange to get its required distribution. + // The result is one partition, but its rows are not clustered by anything. override def outputPartitioning: Partitioning = UnknownPartitioning(1) // The ordering of the whole result, which admission derives from the block's partitions. diff --git a/spark/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala b/spark/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala index 491f8319e82..545b39400ec 100644 --- a/spark/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala +++ b/spark/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala @@ -32,34 +32,29 @@ import org.apache.spark.sql.catalyst.plans.physical.Partitioning import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, UnaryExecNode} /** - * Row result boundary of a local query. Spark's collect encodes and compresses every row of a - * result partition in its task, then decodes them on the driver. A local query has one result - * partition, so that work would run on a single thread. Admission requires an in-process local - * master, so collect still runs one Spark task (keeping cancellation, job groups and SQL metrics) - * but hands copied rows to the driver in this JVM. + * Row result boundary of a local query. Spark's collect encodes and compresses every row of each + * result partition in its task, then decodes them on the driver. A local collect instead runs the + * whole query as one native graph in a single Spark task (keeping cancellation, job groups and + * SQL metrics), which hands copied rows to the driver in this JVM, since admission requires an + * in-process local master. + * + * Only `collect` runs that way. Everything else reads the result as an RDD: a take requests + * partitions in growing batches, an iterator fetches one partition at a time, a tail scans from + * the last one, and a cache, `df.rdd` or PySpark's collect keep its partitions. So the node + * executes as the ordinary native block otherwise, one Spark task per partition, and reports the + * block's partitioning and ordering. */ case class CometLocalResultExec(child: SparkPlan) extends UnaryExecNode { override def output: Seq[Attribute] = child.output - override def outputPartitioning: Partitioning = child.outputPartitioning - override def outputOrdering: Seq[SortOrder] = child.outputOrdering + override def outputPartitioning: Partitioning = ordinary.outputPartitioning + override def outputOrdering: Seq[SortOrder] = ordinary.outputOrdering - override protected def doExecute(): RDD[InternalRow] = child.execute() + override protected def doExecute(): RDD[InternalRow] = ordinary.execute() override def executeCollect(): Array[InternalRow] = - if (sparkContext.isLocal) LocalResultHandoff.collect(execute(), maxResultSize) + if (sparkContext.isLocal) LocalResultHandoff.collect(child.execute(), maxResultSize) else super.executeCollect() - // A take requests partitions in growing batches until it has enough rows, an iterator fetches - // one result partition at a time to bound what it holds, and a tail scans partitions from the - // last one. A local query has a single partition holding the whole result, so these run the - // native block as ordinary Comet execution instead, one Spark task per partition, as without - // local execution. Which partitions they evaluate, and so which errors they raise, match Spark. - override def executeTake(n: Int): Array[InternalRow] = ordinary.executeTake(n) - - override def executeToIterator(): Iterator[InternalRow] = ordinary.executeToIterator() - - override def executeTail(n: Int): Array[InternalRow] = ordinary.executeTail(n) - private def ordinary: SparkPlan = child .collectFirst { case local: CometLocalExec => ColumnarToRowExec(local.block) } diff --git a/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala b/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala index 8c005a9485a..1fc964c6165 100644 --- a/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala @@ -194,18 +194,18 @@ class CometLocalExecutionSuite unbounded.show(2) } - test("a query over a cached local result keeps its required distribution") { + test("a cached result keeps the partitions of ordinary execution") { withConf("spark.sql.adaptive.enabled", "false") { val cached = spark.range(0, 1000, 1, 7).selectExpr("id", "id % 5 AS k") assert(localNodes(cached).size == 1) cached.cache() try { val counts = cached.groupBy("k").count() - // The cache holds the local result, a single partition with no clustering. - val cachedPlans = counts.queryExecution.withCachedData.collect { - case relation: InMemoryRelation => relation.cacheBuilder.cachedPlan + val relations = counts.queryExecution.withCachedData.collect { + case relation: InMemoryRelation => relation } - assert(cachedPlans.flatMap(_.collect { case p: CometLocalExec => p }).size == 1) + assert(relations.size == 1) + assert(relations.head.cacheBuilder.cachedColumnBuffers.getNumPartitions == 7) val plan = counts.queryExecution.executedPlan assert(plan.collect { case e: ShuffleExchangeLike => e }.nonEmpty, plan.toString) assert( @@ -215,6 +215,17 @@ class CometLocalExecutionSuite } } + test("RDD consumers read the partitions of ordinary execution") { + val query = spark.range(0, 300, 1, 6).selectExpr("id", "id * 2 AS v") + assert(localNodes(query).size == 1) + assert(query.rdd.getNumPartitions == 6) + assert(query.queryExecution.toRdd.getNumPartitions == 6) + assert(query.queryExecution.executedPlan.outputPartitioning.numPartitions == 6) + assert( + query.rdd.mapPartitions(rows => Iterator(rows.size)).collect().toSeq == Seq.fill(6)(50)) + assert(query.localCheckpoint().rdd.getNumPartitions == 6) + } + test("only an ordering that holds across concatenated partitions is reported") { // A range's slices are ranges of its ordering, so the result keeps it. assert(localNodes(spark.range(0, 100, 1, 7).toDF()).head.outputOrdering.nonEmpty) @@ -390,14 +401,14 @@ class CometLocalExecutionSuite assert(plan.isInstanceOf[CometLocalResultExec], plan.toString) // Each two-long UnsafeRow is 24 bytes, so 1,000 bytes fit 41 rows and the 42nd exceeds it. val failure = intercept[SparkException] { - LocalResultHandoff.collect(plan.execute(), 1000) + LocalResultHandoff.collect(localNodes(query).head.execute(), 1000) } assert(failure.getMessage.contains("spark.driver.maxResultSize"), failure.getMessage) assert(failure.getMessage.contains("at least 1008 bytes"), failure.getMessage) assert(new NativeLocal().activeQueries() == 0) // A result within the limit still succeeds against the same bound. def bounded(rows: Long) = - spark.range(0, rows, 1, 7).selectExpr("id", "id AS other").queryExecution.executedPlan + localNodes(spark.range(0, rows, 1, 7).selectExpr("id", "id AS other")).head assert(LocalResultHandoff.collect(bounded(41).execute(), 1000).length == 41) // Non-positive limits mean unlimited, as in Spark. assert(LocalResultHandoff.collect(bounded(100).execute(), 0).length == 100) @@ -410,7 +421,7 @@ class CometLocalExecutionSuite val future = executor.submit(new Callable[Unit] { override def call(): Unit = { spark.sparkContext.setJobGroup("comet-local-cancel", "local cancellation", true) - query.queryExecution.executedPlan.execute().foreachPartition { rows => + localNodes(query).head.execute().foreachPartition { rows => while (rows.hasNext) rows.next() } } From 28d5702dd6a17b7e77b827942d15f96b84511664 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Thu, 29 Oct 2026 16:54:42 -0700 Subject: [PATCH 10/12] refactor: leave limits to ordinary execution and poll local queries on the reader thread A root CollectLimitExec collects through executeTake, which no longer runs locally, so admitting it only changed its doExecute. It now keeps ordinary Comet execution. Partition order is the only delivery mode of LocalQuery, since Comet always used it. The coalescing mode and its exchange tests are removed. QueryHandle no longer runs its own producer task and channel: the JNI pull polls the query's stream on its own thread, parking between polls rather than waiting on a Tokio timer, and cancel wakes a waiting pull. --- .../contributor-guide/local-execution.md | 28 +- native/core/src/local.rs | 8 +- native/local/src/handle.rs | 241 ++++++++---------- native/local/src/lib.rs | 60 ++--- native/local/tests/bridge.rs | 65 +++-- native/local/tests/execution.rs | 198 +++----------- .../apache/comet/local/CometLocalRule.scala | 21 +- .../local/CometLocalExecutionSuite.scala | 24 +- 8 files changed, 272 insertions(+), 373 deletions(-) diff --git a/docs/source/contributor-guide/local-execution.md b/docs/source/contributor-guide/local-execution.md index ba5da78da70..5f04ffac910 100644 --- a/docs/source/contributor-guide/local-execution.md +++ b/docs/source/contributor-guide/local-execution.md @@ -42,9 +42,10 @@ admitted when: which the serialized plan is also checked for; - it has between 1 and 1024 partitions. -A root `CollectLimitExec`, which Spark plans for `take`, `head` and `show`, is kept above the -local query, and reads it as described under result delivery. Such plans have no exchange, so AQE does not wrap them and they are admitted -whether or not AQE is enabled. +A root `CollectLimitExec`, which Spark plans for `take`, `head`, `show` and `limit`, is not +admitted: it collects through `executeTake`, which does not run locally, so such a query keeps +ordinary Comet execution. Admitted plans have no exchange, so AQE does not wrap them and they are +admitted whether or not AQE is enabled. ## Execution @@ -90,10 +91,12 @@ metrics; `CometLocalExec` counts the rows and batches it delivers. ## Result delivery -The native side hands batches to the JVM through a queue that holds at most one batch. Exported +The JVM's pull polls the query's stream on its own thread, so a batch reaches it straight from +its partition's task; the read-ahead bound above is the only buffering. Exported Arrow arrays use `prepare_output` and `NativeUtil`; an imported batch stays valid after its native query is closed, until the JVM closes it. A pull returns "pending" after 50 ms, so that -the task can check whether it was killed. +the task can check whether it was killed. It waits by parking its thread rather than on a Tokio +timer, which only fires once a worker is free. `CometLocalResultExec` sits above Spark's columnar-to-row transition. For `collect` it still runs the single task as a Spark job, which keeps cancellation, job groups and SQL metrics, but @@ -104,9 +107,8 @@ bytes and stops reading, which closes the native query, once they exceed so this check is the stricter one. On a non-local master the node uses Spark's ordinary collect. -Every other action reads the result as an RDD of the block's partitions: a take requests -partitions in growing batches until it has enough rows, an iterator fetches one partition at a -time so that it holds at most one, a tail scans partitions from the last one, and a cache, +Every other action reads the result as an RDD of the block's partitions: an iterator fetches +one partition at a time so that it holds at most one, a tail scans partitions from the last one, and a cache, `df.rdd`, `foreachPartition`, `localCheckpoint` or PySpark's collect and `toLocalIterator` keep the partitions for later reads. So `CometLocalResultExec` executes as the ordinary native block, one Spark task per partition, for everything but `collect`. The partitions these actions @@ -117,9 +119,9 @@ is scanned as many partitions as without local execution. The JNI registry holds numeric, never reused IDs for at most 1024 live queries. EOF, errors, early termination and task completion remove an entry. A concurrent pull holds an `Arc`, so a -close cannot free an entry that is being read, and cancellation does not take the reader lock. -A query records when it has delivered its last batch; if the runtime stops it before that, the -reader reports an error instead of a truncated result. +close cannot free an entry that is being read, and cancellation wakes a waiting pull, so it waits +only for a poll the pull is already in. Each partition records when it has sent its last batch; +if the runtime stops it before that, the reader reports an error instead of a truncated result. Closing a query drops its streams on the calling thread, as `releasePlan` does for an ordinary plan without JVM input: cancelling a Tokio task only takes effect once a worker is free, and @@ -138,7 +140,7 @@ cd .. ``` The native tests cover partition-order delivery, the bound on running partitions, read-ahead, errors reported in a partition's turn, cancellation with busy workers, -the one-batch handoff and runtime shutdown. `CometLocalExecutionSuite` compares admitted queries +the pull that waits on its own thread, and runtime shutdown. `CometLocalExecutionSuite` compares admitted queries with Spark's results and order, and covers ranges Spark would overflow, declined expressions, -AQE, reported orderings, takes and limits, iterators, tails, many partitions with sorts in a small memory pool, caching and other RDD consumers, the result size limit, job and task cancellation, and native errors. After +AQE, reported orderings, limits left to ordinary execution, iterators, tails, many partitions with sorts in a small memory pool, caching and other RDD consumers, the result size limit, job and task cancellation, and native errors. After each test it checks that no native query, imported Arrow memory or result slot is left. diff --git a/native/core/src/local.rs b/native/core/src/local.rs index 5c734887afa..f24ea70c38e 100644 --- a/native/core/src/local.rs +++ b/native/core/src/local.rs @@ -135,8 +135,12 @@ pub extern "system" fn Java_org_apache_comet_local_NativeLocal_createQuery( "Local output schema width mismatch".into(), )); } - let query = LocalQuery::new(plan, context.session_ctx.task_ctx()) - .in_partition_order(max_running, read_ahead_bytes.max(0) as usize); + let query = LocalQuery::new( + plan, + context.session_ctx.task_ctx(), + max_running, + read_ahead_bytes.max(0) as usize, + ); let query = QueryHandle::start(query, &get_runtime())?; register( id, diff --git a/native/local/src/handle.rs b/native/local/src/handle.rs index b5f1f3f56a6..fa1c63da60d 100644 --- a/native/local/src/handle.rs +++ b/native/local/src/handle.rs @@ -15,47 +15,50 @@ // specific language governing permissions and limitations // under the License. -//! Bounded handoff to blocking callers, with cancellation independent of the reader lock. +//! Delivery to blocking callers, with cancellation that does not wait for a reader's wait. use std::panic::AssertUnwindSafe; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::{mpsc, Arc, Mutex}; -use std::task::Poll; -use std::time::Duration; +use std::sync::{Arc, Mutex, TryLockError}; +use std::task::{Context, Poll, Wake, Waker}; +use std::thread::Thread; +use std::time::{Duration, Instant}; use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::{DataFusionError, Result}; -use futures::future::poll_fn; -use futures::{FutureExt, StreamExt}; +use futures::StreamExt; use tokio::runtime::Handle; -use tokio::sync::{OwnedSemaphorePermit, Semaphore}; -use tokio::task::JoinHandle; use crate::{lock, LocalQuery, LocalQueryStream}; -struct Delivery { - batch: Result, - _permit: OwnedSemaphorePermit, -} - pub enum QueryPoll { Batch(RecordBatch), Pending, Finished, } -/// Holds at most one queued batch, plus the batch owned by the external consumer. -/// Cancelling or dropping the handle drops the query's stream on the calling thread and -/// unread batches with it. The runtime must remain alive until DataFusion finishes aborting -/// the tasks its operators spawn. +/// A started query, read by one blocking caller at a time. The reader polls the query's stream +/// on its own thread, so a batch reaches it straight from its partition's task, and waits by +/// parking rather than on a Tokio timer, which needs a free worker. Cancelling or dropping the +/// handle drops the query's stream on the calling thread and unread batches with it. The runtime +/// must remain alive until DataFusion finishes aborting the tasks its operators spawn. pub struct QueryHandle { - receiver: Mutex>, - /// The query's stream. The producer locks it only while polling it, so `cancel` can take it - /// whenever the producer is waiting. - stream: Arc>>, - producer: JoinHandle<()>, + /// The query's stream. A reader holds it for the whole of a poll; `cancel` wakes the reader + /// and so waits only for a poll of the stream it is already in. + stream: Mutex>, + /// The thread of the reader in a poll, which `cancel` wakes. + reader: Mutex>, cancelled: AtomicBool, - completed: Arc, + runtime: Handle, +} + +/// Wakes a reader parked in `QueryHandle::poll`. +struct Unpark(Thread); + +impl Wake for Unpark { + fn wake(self: Arc) { + self.0.unpark(); + } } impl QueryHandle { @@ -64,113 +67,79 @@ impl QueryHandle { let _guard = runtime.enter(); query.execute()? }; - let (sender, receiver) = mpsc::channel(); - let completed = Arc::new(AtomicBool::new(false)); - let worker_completed = Arc::clone(&completed); - let stream = Arc::new(Mutex::new(Some(stream))); - let polled = Arc::clone(&stream); - let producer = runtime.spawn(async move { - let capacity = Arc::new(Semaphore::new(1)); - let result = AssertUnwindSafe(async { - loop { - let permit = Arc::clone(&capacity) - .acquire_owned() - .await - .map_err(|e| DataFusionError::Execution(e.to_string()))?; - let next = poll_fn(|cx| match lock(&polled).as_mut() { - Some(stream) => stream.poll_next_unpin(cx).map(Some), - // Cancelled: the stream was taken, so this is not completion. - None => Poll::Ready(None), - }) - .await; - let Some(next) = next else { - return Ok(false); - }; - match next { - Some(batch) => { - let failed = batch.is_err(); - if sender - .send(Delivery { - batch, - _permit: permit, - }) - .is_err() - || failed - { - break; - } - } - None => break, - } - } - Ok::<_, DataFusionError>(true) - }) - .catch_unwind() - .await; - let error = match result { - Ok(Ok(true)) => None, - Ok(Ok(false)) => return, - Ok(Err(e)) => Some(e), - Err(_) => Some(DataFusionError::Execution( - "Local query worker panicked".into(), - )), - }; - if let Some(error) = error { - if let Ok(permit) = capacity.acquire_owned().await { - let _ = sender.send(Delivery { - batch: Err(error), - _permit: permit, - }); - } - } - // Publish before dropping sender. Runtime shutdown/abort skips this store; - // channel disconnection alone must never turn a truncated result into EOF. - worker_completed.store(true, Ordering::Release); - }); Ok(Self { - receiver: Mutex::new(receiver), - stream, - producer, + stream: Mutex::new(Some(stream)), + reader: Mutex::new(None), cancelled: AtomicBool::new(false), - completed, + runtime: runtime.clone(), }) } /// A bounded wait lets a JVM caller check Spark task interruption between polls. - /// Only one reader may call this at a time. Cancellation never needs this lock. + /// Only one reader may call this at a time. pub fn poll(&self, wait: Duration) -> Result { - if self.cancelled.load(Ordering::Acquire) { - return Err(DataFusionError::Execution("Local query cancelled".into())); - } - let receiver = self.receiver.try_lock().map_err(|_| { - DataFusionError::Execution("Concurrent local query result readers".into()) - })?; - let delivery = receiver.recv_timeout(wait); - if self.cancelled.load(Ordering::Acquire) { - return Err(DataFusionError::Execution("Local query cancelled".into())); - } - match delivery { - Ok(delivery) => delivery.batch.map(QueryPoll::Batch), - Err(mpsc::RecvTimeoutError::Timeout) => Ok(QueryPoll::Pending), - Err(mpsc::RecvTimeoutError::Disconnected) => { - if self.completed.load(Ordering::Acquire) { - Ok(QueryPoll::Finished) - } else { - Err(DataFusionError::Execution( - "Local query worker stopped without completing".into(), - )) - } + let mut stream = match self.stream.try_lock() { + Ok(stream) => stream, + Err(TryLockError::Poisoned(poisoned)) => poisoned.into_inner(), + Err(TryLockError::WouldBlock) if self.cancelled.load(Ordering::Acquire) => { + return Err(cancelled()) + } + Err(TryLockError::WouldBlock) => { + return Err(DataFusionError::Execution( + "Concurrent local query result readers".into(), + )) } + }; + *lock(&self.reader) = Some(std::thread::current()); + let result = self.poll_stream(&mut stream, wait); + *lock(&self.reader) = None; + result + } + + fn poll_stream( + &self, + stream: &mut Option, + wait: Duration, + ) -> Result { + let deadline = Instant::now() + wait; + let waker = Waker::from(Arc::new(Unpark(std::thread::current()))); + let mut cx = Context::from_waker(&waker); + // Starting a partition executes its operators, which may spawn tasks. + let _guard = self.runtime.enter(); + loop { + // Checked after the reader is registered, so a cancel either sees the reader and + // wakes it, or is seen here. + if self.cancelled.load(Ordering::Acquire) { + return Err(cancelled()); + } + let Some(query) = stream.as_mut() else { + return Err(cancelled()); + }; + // A panic must not unwind into the JNI caller. + let next = + std::panic::catch_unwind(AssertUnwindSafe(|| query.poll_next_unpin(&mut cx))) + .map_err(|_| DataFusionError::Execution("Local query panicked".into()))?; + match next { + Poll::Ready(Some(batch)) => return batch.map(QueryPoll::Batch), + Poll::Ready(None) => return Ok(QueryPoll::Finished), + Poll::Pending => (), + } + let now = Instant::now(); + if now >= deadline { + return Ok(QueryPoll::Pending); + } + std::thread::park_timeout(deadline - now); } } /// Stops the query and drops its stream on the calling thread, rather than leaving that to - /// the cancelled producer: a cancellation only runs once a Tokio worker is free, and every - /// worker can be waiting, for instance in Spark's `acquireMemory` for memory this stream - /// holds. Taking the stream waits only for a poll the producer is already in. + /// a Tokio task: a task only runs once a worker is free, and every worker can be waiting, + /// for instance in Spark's `acquireMemory` for memory this stream holds. pub fn cancel(&self) { self.cancelled.store(true, Ordering::Release); - self.producer.abort(); + if let Some(reader) = lock(&self.reader).as_ref() { + reader.unpark(); + } let stream = lock(&self.stream).take(); // A panic while dropping must not unwind into the JNI caller. let _ = std::panic::catch_unwind(AssertUnwindSafe(|| drop(stream))); @@ -183,6 +152,10 @@ impl Drop for QueryHandle { } } +fn cancelled() -> DataFusionError { + DataFusionError::Execution("Local query cancelled".into()) +} + #[cfg(test)] mod tests { use super::*; @@ -191,36 +164,40 @@ mod tests { use datafusion::physical_plan::test::exec::BlockingExec; #[test] - fn cancel_interrupts_a_blocked_reader_without_its_mutex() { + fn cancel_wakes_a_waiting_reader() { let runtime = tokio::runtime::Builder::new_multi_thread() .worker_threads(1) .build() .unwrap(); - let plan = Arc::new(BlockingExec::new(Arc::new(Schema::empty()), 2)); + let plan = BlockingExec::new(Arc::new(Schema::empty()), 2); + let refs = plan.refs(); let query = Arc::new( QueryHandle::start( - LocalQuery::new(plan, Arc::new(TaskContext::default())), + LocalQuery::new( + Arc::new(plan), + Arc::new(TaskContext::default()), + usize::MAX, + usize::MAX, + ), runtime.handle(), ) .unwrap(), ); let reader = Arc::clone(&query); - let (started_tx, started_rx) = mpsc::channel(); - let (sender, receiver) = mpsc::channel(); + let (sender, receiver) = std::sync::mpsc::channel(); let thread = std::thread::spawn(move || { - // Hold exactly the receiver lock used by poll. A handshake avoids racing - // poll's try_lock with a test thread inspecting that same lock. - let incoming = reader.receiver.lock().unwrap(); - started_tx.send(()).unwrap(); - let result = incoming.recv_timeout(Duration::from_secs(60)); - sender - .send(matches!(result, Err(mpsc::RecvTimeoutError::Disconnected))) - .unwrap(); + sender.send(reader.poll(Duration::from_secs(60))).unwrap(); }); - started_rx.recv_timeout(Duration::from_secs(5)).unwrap(); - assert!(query.poll(Duration::ZERO).is_err()); + while lock(&query.reader).is_none() { + std::thread::yield_now(); + } + assert!(query.poll(Duration::ZERO).is_err(), "a second reader"); + let started = Instant::now(); query.cancel(); - assert!(receiver.recv_timeout(Duration::from_secs(5)).unwrap()); + let result = receiver.recv_timeout(Duration::from_secs(5)).unwrap(); + assert!(result.err().unwrap().to_string().contains("cancelled")); + assert!(started.elapsed() < Duration::from_secs(5)); + assert_eq!(refs.strong_count(), 0); thread.join().unwrap(); } } diff --git a/native/local/src/lib.rs b/native/local/src/lib.rs index 3efc81a114c..33a4273246a 100644 --- a/native/local/src/lib.rs +++ b/native/local/src/lib.rs @@ -36,41 +36,29 @@ use datafusion::common::{DataFusionError, Result}; use datafusion::execution::TaskContext; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ - execute_stream, ExecutionPlan, ExecutionPlanProperties, RecordBatchStream, - SendableRecordBatchStream, + ExecutionPlan, ExecutionPlanProperties, RecordBatchStream, SendableRecordBatchStream, }; use futures::future::poll_fn; use futures::{FutureExt, Stream, StreamExt}; use tokio::sync::Notify; use tokio::task::JoinHandle; -/// Batches each root partition may always hold ahead of delivery, when results are delivered in -/// partition order, whatever the read-ahead limit. +/// Batches each root partition may always hold ahead of delivery, whatever the read-ahead +/// limit. const PARTITION_BUFFER: usize = 2; /// One execution's graph and resources. Consumed when execution starts; not a plan cache. pub struct LocalQuery { plan: Arc, context: Arc, - /// Delivers the partitions in order, with at most this many running at once and this many - /// bytes read ahead. - in_partition_order: Option<(usize, usize)>, + max_running: usize, + read_ahead: usize, } impl LocalQuery { /// Takes ownership of a fresh graph and its query-scoped context. /// - /// The planner must satisfy operator distribution and ordering requirements. It must not - /// reuse stateful operator instances from another query or attempt. - pub fn new(plan: Arc, context: Arc) -> Self { - Self { - plan, - context, - in_partition_order: None, - } - } - - /// Delivers the root partitions one after another, in partition order, as Spark returns + /// The root partitions are delivered one after another, in partition order, as Spark returns /// the partitions of a result. Up to `max_running` partitions run at once, and the next one /// starts when the first is delivered, as Spark runs at most as many tasks as it has cores. /// A partition's operators register their memory consumers when it starts, so this also @@ -78,28 +66,34 @@ impl LocalQuery { /// /// The running partitions do not wait for their turn: each runs ahead of delivery while the /// batches they all hold total less than `read_ahead` bytes, and up to a few batches beyond - /// that. Only for roots whose partitions do not share an exchange: a partition that waits to - /// be delivered would otherwise leave its exchange buffering for the others. - pub fn in_partition_order(mut self, max_running: usize, read_ahead: usize) -> Self { - self.in_partition_order = Some((max_running.max(1), read_ahead)); - self + /// that. The root's partitions must not share an exchange: a partition that waits to be + /// delivered would otherwise leave its exchange buffering for the others. + /// + /// The planner must satisfy operator distribution and ordering requirements. It must not + /// reuse stateful operator instances from another query or attempt. + pub fn new( + plan: Arc, + context: Arc, + max_running: usize, + read_ahead: usize, + ) -> Self { + Self { + plan, + context, + max_running: max_running.max(1), + read_ahead, + } } - /// Starts all root partitions and exposes one streaming result. - /// - /// By default the partitions are coalesced as they produce batches, which does not - /// preserve their order. A single-partition root preserves its ordering either way. No - /// result-wide collection is performed here. + /// Starts the first root partitions and exposes one streaming result. No result-wide + /// collection is performed here. pub fn execute(self) -> Result { tokio::runtime::Handle::try_current().map_err(|_| { DataFusionError::Execution("LocalQuery requires an active Tokio runtime".into()) })?; let schema = self.plan.schema(); - let stream = if let Some((max_running, read_ahead)) = self.in_partition_order { - in_partition_order(&self.plan, &self.context, max_running, read_ahead)? - } else { - execute_stream(Arc::clone(&self.plan), Arc::clone(&self.context))? - }; + let stream = + in_partition_order(&self.plan, &self.context, self.max_running, self.read_ahead)?; Ok(LocalQueryStream { stream: Some(stream), query: Some(self), diff --git a/native/local/tests/bridge.rs b/native/local/tests/bridge.rs index 2d10a30ff6b..e3642a8195b 100644 --- a/native/local/tests/bridge.rs +++ b/native/local/tests/bridge.rs @@ -73,13 +73,17 @@ fn values(batch: &RecordBatch) -> &[i64] { } #[test] -fn partition_order_is_preserved_through_the_handoff() { +fn partition_order_is_preserved_through_the_handle() { let runtime = runtime(); let plan = ranges(&[(0, 10), (10, 10), (10, 25), (25, 26)], 3); let refs = Arc::downgrade(&plan); let query = QueryHandle::start( - LocalQuery::new(plan, Arc::new(TaskContext::default())) - .in_partition_order(usize::MAX, usize::MAX), + LocalQuery::new( + plan, + Arc::new(TaskContext::default()), + usize::MAX, + usize::MAX, + ), runtime.handle(), ) .unwrap(); @@ -107,7 +111,12 @@ fn pending_is_distinct_from_eof_and_drop_releases_graph() { let plan = BlockingExec::new(schema(), 4); let refs = plan.refs(); let query = QueryHandle::start( - LocalQuery::new(Arc::new(plan), Arc::new(TaskContext::default())), + LocalQuery::new( + Arc::new(plan), + Arc::new(TaskContext::default()), + usize::MAX, + usize::MAX, + ), runtime.handle(), ) .unwrap(); @@ -132,7 +141,12 @@ fn worker_panic_is_an_error_not_successful_eof() { let runtime = runtime(); let plan = Arc::new(PanicExec::new(schema(), 2)); let query = QueryHandle::start( - LocalQuery::new(plan, Arc::new(TaskContext::default())), + LocalQuery::new( + plan, + Arc::new(TaskContext::default()), + usize::MAX, + usize::MAX, + ), runtime.handle(), ) .unwrap(); @@ -155,7 +169,12 @@ fn startup_failure_is_returned() { let runtime = runtime(); let plan = Arc::new(ErrorExec::new()); assert!(QueryHandle::start( - LocalQuery::new(plan, Arc::new(TaskContext::default())), + LocalQuery::new( + plan, + Arc::new(TaskContext::default()), + usize::MAX, + usize::MAX + ), runtime.handle() ) .is_err()); @@ -168,6 +187,8 @@ fn runtime_shutdown_is_an_error_not_successful_eof() { LocalQuery::new( Arc::new(BlockingExec::new(schema(), 4)), Arc::new(TaskContext::default()), + usize::MAX, + usize::MAX, ), runtime.handle(), ) @@ -190,18 +211,19 @@ fn backpressured_query_does_not_starve_other_queries_and_cancel_releases_graph() let stalled = ranges(&[(0, 100_000), (100_000, 200_000)], 5); let refs = Arc::downgrade(&stalled); let stalled = QueryHandle::start( - LocalQuery::new(stalled, Arc::new(TaskContext::default())) - .in_partition_order(usize::MAX, usize::MAX), + LocalQuery::new(stalled, Arc::new(TaskContext::default()), usize::MAX, 0), runtime.handle(), ) .unwrap(); - // Leave its one-batch result channel unread while a second query makes progress. + // Leave it unread, with its partitions at their read-ahead bound, while a second query + // makes progress. let healthy = QueryHandle::start( LocalQuery::new( ranges(&[(100, 150), (150, 200), (200, 237)], 3), Arc::new(TaskContext::default()), - ) - .in_partition_order(usize::MAX, usize::MAX), + usize::MAX, + usize::MAX, + ), runtime.handle(), ) .unwrap(); @@ -239,7 +261,13 @@ fn backpressured_query_does_not_starve_other_queries_and_cancel_releases_graph() fn shutdown_after_partial_output_never_silently_truncates_the_result() { let runtime = runtime(); let query = QueryHandle::start( - LocalQuery::new(ranges(&[(0, 100_000)], 5), Arc::new(TaskContext::default())), + // No read-ahead, so the partition holds at most two batches ahead of delivery. + LocalQuery::new( + ranges(&[(0, 100_000)], 5), + Arc::new(TaskContext::default()), + usize::MAX, + 0, + ), runtime.handle(), ) .unwrap(); @@ -260,7 +288,10 @@ fn shutdown_after_partial_output_never_silently_truncates_the_result() { match query.poll(Duration::from_millis(20)) { Ok(QueryPoll::Batch(_)) => { queued += 1; - assert!(queued <= 1, "handoff exceeded its one-batch bound"); + assert!( + queued <= 2, + "the partition held more than two batches ahead" + ); } Ok(QueryPoll::Pending) => (), Ok(QueryPoll::Finished) => panic!("partial output was reported as complete"), @@ -278,8 +309,12 @@ fn cancel_releases_the_graph_while_every_worker_is_busy() { let plan = BlockingExec::new(schema(), 3); let refs = plan.refs(); let query = QueryHandle::start( - LocalQuery::new(Arc::new(plan), Arc::new(TaskContext::default())) - .in_partition_order(usize::MAX, usize::MAX), + LocalQuery::new( + Arc::new(plan), + Arc::new(TaskContext::default()), + usize::MAX, + usize::MAX, + ), runtime.handle(), ) .unwrap(); diff --git a/native/local/tests/execution.rs b/native/local/tests/execution.rs index de45504f8c7..46b502db929 100644 --- a/native/local/tests/execution.rs +++ b/native/local/tests/execution.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use std::collections::{BTreeMap, BTreeSet}; +use std::collections::BTreeMap; use std::fmt; use std::sync::{Arc, Mutex, Weak}; use std::time::Duration; @@ -27,17 +27,14 @@ use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{DataFusionError, Result}; use datafusion::execution::config::SessionConfig; use datafusion::execution::TaskContext; -use datafusion::physical_expr::{expressions::Column, PhysicalExpr}; -use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_plan::empty::EmptyExec; -use datafusion::physical_plan::limit::GlobalLimitExec; -use datafusion::physical_plan::repartition::RepartitionExec; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::test::exec::{BlockingExec, ErrorExec, MockExec}; use datafusion::physical_plan::union::UnionExec; use datafusion::physical_plan::{ - ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, - PlanProperties, RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, }; use datafusion_comet_local::LocalQuery; use futures::{StreamExt, TryStreamExt}; @@ -86,21 +83,13 @@ fn input() -> Arc { UnionExec::try_new(inputs).unwrap() } -fn repartition( - input: Arc, - partitioning: Partitioning, -) -> Arc { - Arc::new(RepartitionExec::try_new(input, partitioning).unwrap()) -} - #[derive(Debug, Default)] struct Observations { executions: Mutex>, batches: Mutex>, - key_partitions: Mutex>>, } -/// Records actual partition execution and routing without changing the data or scheduling. +/// Records actual partition execution without changing the data or scheduling. #[derive(Debug)] struct ObserveExec { input: Arc, @@ -163,21 +152,12 @@ impl ExecutionPlan for ObserveExec { self.observations.executions.lock().unwrap().push(partition); let observations = Arc::clone(&self.observations); let stream = self.input.execute(partition, ctx)?.map_ok(move |batch| { - let keys = batch - .column(0) - .as_any() - .downcast_ref::() - .unwrap(); *observations .batches .lock() .unwrap() .entry(partition) .or_default() += 1; - let mut locations = observations.key_partitions.lock().unwrap(); - for key in keys.values() { - locations.entry(*key).or_default().insert(partition); - } batch }); Ok(Box::pin(RecordBatchStreamAdapter::new( @@ -225,78 +205,11 @@ async fn wait_started(observations: &Observations, partitions: usize) { assert_executed_once(observations, partitions); } -async fn check_exchange(hash: bool, nested: bool) { - let (input, source_observations) = observe(input()); - let partitioning = if hash { - Partitioning::Hash(vec![Arc::new(Column::new("key", 0))], 7) - } else { - Partitioning::RoundRobinBatch(7) - }; - let (mut plan, exchange_observations) = observe(repartition(input, partitioning)); - if nested { - plan = repartition(plan, Partitioning::RoundRobinBatch(11)); - } - let plan_ref = Arc::downgrade(&plan); - let ctx = context(); - let ctx_ref = Arc::downgrade(&ctx); - let pool = Arc::clone(ctx.memory_pool()); - let mut stream = LocalQuery::new(plan, ctx).execute().unwrap(); - let batches: Vec = timeout(DEADLINE, stream.by_ref().try_collect()) - .await - .expect("repartition stalled") - .unwrap(); - let mut ids: Vec = batches - .iter() - .flat_map(|batch| { - batch - .column(1) - .as_any() - .downcast_ref::() - .unwrap() - .values() - .to_vec() - }) - .collect(); - ids.sort_unstable(); - assert_eq!( - ids, - (0..(INPUTS * ROWS_PER_INPUT) as i64).collect::>() - ); - assert_executed_once(&source_observations, INPUTS); - assert_executed_once(&exchange_observations, 7); - if hash { - let locations = exchange_observations.key_partitions.lock().unwrap(); - assert_eq!(locations.len(), 17); - assert!(locations.values().all(|partitions| partitions.len() == 1)); - } - // EOF releases resources even if the JVM would retain its exhausted iterator. - assert!(stream.next().await.is_none()); - assert_eq!(stream.schema(), schema()); - wait_released(&plan_ref).await; - wait_released(&ctx_ref).await; - assert_eq!(pool.reserved(), 0); -} - -#[tokio::test(flavor = "current_thread")] -async fn round_robin_more_partitions_than_workers() { - check_exchange(false, false).await; -} - -#[tokio::test(flavor = "current_thread")] -async fn hash_routes_all_input_partitions_through_one_exchange() { - check_exchange(true, false).await; -} - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn nested_exchanges_more_partitions_than_workers() { - check_exchange(true, true).await; -} - #[tokio::test] async fn single_partition_preserves_order() { let expected = batch(0..12); let plan = Arc::new(MockExec::new(vec![Ok(expected.clone())], schema())); - let batches: Vec = LocalQuery::new(plan, context()) + let batches: Vec = LocalQuery::new(plan, context(), usize::MAX, usize::MAX) .execute() .unwrap() .try_collect() @@ -309,29 +222,14 @@ async fn single_partition_preserves_order() { async fn empty_result_finishes_with_schema() { for partitions in [0, 1, 7] { let plan = Arc::new(EmptyExec::new(schema()).with_partitions(partitions)); - let mut stream = LocalQuery::new(plan, context()).execute().unwrap(); + let mut stream = LocalQuery::new(plan, context(), usize::MAX, usize::MAX) + .execute() + .unwrap(); assert!(timeout(DEADLINE, stream.next()).await.unwrap().is_none()); assert_eq!(stream.schema(), schema()); } } -#[tokio::test] -async fn cancel_pending_exchange_releases_inputs() { - let input = BlockingExec::new(schema(), 4); - let refs = input.refs(); - let (input, observations) = observe(Arc::new(input)); - let plan = repartition(input, Partitioning::RoundRobinBatch(7)); - let mut stream = LocalQuery::new(plan, context()).execute().unwrap(); - // Prove cancellation tears down running streams, not just an unstarted graph. - wait_started(&observations, 4).await; - // Drive the execution to Pending, then cancel without waiting for a batch. - assert!(futures::poll!(stream.next()).is_pending()); - stream.cancel(); - stream.cancel(); - assert!(stream.next().await.is_none()); - wait_released(&refs).await; -} - fn partially_blocked_input(fail: bool) -> (Arc, Weak<()>, Arc) { let blocking = BlockingExec::new(schema(), 3); let refs = blocking.refs(); @@ -346,17 +244,15 @@ fn partially_blocked_input(fail: bool) -> (Arc, Weak<()>, Arc blocking, ]) .unwrap(); - ( - repartition(input, Partitioning::RoundRobinBatch(7)), - refs, - observations, - ) + (input, refs, observations) } #[tokio::test] async fn drop_after_first_batch_aborts_sibling_partitions() { let (plan, refs, observations) = partially_blocked_input(false); - let mut stream = LocalQuery::new(plan, context()).execute().unwrap(); + let mut stream = LocalQuery::new(plan, context(), usize::MAX, usize::MAX) + .execute() + .unwrap(); wait_started(&observations, 3).await; timeout(DEADLINE, stream.next()) .await @@ -373,7 +269,9 @@ async fn first_error_cancels_siblings_and_is_terminal() { let ctx = context(); let ctx_ref = Arc::downgrade(&ctx); let pool = Arc::clone(ctx.memory_pool()); - let mut stream = LocalQuery::new(plan, ctx).execute().unwrap(); + let mut stream = LocalQuery::new(plan, ctx, usize::MAX, usize::MAX) + .execute() + .unwrap(); wait_started(&observations, 3).await; let error = timeout(DEADLINE, stream.next()) .await @@ -387,37 +285,18 @@ async fn first_error_cancels_siblings_and_is_terminal() { assert_eq!(pool.reserved(), 0); } -#[tokio::test] -async fn global_limit_stops_exchange_without_draining_inputs() { - let (input, refs, observations) = partially_blocked_input(false); - let plan = Arc::new(GlobalLimitExec::new( - Arc::new(CoalescePartitionsExec::new(input)), - 0, - Some(1), - )); - let mut stream = LocalQuery::new(plan, context()).execute().unwrap(); - wait_started(&observations, 3).await; - let batches: Vec = timeout(DEADLINE, stream.by_ref().try_collect()) - .await - .unwrap() - .unwrap(); - assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 1); - wait_released(&refs).await; -} - #[tokio::test] async fn cancelling_one_query_does_not_cancel_another() { let blocking_input = BlockingExec::new(schema(), 4); let refs = blocking_input.refs(); let (blocking_input, observations) = observe(Arc::new(blocking_input)); - let mut blocked = LocalQuery::new( - repartition(blocking_input, Partitioning::RoundRobinBatch(7)), - context(), - ) - .execute() - .unwrap(); + let mut blocked = LocalQuery::new(blocking_input, context(), usize::MAX, usize::MAX) + .execute() + .unwrap(); wait_started(&observations, 4).await; - let other = LocalQuery::new(input(), context()).execute().unwrap(); + let other = LocalQuery::new(input(), context(), usize::MAX, usize::MAX) + .execute() + .unwrap(); blocked.cancel(); let batches: Vec = timeout(DEADLINE, other.try_collect()) .await @@ -434,7 +313,10 @@ async fn cancelling_one_query_does_not_cancel_another() { async fn startup_error_releases_query() { let plan = Arc::new(ErrorExec::new()); let refs = Arc::downgrade(&plan); - let error = LocalQuery::new(plan, context()).execute().err().unwrap(); + let error = LocalQuery::new(plan, context(), usize::MAX, usize::MAX) + .execute() + .err() + .unwrap(); assert!(error.to_string().contains("ErrorExec")); wait_released(&refs).await; } @@ -443,7 +325,10 @@ async fn startup_error_releases_query() { fn missing_runtime_is_an_error_not_a_panic() { let plan = Arc::new(EmptyExec::new(schema())); let refs = Arc::downgrade(&plan); - let error = LocalQuery::new(plan, context()).execute().err().unwrap(); + let error = LocalQuery::new(plan, context(), usize::MAX, usize::MAX) + .execute() + .err() + .unwrap(); assert!(error.to_string().contains("active Tokio runtime")); assert_eq!(refs.strong_count(), 0); } @@ -466,8 +351,7 @@ async fn partition_order_delivers_partitions_in_turn() { let plan = ordered_input(7, 9); let batches: Vec = timeout( DEADLINE, - LocalQuery::new(plan, context()) - .in_partition_order(usize::MAX, usize::MAX) + LocalQuery::new(plan, context(), usize::MAX, usize::MAX) .execute() .unwrap() .try_collect(), @@ -502,8 +386,7 @@ fn failing() -> Arc { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn partition_order_reports_a_later_failure_in_its_turn() { let plan = UnionExec::try_new(vec![ordered_input(1, 3), failing()]).unwrap(); - let mut stream = LocalQuery::new(plan, context()) - .in_partition_order(usize::MAX, usize::MAX) + let mut stream = LocalQuery::new(plan, context(), usize::MAX, usize::MAX) .execute() .unwrap(); // The earlier partition's batches come first, even though the later one failed at once. @@ -528,8 +411,7 @@ async fn partition_order_never_reports_the_failure_of_a_partition_not_reached() let blocking = BlockingExec::new(schema(), 1); let refs = blocking.refs(); let plan = UnionExec::try_new(vec![Arc::new(blocking), failing()]).unwrap(); - let mut stream = LocalQuery::new(plan, context()) - .in_partition_order(usize::MAX, usize::MAX) + let mut stream = LocalQuery::new(plan, context(), usize::MAX, usize::MAX) .execute() .unwrap(); // A consumer that stops in the first partition, as a limit does, sees no error. @@ -546,8 +428,7 @@ async fn partition_order_cancel_releases_running_partitions() { let blocking = BlockingExec::new(schema(), 3); let refs = blocking.refs(); let (blocking, observations) = observe(Arc::new(blocking)); - let mut stream = LocalQuery::new(blocking, context()) - .in_partition_order(usize::MAX, usize::MAX) + let mut stream = LocalQuery::new(blocking, context(), usize::MAX, usize::MAX) .execute() .unwrap(); wait_started(&observations, 3).await; @@ -561,8 +442,7 @@ async fn partition_order_cancel_releases_running_partitions() { async fn partition_order_runs_at_most_the_given_number_of_partitions() { let (plan, observations) = observe(ordered_input(6, 3)); let executed = || observations.executions.lock().unwrap().clone(); - let mut stream = LocalQuery::new(plan, context()) - .in_partition_order(2, usize::MAX) + let mut stream = LocalQuery::new(plan, context(), 2, usize::MAX) .execute() .unwrap(); // The first two partitions start with the query; the rest wait for their turn. @@ -604,8 +484,7 @@ fn produced(observations: &Observations, partition: usize) -> usize { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn partition_order_runs_later_partitions_ahead_within_the_limit() { let (plan, observations) = behind_a_blocked_partition(); - let mut stream = LocalQuery::new(plan, context()) - .in_partition_order(2, usize::MAX) + let mut stream = LocalQuery::new(plan, context(), 2, usize::MAX) .execute() .unwrap(); // The second partition runs to its end while the first has delivered nothing. @@ -622,10 +501,7 @@ async fn partition_order_runs_later_partitions_ahead_within_the_limit() { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn partition_order_holds_a_few_batches_beyond_the_limit() { let (plan, observations) = behind_a_blocked_partition(); - let mut stream = LocalQuery::new(plan, context()) - .in_partition_order(2, 0) - .execute() - .unwrap(); + let mut stream = LocalQuery::new(plan, context(), 2, 0).execute().unwrap(); tokio::time::sleep(Duration::from_millis(200)).await; assert_eq!(produced(&observations, 1), 2); stream.cancel(); diff --git a/spark/src/main/scala/org/apache/comet/local/CometLocalRule.scala b/spark/src/main/scala/org/apache/comet/local/CometLocalRule.scala index 69d49b21bef..2e64994042d 100644 --- a/spark/src/main/scala/org/apache/comet/local/CometLocalRule.scala +++ b/spark/src/main/scala/org/apache/comet/local/CometLocalRule.scala @@ -25,7 +25,7 @@ import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.expressions.{SortOrder, SubqueryExpression} import org.apache.spark.sql.catalyst.plans.physical.RangePartitioning import org.apache.spark.sql.comet.{CometNativeExec, CometRangeExec} -import org.apache.spark.sql.execution.{CollectLimitExec, QueryExecution, SparkPlan} +import org.apache.spark.sql.execution.{QueryExecution, SparkPlan} import org.apache.spark.sql.internal.SQLConf import com.google.protobuf.Descriptors.FieldDescriptor.JavaType @@ -48,17 +48,16 @@ private[comet] object CometLocalRule { /** Expressions that call back into the JVM, which a local query does not do. */ private val JvmCallbacks = Set(Expr.ExprStructCase.SUBQUERY, Expr.ExprStructCase.JVM_SCALAR_UDF) - /** An already admitted local query, with or without its result boundary or limit wrapper. */ + /** An already admitted local query, with or without its result boundary. */ def isLocal(plan: SparkPlan): Boolean = plan match { case _: CometLocalExec | _: CometLocalResultExec => true - case limit: CollectLimitExec => isLocal(limit.child) case _ => false } /** * `plan` as a local query if it is admitted, and otherwise as `convert`, Comet's ordinary - * conversion, produces it. A root `CollectLimitExec` (from `take`, `head` or `show`) is kept: - * it reads the result through `CometLocalResultExec.executeTake`. + * conversion, produces it. A root `CollectLimitExec` (from `take`, `head` or `show`) is not + * admitted: only `collect` runs locally, and a limit collects through `executeTake`. */ def apply( session: SparkSession, @@ -68,19 +67,13 @@ private[comet] object CometLocalRule { .get(SQLConf.get) || !environmentAdmits(session, plan)) { return convert(plan) } - val limit = plan match { - case limit: CollectLimitExec if limit.offset == 0 => Some(limit) - case _ => None - } - convert(limit.map(_.child).getOrElse(plan)) match { + convert(plan) match { case block: CometNativeExec if admits(block) => val partitions = block.collectLeaves().head.outputPartitioning.numPartitions val spec = LocalQuerySpec(block.serializedPlanOpt.plan.get, partitions, block.output.size) - val local = CometLocalResultExec( + CometLocalResultExec( CometLocalExec(block.output, spec, resultOrdering(block, partitions), block)) - limit.map(_.copy(child = local)).getOrElse(local) - case converted if limit.isEmpty => converted - case _ => convert(plan) + case converted => converted } } diff --git a/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala b/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala index 1fc964c6165..4d85b95d351 100644 --- a/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala @@ -179,8 +179,11 @@ class CometLocalExecutionSuite } } - test("limits read partitions in order and take, head and show close early") { + test("limits, takes, heads and shows keep ordinary execution") { val query = spark.range(0, 1000, 1, 7).toDF() + assert( + localNodes(query.limit(10)).isEmpty, + query.limit(10).queryExecution.executedPlan.toString) // Unordered limits return the first rows of the first partitions, as Spark does. for (_ <- 0 until 3) { assert( @@ -268,7 +271,7 @@ class CometLocalExecutionSuite .filter("id >= 99999") .selectExpr("100L DIV (100000L - id) AS v") .limit(1) - assert(localNodes(query()).size == 1, query().queryExecution.executedPlan.toString) + assert(localNodes(query()).isEmpty, query().queryExecution.executedPlan.toString) // Spark reads only the first partition; the second divides by zero at id 100000. assert(query().collect().toSeq == Seq(Row(100L))) assert(query().collect().toSeq == expected(query())) @@ -280,6 +283,21 @@ class CometLocalExecutionSuite .selectExpr("100L DIV (100000L - id) AS v") val failure = intercept[Exception](all().collect()) assert(failure.getMessage.contains("DIVIDE_BY_ZERO"), failure.getMessage) + // A limit read as an RDD takes rows in every partition, as Spark does, so it reaches the + // third partition's division by zero at id 4. + def iterated(): Seq[Row] = + spark + .range(0, 6, 1, 3) + .selectExpr("100L DIV (4L - id) AS v") + .limit(1) + .toLocalIterator() + .asScala + .toSeq + val iteratorFailure = intercept[Exception](iterated()) + assert(iteratorFailure.getMessage.contains("DIVIDE_BY_ZERO"), iteratorFailure.getMessage) + val sparkFailure = + withConf("spark.comet.enabled", "false")(intercept[Exception](iterated())) + assert(sparkFailure.getMessage.contains("DIVIDE_BY_ZERO"), sparkFailure.getMessage) } } @@ -359,7 +377,7 @@ class CometLocalExecutionSuite .filter("id % 2 = 0") .selectExpr("100L DIV (4L - id) AS v") .limit(2) - assert(localNodes(query()).size == 1, query().queryExecution.executedPlan.toString) + assert(localNodes(query()).isEmpty, query().queryExecution.executedPlan.toString) val failure = intercept[Exception](query().collect()) assert(failure.getMessage.contains("DIVIDE_BY_ZERO"), failure.getMessage) val sparkFailure = withConf("spark.comet.enabled", "false") { From f99686ccd2607c8c629ce3f0d88c557dea2c6663 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sat, 31 Oct 2026 01:02:39 -0700 Subject: [PATCH 11/12] fix: share task slots with other native tasks and start local partitions off the reader Local queries run at most spark.executor.cores / spark.task.cpus partitions at once, divided again among the executor's tasks running native plans whenever another partition would start, so that concurrent local collects leave each partition a task's share of memory. A partition starts its operators on its own task, so a start waiting for a memory grant no longer holds up the reader's bounded poll. Closing a local query releases its context through release_plan, waiting for outstanding reservations. Exported batches honor spark.comet.debug.enabled and imported ones the strict timestamp label check. RDD actions run the block as ordinary planning prepares it, and child-first executor class loading is not admitted. --- .../contributor-guide/local-execution.md | 64 ++++++---- native/core/src/execution/jni_api.rs | 4 +- native/core/src/execution/memory_pools/mod.rs | 6 + native/core/src/local.rs | 64 ++++++++-- native/local/src/handle.rs | 18 ++- native/local/src/lib.rs | 113 ++++++++++++------ native/local/tests/bridge.rs | 22 +++- native/local/tests/execution.rs | 81 ++++++++++++- .../apache/comet/local/CometLocalExec.scala | 8 +- .../comet/local/CometLocalResultExec.scala | 24 ++-- .../apache/comet/local/CometLocalRule.scala | 7 +- .../comet/local/LocalQueryIterator.scala | 2 + .../local/CometLocalExecutionSuite.scala | 22 ++++ 13 files changed, 337 insertions(+), 98 deletions(-) diff --git a/docs/source/contributor-guide/local-execution.md b/docs/source/contributor-guide/local-execution.md index 5f04ffac910..7cb9ea9b1d4 100644 --- a/docs/source/contributor-guide/local-execution.md +++ b/docs/source/contributor-guide/local-execution.md @@ -33,6 +33,8 @@ Admission runs in `CometRule`, on the plan that Comet's ordinary conversion prod admitted when: - the application's `SparkContext` is local (`local[N]`), so the query runs in this process; +- `spark.executor.userClassPathFirst` is off, since a child-first task class loader would load + its own copy of the driver's result slots; - the plan is a batch query and not a subquery being prepared; - the converted plan is one native block, a tree of Comet native operators whose root carries the block's serialized plan, with no JVM input, broadcast or shuffle inside it; @@ -61,11 +63,17 @@ single partition. That partition's task creates the native graph: for partition `i`, so per-partition state such as a range slice is the same as with ordinary Comet execution. - `native/local` runs the partitions on Comet's Tokio runtime and delivers them in partition - order. At most `spark.executor.cores` partitions (N for `local[N]`) run at once, and the next - starts when the first is delivered, as Spark runs at most that many tasks. A partition's - operators register their memory consumers when it starts, so this also bounds how many share - the task's memory pool. The result therefore has Spark's row order. A partition's error is - reported when its turn comes. + order, so the result has Spark's row order. Each partition starts its operators on its own + task, so a start that waits, for instance for a memory grant, does not hold up the pull, and a + partition's error, including a failure to start, is reported when its turn comes. +- At most as many partitions run at once as Spark runs tasks, `spark.executor.cores` divided by + `spark.task.cpus`, and the next starts when the first is delivered. A partition's operators + register their memory consumers when it starts, so this also bounds how many share the task's + memory pool. Whenever another partition would start, the bound is further divided among the + executor's tasks running native plans, this one included: Comet's pools give each of those + tasks an equal share of Spark's memory, so a partition gets what a task would. Partitions + already running keep running, so a job that starts after a local query has started all its + partitions still leaves those partitions less than a task's share, until they end. - The running partitions do not wait for their turn to produce batches, so they all make progress, as Spark's tasks do. What they hold ahead of delivery is bounded by `spark.driver.maxResultSize`, the most a collect may bring to the driver; past it, each holds @@ -100,20 +108,24 @@ timer, which only fires once a worker is free. `CometLocalResultExec` sits above Spark's columnar-to-row transition. For `collect` it still runs the single task as a Spark job, which keeps cancellation, job groups and SQL metrics, but -the task copies its rows into a driver-side slot instead of returning an encoded task result. Spark would otherwise encode and compress every row in that one task and decode -them on the driver, all on one thread. While copying, the task counts uncompressed `UnsafeRow` -bytes and stops reading, which closes the native query, once they exceed -`spark.driver.maxResultSize`; the driver then fails the action. Spark counts compressed bytes, -so this check is the stricter one. On a non-local master the node uses Spark's ordinary -collect. - -Every other action reads the result as an RDD of the block's partitions: an iterator fetches -one partition at a time so that it holds at most one, a tail scans partitions from the last one, and a cache, -`df.rdd`, `foreachPartition`, `localCheckpoint` or PySpark's collect and `toLocalIterator` keep -the partitions for later reads. So `CometLocalResultExec` executes as the ordinary native block, -one Spark task per partition, for everything but `collect`. The partitions these actions -evaluate, and so the rows they return and the errors they raise, are Spark's, and a cached result -is scanned as many partitions as without local execution. +the task copies its rows into a driver-side slot instead of returning an encoded task result. +Spark would otherwise encode and compress every row in that one task and decode them on the +driver, all on one thread. While copying, the task counts uncompressed `UnsafeRow` bytes and +stops reading, which closes the native query, once they exceed `spark.driver.maxResultSize`; +the driver then fails the action. Spark instead counts each task's serialized, compressed +result, overhead included, so either can reach the limit first: rows that compress well reach +it here, and many small partitions reach Spark's. Classic PySpark's `collect()` goes through +the same path. On a non-local master the node uses Spark's ordinary collect. + +Every other action reads the result as an RDD of the block's partitions: an iterator, +PySpark's `toLocalIterator()` included, fetches one partition at a time so that it holds at +most one, a tail scans partitions from the last one, and a cache, `df.rdd`, `foreachPartition` +or `localCheckpoint` keep the partitions for later reads. So `CometLocalResultExec` executes as +the ordinary native block, one Spark task per partition, for everything but `collect`, prepared +as ordinary planning prepares it: with the columnar-to-row transition Comet picks, and as +Spark's `RangeExec` for a bare range. The partitions these actions evaluate, and so the rows +they return and the errors they raise, are Spark's, and a cached result is scanned as many +partitions as without local execution. ## Lifecycle and cancellation @@ -126,21 +138,27 @@ if the runtime stops it before that, the reader reports an error instead of a tr Closing a query drops its streams on the calling thread, as `releasePlan` does for an ordinary plan without JVM input: cancelling a Tokio task only takes effect once a worker is free, and every worker can be waiting, for instance in Spark's `acquireMemory` for memory those streams -hold. Tasks that DataFusion operators spawn themselves still end asynchronously. +hold. It then releases the context as `releasePlan` does, waiting briefly for the reservations +that tasks DataFusion operators spawned still hold, since they give them back only the next time +they yield. This applies however the query ends, including early on cancellation or the result +size limit. ## Testing -Build native code before the JVM suite, and never use Maven `-pl`: +Build native code, including the JNI library, before the JVM suite, and never use Maven +`-pl`: ```shell cd native cargo test -p datafusion-comet-local cd .. +make core ./mvnw test -Dsuites=org.apache.comet.local.CometLocalExecutionSuite ``` -The native tests cover partition-order delivery, the bound on running partitions, read-ahead, errors reported in a partition's turn, cancellation with busy workers, -the pull that waits on its own thread, and runtime shutdown. `CometLocalExecutionSuite` compares admitted queries +The native tests cover partition-order delivery, the bound on running partitions and its +changes, partition starts that wait, read-ahead, errors reported in a partition's turn, +cancellation with busy workers, the pull that waits on its own thread, and runtime shutdown. `CometLocalExecutionSuite` compares admitted queries with Spark's results and order, and covers ranges Spark would overflow, declined expressions, AQE, reported orderings, limits left to ordinary execution, iterators, tails, many partitions with sorts in a small memory pool, caching and other RDD consumers, the result size limit, job and task cancellation, and native errors. After each test it checks that no native query, imported Arrow memory or result slot is left. diff --git a/native/core/src/execution/jni_api.rs b/native/core/src/execution/jni_api.rs index c2cb9ab470b..2cac7f0a4a6 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -123,7 +123,7 @@ use tokio::sync::mpsc; use tokio::task::JoinHandle; // Re-exported for local execution, whose memory pool is the one `createPlan` creates. -pub(crate) use crate::execution::memory_pools::create_memory_pool; +pub(crate) use crate::execution::memory_pools::{create_memory_pool, running_native_tasks}; use crate::execution::memory_pools::{ overcommit, parse_memory_pool_config, MemoryPoolConfig, PlanMemoryPool, }; @@ -1650,7 +1650,7 @@ pub extern "system" fn Java_org_apache_comet_Native_releasePlan( /// Releases a plan: stops it, publishes its final metrics through `publish_metrics`, drops it, /// and waits briefly for its memory reservations to be returned. The context is freed even if a /// step fails. Core of `Native.releasePlan`. -fn release_plan( +pub(crate) fn release_plan( mut execution_context: Box, publish_metrics: &mut PublishMetrics, ) -> CometResult<()> { diff --git a/native/core/src/execution/memory_pools/mod.rs b/native/core/src/execution/memory_pools/mod.rs index 75339aa2558..7e93c02f788 100644 --- a/native/core/src/execution/memory_pools/mod.rs +++ b/native/core/src/execution/memory_pools/mod.rs @@ -117,6 +117,12 @@ fn create_pool( } } +/// The tasks of this executor running native plans, among which Comet's pools divide a task's +/// share of Spark's memory. +pub(crate) fn running_native_tasks() -> usize { + RunningTasks::of_executor().count() +} + /// The bytes that `pool` has recorded beyond what Spark granted it, which it carries as overcommit /// until Spark grants them or the pool frees memory; see [`SparkMemory`]. This looks through the /// wrappers that [`create_memory_pool`] puts around a Comet pool, and is zero for a pool that diff --git a/native/core/src/local.rs b/native/core/src/local.rs index f24ea70c38e..9081fa0fa59 100644 --- a/native/core/src/local.rs +++ b/native/core/src/local.rs @@ -33,22 +33,40 @@ use datafusion_comet_local::LocalQuery; use jni::objects::{JByteArray, JLongArray, JObject, JObjectArray, JString, ReleaseMode}; use jni::sys::{jboolean, jint, jlong, JNI_FALSE}; use jni::{Env, EnvUnowned}; +use log::warn; use parking_lot::Mutex; use crate::errors::{try_unwrap_or_throw, CometError, CometResult}; use crate::execution::jni_api::{ create_execution_context, create_memory_pool, get_runtime, prepare_output, - prepare_plan_settings, ExecutionContext, PlanInputs, PlanJvmRefs, + prepare_plan_settings, release_plan, running_native_tasks, ExecutionContext, PlanInputs, + PlanJvmRefs, }; use crate::execution::spark_config::{SparkConfig, SPARK_EXECUTOR_CORES}; struct Entry { query: QueryHandle, columns: usize, - // Holds the plan's memory pool and session until the query is closed. Dropped after the + /// Validates exported batches, as `spark.comet.debug.enabled` does for any Comet plan. + validate: bool, + // Holds the plan's memory pool and session until the query is closed. Released after the // query, which cancels the graph that reserves from the pool. Only the creating thread // uses it; the mutex makes the entry shareable with result pulls. - _context: Mutex>, + context: Mutex>>, +} + +impl Drop for Entry { + /// Releases the context as `Native.releasePlan` does, waiting briefly for the reservations + /// that the cancelled graph's tasks still hold, so that they are returned before the Spark + /// task can end. + fn drop(&mut self) { + self.query.cancel(); + if let Some(context) = self.context.get_mut().take() { + if let Err(error) = release_plan(context, &mut |_| Ok(())) { + warn!("Releasing a local query failed: {error}"); + } + } + } } // JNI carries numeric IDs instead of raw pointers, so close racing with a result pull @@ -94,9 +112,11 @@ pub extern "system" fn Java_org_apache_comet_local_NativeLocal_createQuery( } let start = Instant::now(); let settings = prepare_plan_settings(&env.convert_byte_array(serialized_spark_configs)?)?; - // Run as many partitions at once as Spark would run tasks, which also bounds how many - // operators share the memory pool. `local[N]` resolves to N here. - let max_running = settings.spark_config.get_usize(SPARK_EXECUTOR_CORES, 1); + // `local[N]` resolves to N cores here. + let slots = task_slots( + settings.spark_config.get_usize(SPARK_EXECUTOR_CORES, 1), + task_cpus as usize, + ); JVMClasses::init(env); let serialized_plan = env.convert_byte_array(serialized_query)?; let task_memory_manager = @@ -135,19 +155,24 @@ pub extern "system" fn Java_org_apache_comet_local_NativeLocal_createQuery( "Local output schema width mismatch".into(), )); } + // Run at most as many partitions at once as Spark runs tasks. Comet's pools split Spark's + // memory equally among the tasks running native plans, this one included, so with T of + // them the partitions each get what a task would only while at most slots / T run. let query = LocalQuery::new( plan, context.session_ctx.task_ctx(), - max_running, + slots, read_ahead_bytes.max(0) as usize, - ); + ) + .with_running_limit(Arc::new(move || slots / running_native_tasks().max(1))); let query = QueryHandle::start(query, &get_runtime())?; register( id, Entry { query, columns: columns as usize, - _context: Mutex::new(context), + validate: context.debug_native, + context: Mutex::new(Some(context)), }, )?; Ok(id) @@ -156,6 +181,11 @@ pub extern "system" fn Java_org_apache_comet_local_NativeLocal_createQuery( const MAX_PARTITIONS: jint = 1024; +/// The tasks Spark runs at once with `cores`, as `LocalSchedulerBackend` counts them. +fn task_slots(cores: usize, task_cpus: usize) -> usize { + (cores / task_cpus.max(1)).max(1) +} + /// One graph for the whole query. Each partition's operators are created by the planner for /// that partition, as a Spark task would create them, so per-partition state such as a range /// slice or a partition-dependent expression is correct. Partition `i` of the graph is the @@ -220,7 +250,7 @@ fn next(env: &mut Env, id: i64, arrays: JLongArray, schemas: JLongArray) -> Come match entry.query.poll(Duration::from_millis(50))? { QueryPoll::Pending => Ok(-2), QueryPoll::Finished => Ok(-1), - QueryPoll::Batch(batch) => prepare_output(env, arrays, schemas, batch, false), + QueryPoll::Batch(batch) => prepare_output(env, arrays, schemas, batch, entry.validate), } } @@ -261,3 +291,17 @@ pub extern "system" fn Java_org_apache_comet_local_NativeLocal_activeQueries( ) -> jlong { try_unwrap_or_throw(&e, |_| Ok(QUERIES.lock().len() as jlong)) } + +#[cfg(test)] +mod tests { + use super::task_slots; + + #[test] + fn task_slots_divide_cores_by_task_cpus() { + assert_eq!(task_slots(32, 16), 2); + assert_eq!(task_slots(8, 3), 2); + assert_eq!(task_slots(4, 1), 4); + assert_eq!(task_slots(2, 4), 1); + assert_eq!(task_slots(4, 0), 4); + } +} diff --git a/native/local/src/handle.rs b/native/local/src/handle.rs index fa1c63da60d..9ca0dd20c23 100644 --- a/native/local/src/handle.rs +++ b/native/local/src/handle.rs @@ -50,6 +50,9 @@ pub struct QueryHandle { reader: Mutex>, cancelled: AtomicBool, runtime: Handle, + /// Set when a reader is about to park, so that a test can cancel a parked reader. + #[cfg(test)] + parking: AtomicBool, } /// Wakes a reader parked in `QueryHandle::poll`. @@ -72,6 +75,8 @@ impl QueryHandle { reader: Mutex::new(None), cancelled: AtomicBool::new(false), runtime: runtime.clone(), + #[cfg(test)] + parking: AtomicBool::new(false), }) } @@ -104,7 +109,7 @@ impl QueryHandle { let deadline = Instant::now() + wait; let waker = Waker::from(Arc::new(Unpark(std::thread::current()))); let mut cx = Context::from_waker(&waker); - // Starting a partition executes its operators, which may spawn tasks. + // Ending the query drops its streams here, which may need the runtime. let _guard = self.runtime.enter(); loop { // Checked after the reader is registered, so a cancel either sees the reader and @@ -128,6 +133,8 @@ impl QueryHandle { if now >= deadline { return Ok(QueryPoll::Pending); } + #[cfg(test)] + self.parking.store(true, Ordering::Release); std::thread::park_timeout(deadline - now); } } @@ -188,7 +195,14 @@ mod tests { let thread = std::thread::spawn(move || { sender.send(reader.poll(Duration::from_secs(60))).unwrap(); }); - while lock(&query.reader).is_none() { + // Cancel only once the reader has passed its last check before parking, so that only + // the wakeup can end its wait. + let waiting = Instant::now(); + while !query.parking.load(Ordering::Acquire) { + assert!( + waiting.elapsed() < Duration::from_secs(5), + "the reader never parked" + ); std::thread::yield_now(); } assert!(query.poll(Duration::ZERO).is_err(), "a second reader"); diff --git a/native/local/src/lib.rs b/native/local/src/lib.rs index 33a4273246a..a9dbe8035b1 100644 --- a/native/local/src/lib.rs +++ b/native/local/src/lib.rs @@ -47,11 +47,15 @@ use tokio::task::JoinHandle; /// limit. const PARTITION_BUFFER: usize = 2; +/// How many partitions may run at once as of now. +pub type RunningLimit = Arc usize + Send + Sync>; + /// One execution's graph and resources. Consumed when execution starts; not a plan cache. pub struct LocalQuery { plan: Arc, context: Arc, max_running: usize, + running_limit: Option, read_ahead: usize, } @@ -64,6 +68,10 @@ impl LocalQuery { /// A partition's operators register their memory consumers when it starts, so this also /// bounds how many share the memory pool. /// + /// A partition starts its operators on its own task, so a start that waits, for instance to + /// register a memory consumer, does not hold up delivery. A failure to start is reported in + /// the partition's turn. + /// /// The running partitions do not wait for their turn: each runs ahead of delivery while the /// batches they all hold total less than `read_ahead` bytes, and up to a few batches beyond /// that. The root's partitions must not share an exchange: a partition that waits to be @@ -81,10 +89,19 @@ impl LocalQuery { plan, context, max_running: max_running.max(1), + running_limit: None, read_ahead, } } + /// Lowers the number of partitions that may run at once to what `limit` returns whenever + /// another partition would start, and to at least one. Partitions already running keep + /// running. + pub fn with_running_limit(mut self, limit: RunningLimit) -> Self { + self.running_limit = Some(limit); + self + } + /// Starts the first root partitions and exposes one streaming result. No result-wide /// collection is performed here. pub fn execute(self) -> Result { @@ -92,8 +109,7 @@ impl LocalQuery { DataFusionError::Execution("LocalQuery requires an active Tokio runtime".into()) })?; let schema = self.plan.schema(); - let stream = - in_partition_order(&self.plan, &self.context, self.max_running, self.read_ahead)?; + let stream = in_partition_order(&self); Ok(LocalQueryStream { stream: Some(stream), query: Some(self), @@ -102,31 +118,23 @@ impl LocalQuery { } } -fn in_partition_order( - plan: &Arc, - context: &Arc, - max_running: usize, - read_ahead: usize, -) -> Result { +fn in_partition_order(query: &LocalQuery) -> SendableRecordBatchStream { let mut order = PartitionOrder { - plan: Arc::clone(plan), - context: Arc::clone(context), + plan: Arc::clone(&query.plan), + context: Arc::clone(&query.context), runtime: tokio::runtime::Handle::current(), next: 0, - max_running, + max_running: query.max_running, + running_limit: query.running_limit.clone(), read_ahead: Arc::new(ReadAhead { - limit: read_ahead, + limit: query.read_ahead, held: AtomicUsize::new(0), delivered: Notify::new(), }), running: VecDeque::new(), }; - // Start the first partitions now, so that a failure to start them fails here. - order.start_partitions()?; - Ok(Box::pin(RecordBatchStreamAdapter::new( - plan.schema(), - order, - ))) + order.start_partitions(); + Box::pin(RecordBatchStreamAdapter::new(query.plan.schema(), order)) } /// What the running partitions hold ahead of delivery, shared by them. @@ -151,25 +159,48 @@ struct Partition { ended: Arc, /// The partition's stream. The task locks it only while polling it, so dropping the /// partition can take it whenever the task is waiting. - stream: Arc>>, + stream: Arc>, task: JoinHandle<()>, } +/// A partition's stream over its lifetime. +enum Slot { + /// The task is starting the partition's operators. + Starting, + Running(SendableRecordBatchStream), + /// The partition was dropped. A task that started the stream afterwards drops it. + Closed, +} + impl Partition { - fn spawn( - runtime: &tokio::runtime::Handle, - stream: SendableRecordBatchStream, - read_ahead: &Arc, - ) -> Self { + fn spawn(order: &PartitionOrder, partition: usize) -> Self { let (sender, batches) = tokio::sync::mpsc::unbounded_channel(); let queued = Arc::new(AtomicUsize::new(0)); let ended = Arc::new(AtomicBool::new(false)); - let stream = Arc::new(Mutex::new(Some(stream))); + let stream = Arc::new(Mutex::new(Slot::Starting)); let (task_queued, task_ended, polled) = (Arc::clone(&queued), Arc::clone(&ended), Arc::clone(&stream)); - let read_ahead = Arc::clone(read_ahead); - let task = runtime.spawn(async move { + let read_ahead = Arc::clone(&order.read_ahead); + let (plan, context) = (Arc::clone(&order.plan), Arc::clone(&order.context)); + let task = order.runtime.spawn(async move { let result = AssertUnwindSafe(async { + let started = plan.execute(partition, context); + drop(plan); + match started { + Ok(started) => { + let mut slot = lock(&polled); + if matches!(*slot, Slot::Closed) { + drop(slot); + drop(started); + return; + } + *slot = Slot::Running(started); + } + Err(error) => { + let _ = sender.send(Err(error)); + return; + } + } loop { // Wait before polling while this partition holds `PARTITION_BUFFER` batches // and the partitions together hold the read-ahead limit. The first partition @@ -183,9 +214,9 @@ impl Partition { } delivered.await; } - let next = poll_fn(|cx| match lock(&polled).as_mut() { - Some(stream) => stream.poll_next_unpin(cx), - None => Poll::Ready(None), + let next = poll_fn(|cx| match &mut *lock(&polled) { + Slot::Running(stream) => stream.poll_next_unpin(cx), + Slot::Starting | Slot::Closed => Poll::Ready(None), }) .await; match next { @@ -232,7 +263,7 @@ impl Drop for Partition { /// effect once a Tokio worker is free. See `QueryHandle::cancel`. fn drop(&mut self) { self.task.abort(); - let stream = lock(&self.stream).take(); + let stream = std::mem::replace(&mut *lock(&self.stream), Slot::Closed); let _ = std::panic::catch_unwind(AssertUnwindSafe(|| drop(stream))); } } @@ -247,21 +278,27 @@ struct PartitionOrder { /// The first partition not started yet. next: usize, max_running: usize, + running_limit: Option, read_ahead: Arc, /// Started partitions, the first of which is being delivered. running: VecDeque, } impl PartitionOrder { - fn start_partitions(&mut self) -> Result<()> { + fn start_partitions(&mut self) { let partitions = self.plan.output_partitioning().partition_count(); - while self.running.len() < self.max_running && self.next < partitions { - let stream = self.plan.execute(self.next, Arc::clone(&self.context))?; - let partition = Partition::spawn(&self.runtime, stream, &self.read_ahead); + if self.next >= partitions { + return; + } + let limit = self + .running_limit + .as_ref() + .map_or(self.max_running, |limit| limit().clamp(1, self.max_running)); + while self.running.len() < limit && self.next < partitions { + let partition = Partition::spawn(self, self.next); self.running.push_back(partition); self.next += 1; } - Ok(()) } fn fail(&mut self, error: DataFusionError) -> Poll>> { @@ -276,9 +313,7 @@ impl Stream for PartitionOrder { fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { loop { - if let Err(error) = self.start_partitions() { - return self.fail(error); - } + self.start_partitions(); let Some(partition) = self.running.front_mut() else { return Poll::Ready(None); }; diff --git a/native/local/tests/bridge.rs b/native/local/tests/bridge.rs index e3642a8195b..ab70ee2c719 100644 --- a/native/local/tests/bridge.rs +++ b/native/local/tests/bridge.rs @@ -167,17 +167,27 @@ fn worker_panic_is_an_error_not_successful_eof() { #[test] fn startup_failure_is_returned() { let runtime = runtime(); - let plan = Arc::new(ErrorExec::new()); - assert!(QueryHandle::start( + let query = QueryHandle::start( LocalQuery::new( - plan, + Arc::new(ErrorExec::new()), Arc::new(TaskContext::default()), usize::MAX, - usize::MAX + usize::MAX, ), - runtime.handle() + runtime.handle(), ) - .is_err()); + .unwrap(); + let deadline = Instant::now() + Duration::from_secs(10); + loop { + match query.poll(Duration::from_millis(50)) { + Err(error) => { + assert!(error.to_string().contains("ErrorExec")); + break; + } + Ok(QueryPoll::Pending) => assert!(Instant::now() < deadline, "no startup failure"), + Ok(_) => panic!("startup failure was lost"), + } + } } #[test] diff --git a/native/local/tests/execution.rs b/native/local/tests/execution.rs index 46b502db929..956b30a3ef5 100644 --- a/native/local/tests/execution.rs +++ b/native/local/tests/execution.rs @@ -17,8 +17,9 @@ use std::collections::BTreeMap; use std::fmt; +use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex, Weak}; -use std::time::Duration; +use std::time::{Duration, Instant}; use datafusion::arrow::array::Int64Array; use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; @@ -94,6 +95,9 @@ struct Observations { struct ObserveExec { input: Arc, observations: Arc, + /// How long starting a partition blocks its caller, as a memory consumer's registration + /// can while the pool waits for Spark. + start_delay: Duration, } impl DisplayAs for ObserveExec { @@ -131,6 +135,7 @@ impl ExecutionPlan for ObserveExec { Ok(Arc::new(Self { input: Arc::clone(&children[0]), observations: Arc::clone(&self.observations), + start_delay: self.start_delay, })) } @@ -149,6 +154,7 @@ impl ExecutionPlan for ObserveExec { partition: usize, ctx: Arc, ) -> Result { + std::thread::sleep(self.start_delay); self.observations.executions.lock().unwrap().push(partition); let observations = Arc::clone(&self.observations); let stream = self.input.execute(partition, ctx)?.map_ok(move |batch| { @@ -168,11 +174,19 @@ impl ExecutionPlan for ObserveExec { } fn observe(input: Arc) -> (Arc, Arc) { + observe_starting_in(input, Duration::ZERO) +} + +fn observe_starting_in( + input: Arc, + start_delay: Duration, +) -> (Arc, Arc) { let observations = Arc::new(Observations::default()); ( Arc::new(ObserveExec { input, observations: Arc::clone(&observations), + start_delay, }), observations, ) @@ -310,17 +324,39 @@ async fn cancelling_one_query_does_not_cancel_another() { } #[tokio::test] -async fn startup_error_releases_query() { +async fn startup_error_is_reported_in_its_turn_and_releases_query() { let plan = Arc::new(ErrorExec::new()); let refs = Arc::downgrade(&plan); - let error = LocalQuery::new(plan, context(), usize::MAX, usize::MAX) + let mut stream = LocalQuery::new(plan, context(), usize::MAX, usize::MAX) .execute() - .err() .unwrap(); + let error = timeout(DEADLINE, stream.next()) + .await + .unwrap() + .unwrap() + .unwrap_err(); assert!(error.to_string().contains("ErrorExec")); + drop(stream); wait_released(&refs).await; } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn a_partition_start_that_waits_does_not_hold_up_the_reader() { + let (plan, observations) = observe_starting_in(ordered_input(2, 1), Duration::from_millis(500)); + let started = Instant::now(); + let mut stream = LocalQuery::new(plan, context(), usize::MAX, usize::MAX) + .execute() + .unwrap(); + assert!(futures::poll!(stream.next()).is_pending()); + assert!(started.elapsed() < Duration::from_millis(250)); + let batches: Vec = timeout(DEADLINE, stream.try_collect()) + .await + .unwrap() + .unwrap(); + assert_eq!(batches.len(), 2); + assert_executed_once(&observations, 2); +} + #[test] fn missing_runtime_is_an_error_not_a_panic() { let plan = Arc::new(EmptyExec::new(schema())); @@ -446,7 +482,7 @@ async fn partition_order_runs_at_most_the_given_number_of_partitions() { .execute() .unwrap(); // The first two partitions start with the query; the rest wait for their turn. - assert_eq!(executed(), vec![0, 1]); + wait_started(&observations, 2).await; for _ in 0..3 { timeout(DEADLINE, stream.next()) .await @@ -461,7 +497,7 @@ async fn partition_order_runs_at_most_the_given_number_of_partitions() { .unwrap() .unwrap() .unwrap(); - assert_eq!(executed(), vec![0, 1, 2]); + wait_started(&observations, 3).await; let rest: Vec = timeout(DEADLINE, stream.try_collect()) .await .unwrap() @@ -506,3 +542,36 @@ async fn partition_order_holds_a_few_batches_beyond_the_limit() { assert_eq!(produced(&observations, 1), 2); stream.cancel(); } + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn partition_order_follows_a_changing_running_limit() { + let (plan, observations) = observe(ordered_input(6, 1)); + let limit = Arc::new(AtomicUsize::new(1)); + let current = Arc::clone(&limit); + let mut stream = LocalQuery::new(plan, context(), 3, usize::MAX) + .with_running_limit(Arc::new(move || current.load(Ordering::Acquire))) + .execute() + .unwrap(); + let executed = || observations.executions.lock().unwrap().len(); + wait_started(&observations, 1).await; + // A limit of zero still runs one partition at a time. + limit.store(0, Ordering::Release); + for partition in 0..2 { + timeout(DEADLINE, stream.next()) + .await + .unwrap() + .unwrap() + .unwrap(); + // The next partition starts only once this one ends. + tokio::time::sleep(Duration::from_millis(50)).await; + assert_eq!(executed(), partition + 1, "started beyond the limit"); + } + // A higher limit applies to the partitions started from then on. + limit.store(usize::MAX, Ordering::Release); + let rest: Vec = timeout(DEADLINE, stream.try_collect()) + .await + .unwrap() + .unwrap(); + assert_eq!(rest.len(), 4); + assert_executed_once(&observations, 6); +} diff --git a/spark/src/main/scala/org/apache/comet/local/CometLocalExec.scala b/spark/src/main/scala/org/apache/comet/local/CometLocalExec.scala index a9cb7a4f648..40297dee2be 100644 --- a/spark/src/main/scala/org/apache/comet/local/CometLocalExec.scala +++ b/spark/src/main/scala/org/apache/comet/local/CometLocalExec.scala @@ -27,10 +27,16 @@ import org.apache.spark.sql.catalyst.plans.QueryPlan import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, UnknownPartitioning} import org.apache.spark.sql.execution.{ColumnarToRowExec, LeafExecNode, SparkPlan} import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.types.DataType import org.apache.spark.sql.vectorized.ColumnarBatch /** What the result task needs to create the native graph: no plan objects. */ -private[local] case class LocalQuerySpec(plan: Array[Byte], partitions: Int, columns: Int) +private[local] case class LocalQuerySpec( + plan: Array[Byte], + partitions: Int, + outputTypes: Seq[DataType]) { + def columns: Int = outputTypes.size +} /** * A whole query run as one native graph by the single Spark task of its result. `block` is the diff --git a/spark/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala b/spark/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala index 545b39400ec..6c0fa10bf52 100644 --- a/spark/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala +++ b/spark/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala @@ -31,6 +31,8 @@ import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder, UnsafeRo import org.apache.spark.sql.catalyst.plans.physical.Partitioning import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, UnaryExecNode} +import org.apache.comet.rules.EliminateRedundantTransitions + /** * Row result boundary of a local query. Spark's collect encodes and compresses every row of each * result partition in its task, then decodes them on the driver. A local collect instead runs the @@ -38,9 +40,10 @@ import org.apache.spark.sql.execution.{ColumnarToRowExec, SparkPlan, UnaryExecNo * SQL metrics), which hands copied rows to the driver in this JVM, since admission requires an * in-process local master. * - * Only `collect` runs that way. Everything else reads the result as an RDD: a take requests - * partitions in growing batches, an iterator fetches one partition at a time, a tail scans from - * the last one, and a cache, `df.rdd` or PySpark's collect keep its partitions. So the node + * Only `collect` runs that way, including classic PySpark's `collect()`, which goes through + * `executeCollect`. Everything else reads the result as an RDD: a take requests partitions in + * growing batches, an iterator (PySpark's `toLocalIterator()` too) fetches one partition at a + * time, a tail scans from the last one, and a cache or `df.rdd` keep its partitions. So the node * executes as the ordinary native block otherwise, one Spark task per partition, and reports the * block's partitioning and ordering. */ @@ -55,12 +58,16 @@ case class CometLocalResultExec(child: SparkPlan) extends UnaryExecNode { if (sparkContext.isLocal) LocalResultHandoff.collect(child.execute(), maxResultSize) else super.executeCollect() - private def ordinary: SparkPlan = + // The block as ordinary planning prepares it, which runs after this node is created: the + // columnar-to-row transition Comet picks, or Spark's RangeExec for a bare range. + @transient private[local] lazy val ordinary: SparkPlan = child - .collectFirst { case local: CometLocalExec => ColumnarToRowExec(local.block) } + .collectFirst { case local: CometLocalExec => + EliminateRedundantTransitions(session)(ColumnarToRowExec(local.block)) + } .getOrElse(child) - // Spark enforces this on serialized task results; this boundary enforces it while copying. + // Spark enforces this on serialized task results; this boundary enforces it on copied rows. private def maxResultSize: Long = sparkContext.getConf.getSizeAsBytes("spark.driver.maxResultSize", "1g") @@ -84,8 +91,9 @@ private[local] object LocalResultHandoff { /** * Runs the single result partition of `rdd` and returns its rows. Stops reading, which closes * the native query, once the copied rows exceed `maxResultSize` bytes (unlimited if not - * positive). Unlike Spark's compressed size, this counts uncompressed UnsafeRow bytes, so it is - * the more conservative of the two. + * positive). This counts the uncompressed UnsafeRow bytes of the rows, while Spark counts each + * task's serialized, compressed result, overhead included. Either can reach the limit first: + * rows that compress well reach it here, and many small partitions reach Spark's. */ def collect(rdd: RDD[InternalRow], maxResultSize: Long): Array[InternalRow] = { require(rdd.getNumPartitions == 1, "Local results have exactly one partition") diff --git a/spark/src/main/scala/org/apache/comet/local/CometLocalRule.scala b/spark/src/main/scala/org/apache/comet/local/CometLocalRule.scala index 2e64994042d..59180453f79 100644 --- a/spark/src/main/scala/org/apache/comet/local/CometLocalRule.scala +++ b/spark/src/main/scala/org/apache/comet/local/CometLocalRule.scala @@ -70,7 +70,10 @@ private[comet] object CometLocalRule { convert(plan) match { case block: CometNativeExec if admits(block) => val partitions = block.collectLeaves().head.outputPartitioning.numPartitions - val spec = LocalQuerySpec(block.serializedPlanOpt.plan.get, partitions, block.output.size) + val spec = LocalQuerySpec( + block.serializedPlanOpt.plan.get, + partitions, + block.output.map(_.dataType)) CometLocalResultExec( CometLocalExec(block.output, spec, resultOrdering(block, partitions), block)) case converted => converted @@ -95,6 +98,8 @@ private[comet] object CometLocalRule { private def environmentAdmits(session: SparkSession, plan: SparkPlan): Boolean = { session.sparkContext.isLocal && + // A child-first task loader would load its own copy of the driver's result slots. + !session.sparkContext.getConf.getBoolean("spark.executor.userClassPathFirst", false) && !ShimCometStreaming.isStreamingPlan(plan) && // As in CometRule's plan-only reporting, this identifies subquery preparation. !Thread diff --git a/spark/src/main/scala/org/apache/comet/local/LocalQueryIterator.scala b/spark/src/main/scala/org/apache/comet/local/LocalQueryIterator.scala index 6dc01a35b49..21f27a1d740 100644 --- a/spark/src/main/scala/org/apache/comet/local/LocalQueryIterator.scala +++ b/spark/src/main/scala/org/apache/comet/local/LocalQueryIterator.scala @@ -37,6 +37,7 @@ private[local] class LocalQueryIterator(spec: LocalQuerySpec, context: TaskConte private var closed = false private var pending: ColumnarBatch = _ private var previous: ColumnarBatch = _ + private val checkTimestampLabels = CometConf.COMET_STRICT_TIMESTAMP_LABELS.get() // Shared by every native plan of the task, and released when the task completes. private val taskMemory = CometExecIterator.taskMemory(context, CometExec.newIterId) @@ -108,6 +109,7 @@ private[local] class LocalQueryIterator(spec: LocalQuerySpec, context: TaskConte }) .orNull if (pending == null) close() + else if (checkTimestampLabels) util.checkTimestampLabels(pending, spec.outputTypes) pending != null } catch { case failure: Throwable => diff --git a/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala b/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala index 4d85b95d351..910cc79af37 100644 --- a/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala +++ b/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala @@ -29,6 +29,7 @@ import org.scalatest.funsuite.AnyFunSuite import org.apache.spark.{SparkEnv, SparkException, TaskContext} import org.apache.spark.sql.{DataFrame, Row, SparkSession} +import org.apache.spark.sql.execution.{ColumnarToRowExec, RangeExec, SparkPlan} import org.apache.spark.sql.execution.columnar.InMemoryRelation import org.apache.spark.sql.execution.exchange.ShuffleExchangeLike import org.apache.spark.sql.expressions.Window @@ -227,6 +228,27 @@ class CometLocalExecutionSuite assert( query.rdd.mapPartitions(rows => Iterator(rows.size)).collect().toSeq == Seq.fill(6)(50)) assert(query.localCheckpoint().rdd.getNumPartitions == 6) + + def ordinary(df: DataFrame): SparkPlan = + df.queryExecution.executedPlan.collectFirst { case r: CometLocalResultExec => + r.ordinary + }.get + // Prepared as ordinary planning prepares the block: a bare range as Spark's RangeExec, + // which no native operator consumes, and otherwise with Comet's columnar-to-row transition. + val range = spark.range(0, 300, 1, 6).toDF() + assert(localNodes(range).size == 1) + assert(ordinary(range).isInstanceOf[RangeExec], ordinary(range)) + assert( + range.rdd.mapPartitions(rows => Iterator(rows.size)).collect().toSeq == Seq.fill(6)(50)) + assert(!ordinary(query).isInstanceOf[ColumnarToRowExec], ordinary(query)) + } + + test("output validation and timestamp label checks apply to local output") { + withConf(CometConf.COMET_DEBUG_ENABLED.key, "true") { + withConf(CometConf.COMET_STRICT_TIMESTAMP_LABELS.key, "true") { + checkLocal(spark.range(0, 100, 1, 3).selectExpr("id", "CAST(id AS TIMESTAMP) AS ts")) + } + } } test("only an ordering that holds across concatenated partitions is reported") { From 200906ea147bf26e3cf5d27e1a787dd797494667 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Sun, 1 Nov 2026 13:36:50 -0800 Subject: [PATCH 12/12] perf: plan each local partition when it starts A local query planned every partition before the first one started, so each partition's plan, and whatever it copied from the serialized plan such as a string literal, was held for the whole query. LocalQuery can now take a function that plans and starts a root partition on its own task, and the JNI side uses it, so only the running partitions hold their plans, as only running Spark tasks do. Partition 0 is still planned up front for the output schema and to reject a plan before the query starts. --- .../contributor-guide/local-execution.md | 6 +- native/core/src/execution/planner.rs | 10 +++ native/core/src/local.rs | 84 +++++++++++++------ native/local/src/lib.rs | 56 ++++++++++--- native/local/tests/execution.rs | 40 ++++++++- 5 files changed, 154 insertions(+), 42 deletions(-) diff --git a/docs/source/contributor-guide/local-execution.md b/docs/source/contributor-guide/local-execution.md index 7cb9ea9b1d4..86539c23270 100644 --- a/docs/source/contributor-guide/local-execution.md +++ b/docs/source/contributor-guide/local-execution.md @@ -61,7 +61,11 @@ single partition. That partition's task creates the native graph: - The block's operators are created once per partition, by the planner for that partition, as a Spark task would create them. Partition `i` of the graph is the output of the operators created for partition `i`, so per-partition state such as a range slice is the same as with ordinary - Comet execution. + Comet execution. A partition's operators are created when it starts, on its own task, as a Spark + task creates its plan only when it runs, so only the running partitions hold operators and + what they copy from the serialized plan, such as a string literal. Partition 0 is also planned + up front, for the output schema and so that a plan the planner rejects fails before the query + starts. - `native/local` runs the partitions on Comet's Tokio runtime and delivers them in partition order, so the result has Spark's row order. Each partition starts its operators on its own task, so a start that waits, for instance for a memory grant, does not hold up the pull, and a diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 10cef2334f1..c5bf8e49bfb 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -232,6 +232,16 @@ impl PhysicalPlanner { } } + /// This planner's settings for `partition`, with its own query context registry, as + /// `ExecutionContext::planner` would create it. The SQL text pool is shared. + pub(crate) fn for_partition(&self, partition: i32) -> Self { + Self { + partition, + query_context_registry: datafusion_comet_spark_expr::create_query_context_map(), + ..self.clone() + } + } + /// Load the SQL text pool from the root operator of the plan about to be planned. Must be /// called with the *root* operator: the JVM only populates the pool there, and /// `QueryContext.sql_text_idx` values are indices into it. diff --git a/native/core/src/local.rs b/native/core/src/local.rs index 9081fa0fa59..8135468a855 100644 --- a/native/core/src/local.rs +++ b/native/core/src/local.rs @@ -25,11 +25,13 @@ use std::sync::atomic::{AtomicI64, Ordering}; use std::sync::{Arc, LazyLock}; use std::time::{Duration, Instant}; -use datafusion::physical_plan::union::UnionExec; +use arrow::datatypes::SchemaRef; +use datafusion::common::DataFusionError; use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use datafusion_comet_jni_bridge::{jni_new_global_ref, JVMClasses}; use datafusion_comet_local::handle::{QueryHandle, QueryPoll}; -use datafusion_comet_local::LocalQuery; +use datafusion_comet_local::{LocalQuery, PartitionStart}; +use datafusion_comet_proto::spark_operator::Operator; use jni::objects::{JByteArray, JLongArray, JObject, JObjectArray, JString, ReleaseMode}; use jni::sys::{jboolean, jint, jlong, JNI_FALSE}; use jni::{Env, EnvUnowned}; @@ -42,6 +44,7 @@ use crate::execution::jni_api::{ prepare_plan_settings, release_plan, running_native_tasks, ExecutionContext, PlanInputs, PlanJvmRefs, }; +use crate::execution::planner::PhysicalPlanner; use crate::execution::spark_config::{SparkConfig, SPARK_EXECUTOR_CORES}; struct Entry { @@ -149,8 +152,8 @@ pub extern "system" fn Java_org_apache_comet_local_NativeLocal_createQuery( let context = create_execution_context(inputs, |config| { create_memory_pool(config, task_memory_manager, task_attempt_id) })?; - let plan = local_plan(&context, partitions)?; - if plan.schema().fields().len() != columns as usize { + let (schema, start) = local_plan(&context, partitions as usize)?; + if schema.fields().len() != columns as usize { return Err(CometError::Internal( "Local output schema width mismatch".into(), )); @@ -158,8 +161,10 @@ pub extern "system" fn Java_org_apache_comet_local_NativeLocal_createQuery( // Run at most as many partitions at once as Spark runs tasks. Comet's pools split Spark's // memory equally among the tasks running native plans, this one included, so with T of // them the partitions each get what a task would only while at most slots / T run. - let query = LocalQuery::new( - plan, + let query = LocalQuery::from_partitions( + schema, + partitions as usize, + start, context.session_ctx.task_ctx(), slots, read_ahead_bytes.max(0) as usize, @@ -186,29 +191,54 @@ fn task_slots(cores: usize, task_cpus: usize) -> usize { (cores / task_cpus.max(1)).max(1) } -/// One graph for the whole query. Each partition's operators are created by the planner for -/// that partition, as a Spark task would create them, so per-partition state such as a range -/// slice or a partition-dependent expression is correct. Partition `i` of the graph is the -/// single output partition of the operators created for partition `i`. -fn local_plan(context: &ExecutionContext, partitions: jint) -> CometResult> { - let mut plans = Vec::with_capacity(partitions as usize); - for partition in 0..partitions { - let (scans, shuffle_scans, root) = context.planner(partition).create_plan( - &context.spark_plan, - &mut vec![], - partitions as usize, - )?; - if !scans.is_empty() - || !shuffle_scans.is_empty() - || root.native_plan.output_partitioning().partition_count() != 1 - { - return Err(CometError::Internal( - "Local query plans must have no JVM input and one output partition".into(), - )); +/// One graph for the whole query, planned a partition at a time. Each partition's operators are +/// created by the planner for that partition, as a Spark task would create them, so per-partition +/// state such as a range slice or a partition-dependent expression is correct. A partition is +/// planned only when it starts, as a Spark task plans only when it runs, so the plans held at once, +/// and whatever each copies from the serialized plan, are those of the running partitions. +/// +/// Partition 0 is planned here too, for the output schema and so that a plan the planner rejects +/// fails the query before it starts. +fn local_plan( + context: &ExecutionContext, + partitions: usize, +) -> CometResult<(SchemaRef, PartitionStart)> { + let planner = context.planner(0); + let spark_plan = Arc::new(context.spark_plan.clone()); + let schema = plan_partition(&planner, &spark_plan, 0, partitions)?.schema(); + let expected = Arc::clone(&schema); + let start: PartitionStart = Arc::new(move |partition, task_context| { + let plan = plan_partition(&planner, &spark_plan, partition, partitions)?; + if plan.schema() != expected { + return Err(DataFusionError::Internal(format!( + "Local query partition {partition} has a different output schema" + ))); } - plans.push(Arc::clone(&root.native_plan)); + plan.execute(0, task_context) + }); + Ok((schema, start)) +} + +/// The operators of `partition`, whose single output partition is that partition of the graph. +fn plan_partition( + planner: &PhysicalPlanner, + spark_plan: &Operator, + partition: usize, + partitions: usize, +) -> CometResult> { + let (scans, shuffle_scans, root) = + planner + .for_partition(partition as i32) + .create_plan(spark_plan, &mut vec![], partitions)?; + if !scans.is_empty() + || !shuffle_scans.is_empty() + || root.native_plan.output_partitioning().partition_count() != 1 + { + return Err(CometError::Internal( + "Local query plans must have no JVM input and one output partition".into(), + )); } - Ok(UnionExec::try_new(plans)?) + Ok(Arc::clone(&root.native_plan)) } fn next_id() -> CometResult { diff --git a/native/local/src/lib.rs b/native/local/src/lib.rs index a9dbe8035b1..dee31918320 100644 --- a/native/local/src/lib.rs +++ b/native/local/src/lib.rs @@ -17,9 +17,9 @@ //! Query-scoped execution, independent of Spark tasks and JNI. //! -//! A caller supplies a fresh, fully planned graph for each execution and receives one -//! result stream. The caller owns the Tokio runtime and must keep it running until -//! cleanup finishes. Spark admission and JNI live outside this crate. +//! A caller supplies a fresh graph for each execution, or the means to plan each of its +//! partitions, and receives one result stream. The caller owns the Tokio runtime and must keep +//! it running until cleanup finishes. Spark admission and JNI live outside this crate. pub mod handle; @@ -50,9 +50,15 @@ const PARTITION_BUFFER: usize = 2; /// How many partitions may run at once as of now. pub type RunningLimit = Arc usize + Send + Sync>; +/// Starts a root partition: its stream, with the query's task context. +pub type PartitionStart = + Arc) -> Result + Send + Sync>; + /// One execution's graph and resources. Consumed when execution starts; not a plan cache. pub struct LocalQuery { - plan: Arc, + schema: SchemaRef, + partitions: usize, + start: PartitionStart, context: Arc, max_running: usize, running_limit: Option, @@ -84,9 +90,28 @@ impl LocalQuery { context: Arc, max_running: usize, read_ahead: usize, + ) -> Self { + let (schema, partitions) = (plan.schema(), plan.output_partitioning().partition_count()); + let start: PartitionStart = + Arc::new(move |partition, context| plan.execute(partition, context)); + Self::from_partitions(schema, partitions, start, context, max_running, read_ahead) + } + + /// As [`Self::new`], for a graph whose root partitions `start` plans as each one starts, so + /// that only the running partitions hold their operators. Each stream must have `schema`. + /// `start` runs on a Tokio worker of the caller's runtime, so it must not block on it. + pub fn from_partitions( + schema: SchemaRef, + partitions: usize, + start: PartitionStart, + context: Arc, + max_running: usize, + read_ahead: usize, ) -> Self { Self { - plan, + schema, + partitions, + start, context, max_running: max_running.max(1), running_limit: None, @@ -108,7 +133,7 @@ impl LocalQuery { tokio::runtime::Handle::try_current().map_err(|_| { DataFusionError::Execution("LocalQuery requires an active Tokio runtime".into()) })?; - let schema = self.plan.schema(); + let schema = Arc::clone(&self.schema); let stream = in_partition_order(&self); Ok(LocalQueryStream { stream: Some(stream), @@ -120,7 +145,8 @@ impl LocalQuery { fn in_partition_order(query: &LocalQuery) -> SendableRecordBatchStream { let mut order = PartitionOrder { - plan: Arc::clone(&query.plan), + partitions: query.partitions, + start: Arc::clone(&query.start), context: Arc::clone(&query.context), runtime: tokio::runtime::Handle::current(), next: 0, @@ -134,7 +160,10 @@ fn in_partition_order(query: &LocalQuery) -> SendableRecordBatchStream { running: VecDeque::new(), }; order.start_partitions(); - Box::pin(RecordBatchStreamAdapter::new(query.plan.schema(), order)) + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&query.schema), + order, + )) } /// What the running partitions hold ahead of delivery, shared by them. @@ -181,11 +210,11 @@ impl Partition { let (task_queued, task_ended, polled) = (Arc::clone(&queued), Arc::clone(&ended), Arc::clone(&stream)); let read_ahead = Arc::clone(&order.read_ahead); - let (plan, context) = (Arc::clone(&order.plan), Arc::clone(&order.context)); + let (start, context) = (Arc::clone(&order.start), Arc::clone(&order.context)); let task = order.runtime.spawn(async move { let result = AssertUnwindSafe(async { - let started = plan.execute(partition, context); - drop(plan); + let started = start(partition, context); + drop(start); match started { Ok(started) => { let mut slot = lock(&polled); @@ -272,7 +301,8 @@ impl Drop for Partition { /// Spark's `executeTake` evaluates a later partition only when it needs that partition's rows: /// a consumer that stops early never sees the failure of a partition it did not reach. struct PartitionOrder { - plan: Arc, + partitions: usize, + start: PartitionStart, context: Arc, runtime: tokio::runtime::Handle, /// The first partition not started yet. @@ -286,7 +316,7 @@ struct PartitionOrder { impl PartitionOrder { fn start_partitions(&mut self) { - let partitions = self.plan.output_partitioning().partition_count(); + let partitions = self.partitions; if self.next >= partitions { return; } diff --git a/native/local/tests/execution.rs b/native/local/tests/execution.rs index 956b30a3ef5..4cf4660b457 100644 --- a/native/local/tests/execution.rs +++ b/native/local/tests/execution.rs @@ -37,7 +37,7 @@ use datafusion::physical_plan::{ ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, }; -use datafusion_comet_local::LocalQuery; +use datafusion_comet_local::{LocalQuery, PartitionStart}; use futures::{StreamExt, TryStreamExt}; use tokio::time::timeout; @@ -506,6 +506,44 @@ async fn partition_order_runs_at_most_the_given_number_of_partitions() { assert_executed_once(&observations, 6); } +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn partitions_from_a_start_function_are_planned_as_they_start() { + let planned = Arc::new(Mutex::new(Vec::new())); + let (plans, source) = (Arc::clone(&planned), ordered_input(6, 1)); + let start: PartitionStart = Arc::new(move |partition, context| { + plans.lock().unwrap().push(partition); + source.execute(partition, context) + }); + let mut stream = LocalQuery::from_partitions(schema(), 6, start, context(), 2, usize::MAX) + .execute() + .unwrap(); + let planned_now = || planned.lock().unwrap().clone(); + timeout(DEADLINE, async { + while planned_now().len() < 2 { + tokio::task::yield_now().await; + } + }) + .await + .expect("the first two partitions were not planned"); + timeout(DEADLINE, stream.next()) + .await + .unwrap() + .unwrap() + .unwrap(); + // The third partition is planned only once the first has been delivered. + let mut first = planned_now(); + first.sort_unstable(); + assert_eq!(first, vec![0, 1]); + let rest: Vec = timeout(DEADLINE, stream.try_collect()) + .await + .unwrap() + .unwrap(); + assert_eq!(rest.len(), 5); + let mut all = planned_now(); + all.sort_unstable(); + assert_eq!(all, (0..6).collect::>()); +} + /// A first partition that never delivers, and a second one of ten batches. fn behind_a_blocked_partition() -> (Arc, Arc) { let blocking = Arc::new(BlockingExec::new(schema(), 1));