diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 381df39aa95..bbefec1994b 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -584,6 +584,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 373e530dbfa..c3759d973a7 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -250,6 +250,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/docs/source/contributor-guide/ffi.md b/docs/source/contributor-guide/ffi.md index 03a9c3a0b9d..2ac96fed7cb 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 @@ -265,6 +271,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/index.md b/docs/source/contributor-guide/index.md index 53acfa1be69..eade022fbbf 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..86539c23270 --- /dev/null +++ b/docs/source/contributor-guide/local-execution.md @@ -0,0 +1,168 @@ + + +# 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, 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. + +## 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; +- `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; +- 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`, `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 + +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. 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 + 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 + 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. + +## Result delivery + +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. 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 +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 + +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 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 +every worker can be waiting, for instance in Spark's `acquireMemory` for memory those streams +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, 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 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/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md index 4e7248f2ba3..26899442626 100644 --- a/docs/source/contributor-guide/memory_management.md +++ b/docs/source/contributor-guide/memory_management.md @@ -613,6 +613,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/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 673325aedfe..54bdebf3ed2 100644 --- a/native/Cargo.lock +++ b/native/Cargo.lock @@ -1998,6 +1998,7 @@ dependencies = [ "datafusion", "datafusion-comet-common", "datafusion-comet-jni-bridge", + "datafusion-comet-local", "datafusion-comet-operators", "datafusion-comet-proto", "datafusion-comet-shuffle", @@ -2081,6 +2082,15 @@ dependencies = [ "thiserror", ] +[[package]] +name = "datafusion-comet-local" +version = "1.3.0" +dependencies = [ + "datafusion", + "futures", + "tokio", +] + [[package]] name = "datafusion-comet-operators" version = "1.3.0" diff --git a/native/Cargo.toml b/native/Cargo.toml index 6f3b771e8ee..af293e4fe4b 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 254e0740e97..f80f1c1d8a4 100644 --- a/native/core/Cargo.toml +++ b/native/core/Cargo.toml @@ -66,6 +66,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 50df7a786b0..455f151ebe4 100644 --- a/native/core/src/execution/jni_api.rs +++ b/native/core/src/execution/jni_api.rs @@ -123,8 +123,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, running_native_tasks}; 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::{ @@ -531,7 +533,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 @@ -720,9 +722,23 @@ 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()) + .with_output_stream(self.output_stream.clone()) + } +} + /// What `createPlan` takes from the Spark configs. -struct PlanSettings { - spark_config: HashMap, +pub(crate) struct PlanSettings { + pub(crate) spark_config: HashMap, debug_native: bool, strict_timestamp_labels: bool, explain_native: bool, @@ -733,7 +749,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(); @@ -756,34 +772,35 @@ 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, - off_heap_size: 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) off_heap_size: 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 @@ -792,7 +809,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> { @@ -1198,7 +1215,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, @@ -1530,21 +1547,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()) - .with_output_stream(exec_context.output_stream.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(), @@ -1643,7 +1651,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/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/lib.rs b/native/core/src/lib.rs index 4f468409466..df6fb79e066 100644 --- a/native/core/src/lib.rs +++ b/native/core/src/lib.rs @@ -61,6 +61,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..8135468a855 --- /dev/null +++ b/native/core/src/local.rs @@ -0,0 +1,337 @@ +// 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 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, 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}; +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, release_plan, running_native_tasks, ExecutionContext, PlanInputs, + PlanJvmRefs, +}; +use crate::execution::planner::PhysicalPlanner; +use crate::execution::spark_config::{SparkConfig, SPARK_EXECUTOR_CORES}; + +struct Entry { + query: QueryHandle, + columns: usize, + /// 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>>, +} + +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 +// 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, + off_heap_size: 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 { + 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)?)?; + // `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 = + 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, + off_heap_size, + 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 (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(), + )); + } + // 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::from_partitions( + schema, + partitions as usize, + start, + context.session_ctx.task_ctx(), + 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, + validate: context.debug_native, + context: Mutex::new(Some(context)), + }, + )?; + Ok(id) + }) +} + +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, 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" + ))); + } + 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(Arc::clone(&root.native_plan)) +} + +fn next_id() -> CometResult { + // 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<()> { + 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, entry.validate), + } +} + +/// -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)) +} + +#[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/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..9ca0dd20c23 --- /dev/null +++ b/native/local/src/handle.rs @@ -0,0 +1,217 @@ +// 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. + +//! 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::{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::StreamExt; +use tokio::runtime::Handle; + +use crate::{lock, LocalQuery, LocalQueryStream}; + +pub enum QueryPoll { + Batch(RecordBatch), + Pending, + Finished, +} + +/// 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 { + /// 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, + 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`. +struct Unpark(Thread); + +impl Wake for Unpark { + fn wake(self: Arc) { + self.0.unpark(); + } +} + +impl QueryHandle { + pub fn start(query: LocalQuery, runtime: &Handle) -> Result { + let stream = { + let _guard = runtime.enter(); + query.execute()? + }; + Ok(Self { + stream: Mutex::new(Some(stream)), + reader: Mutex::new(None), + cancelled: AtomicBool::new(false), + runtime: runtime.clone(), + #[cfg(test)] + parking: AtomicBool::new(false), + }) + } + + /// A bounded wait lets a JVM caller check Spark task interruption between polls. + /// Only one reader may call this at a time. + pub fn poll(&self, wait: Duration) -> Result { + 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); + // 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 + // 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); + } + #[cfg(test)] + self.parking.store(true, Ordering::Release); + std::thread::park_timeout(deadline - now); + } + } + + /// Stops the query and drops its stream on the calling thread, rather than leaving that to + /// 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); + 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))); + } +} + +impl Drop for QueryHandle { + fn drop(&mut self) { + self.cancel(); + } +} + +fn cancelled() -> DataFusionError { + DataFusionError::Execution("Local query cancelled".into()) +} + +#[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_wakes_a_waiting_reader() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .build() + .unwrap(); + let plan = BlockingExec::new(Arc::new(Schema::empty()), 2); + let refs = plan.refs(); + let query = Arc::new( + QueryHandle::start( + LocalQuery::new( + Arc::new(plan), + Arc::new(TaskContext::default()), + usize::MAX, + usize::MAX, + ), + runtime.handle(), + ) + .unwrap(), + ); + let reader = Arc::clone(&query); + let (sender, receiver) = std::sync::mpsc::channel(); + let thread = std::thread::spawn(move || { + sender.send(reader.poll(Duration::from_secs(60))).unwrap(); + }); + // 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"); + let started = Instant::now(); + query.cancel(); + 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 new file mode 100644 index 00000000000..dee31918320 --- /dev/null +++ b/native/local/src/lib.rs @@ -0,0 +1,416 @@ +// 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 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; + +use std::collections::VecDeque; +use std::panic::AssertUnwindSafe; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, AtomicUsize, 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::{ + 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, whatever the read-ahead +/// limit. +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 { + schema: SchemaRef, + partitions: usize, + start: PartitionStart, + context: Arc, + max_running: usize, + running_limit: Option, + read_ahead: usize, +} + +impl LocalQuery { + /// Takes ownership of a fresh graph and its query-scoped context. + /// + /// 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 + /// 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 + /// 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 { + 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 { + schema, + partitions, + start, + 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 { + tokio::runtime::Handle::try_current().map_err(|_| { + DataFusionError::Execution("LocalQuery requires an active Tokio runtime".into()) + })?; + let schema = Arc::clone(&self.schema); + let stream = in_partition_order(&self); + Ok(LocalQueryStream { + stream: Some(stream), + query: Some(self), + schema, + }) + } +} + +fn in_partition_order(query: &LocalQuery) -> SendableRecordBatchStream { + let mut order = PartitionOrder { + partitions: query.partitions, + start: Arc::clone(&query.start), + context: Arc::clone(&query.context), + runtime: tokio::runtime::Handle::current(), + next: 0, + max_running: query.max_running, + running_limit: query.running_limit.clone(), + read_ahead: Arc::new(ReadAhead { + limit: query.read_ahead, + held: AtomicUsize::new(0), + delivered: Notify::new(), + }), + running: VecDeque::new(), + }; + order.start_partitions(); + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&query.schema), + order, + )) +} + +/// 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 { + /// 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. + 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<()>, +} + +/// 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(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(Slot::Starting)); + let (task_queued, task_ended, polled) = + (Arc::clone(&queued), Arc::clone(&ended), Arc::clone(&stream)); + let read_ahead = Arc::clone(&order.read_ahead); + let (start, context) = (Arc::clone(&order.start), Arc::clone(&order.context)); + let task = order.runtime.spawn(async move { + let result = AssertUnwindSafe(async { + let started = start(partition, context); + drop(start); + 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 + // 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 &mut *lock(&polled) { + Slot::Running(stream) => stream.poll_next_unpin(cx), + Slot::Starting | Slot::Closed => Poll::Ready(None), + }) + .await; + match next { + 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); + return; + } + } + } + }) + .catch_unwind() + .await; + if result.is_err() { + let _ = sender.send(Err(DataFusionError::Execution( + "Local query partition panicked".into(), + ))); + } + }); + Self { + batches, + queued, + 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 = std::mem::replace(&mut *lock(&self.stream), Slot::Closed); + let _ = std::panic::catch_unwind(AssertUnwindSafe(|| drop(stream))); + } +} + +/// 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: usize, + start: PartitionStart, + context: Arc, + runtime: tokio::runtime::Handle, + /// 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) { + let partitions = self.partitions; + 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; + } + } + + fn fail(&mut self, error: DataFusionError) -> Poll>> { + self.running.clear(); + self.next = usize::MAX; + Poll::Ready(Some(Err(error))) + } +} + +impl Stream for PartitionOrder { + type Item = Result; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + loop { + self.start_partitions(); + let Some(partition) = self.running.front_mut() else { + return Poll::Ready(None); + }; + match partition.batches.poll_recv(cx) { + 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(); + } + Poll::Ready(None) => { + return self.fail(DataFusionError::Execution( + "Local query partition stopped without completing".into(), + )) + } + 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..ab70ee2c719 --- /dev/null +++ b/native/local/tests/bridge.rs @@ -0,0 +1,351 @@ +// 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_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()), + usize::MAX, + usize::MAX, + ), + 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()), + usize::MAX, + usize::MAX, + ), + 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()), + usize::MAX, + usize::MAX, + ), + 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 query = QueryHandle::start( + LocalQuery::new( + Arc::new(ErrorExec::new()), + Arc::new(TaskContext::default()), + usize::MAX, + usize::MAX, + ), + runtime.handle(), + ) + .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] +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()), + usize::MAX, + usize::MAX, + ), + 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()), usize::MAX, 0), + runtime.handle(), + ) + .unwrap(); + // 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()), + usize::MAX, + usize::MAX, + ), + 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( + // 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(); + 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 <= 2, + "the partition held more than two batches ahead" + ); + } + 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()), + usize::MAX, + usize::MAX, + ), + 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..4cf4660b457 --- /dev/null +++ b/native/local/tests/execution.rs @@ -0,0 +1,615 @@ +// 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; +use std::fmt; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::{Arc, Mutex, Weak}; +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::common::tree_node::TreeNodeRecursion; +use datafusion::common::{DataFusionError, Result}; +use datafusion::execution::config::SessionConfig; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::empty::EmptyExec; +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, PlanProperties, + RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, +}; +use datafusion_comet_local::{LocalQuery, PartitionStart}; +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() +} + +#[derive(Debug, Default)] +struct Observations { + executions: Mutex>, + batches: Mutex>, +} + +/// Records actual partition execution without changing the data or scheduling. +#[derive(Debug)] +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 { + 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), + start_delay: self.start_delay, + })) + } + + 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 { + 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| { + *observations + .batches + .lock() + .unwrap() + .entry(partition) + .or_default() += 1; + batch + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream, + ))) + } +} + +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, + ) +} + +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); +} + +#[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(), usize::MAX, usize::MAX) + .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(), usize::MAX, usize::MAX) + .execute() + .unwrap(); + assert!(timeout(DEADLINE, stream.next()).await.unwrap().is_none()); + assert_eq!(stream.schema(), schema()); + } +} + +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(); + (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(), usize::MAX, usize::MAX) + .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, usize::MAX, usize::MAX) + .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 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(blocking_input, context(), usize::MAX, usize::MAX) + .execute() + .unwrap(); + wait_started(&observations, 4).await; + let other = LocalQuery::new(input(), context(), usize::MAX, usize::MAX) + .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_is_reported_in_its_turn_and_releases_query() { + let plan = Arc::new(ErrorExec::new()); + let refs = Arc::downgrade(&plan); + let mut stream = LocalQuery::new(plan, context(), usize::MAX, usize::MAX) + .execute() + .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())); + let refs = Arc::downgrade(&plan); + 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); +} + +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(), usize::MAX, usize::MAX) + .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::>()); +} + +fn failing() -> Arc { + Arc::new(MockExec::new( + vec![Err(DataFusionError::Execution( + "injected late failure".into(), + ))], + schema(), + )) +} + +#[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(), usize::MAX, usize::MAX) + .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() + .unwrap() + .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(), usize::MAX, usize::MAX) + .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; +} + +#[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(), usize::MAX, usize::MAX) + .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; +} + +#[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(), 2, usize::MAX) + .execute() + .unwrap(); + // The first two partitions start with the query; the rest wait for their turn. + wait_started(&observations, 2).await; + 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(); + wait_started(&observations, 3).await; + let rest: Vec = timeout(DEADLINE, stream.try_collect()) + .await + .unwrap() + .unwrap(); + assert_eq!(rest.len(), 6 * 3 - 4); + 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)); + 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(), 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(), 2, 0).execute().unwrap(); + tokio::time::sleep(Duration::from_millis(200)).await; + 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/pom.xml b/pom.xml index a215e2a818f..7508a792983 100644 --- a/pom.xml +++ b/pom.xml @@ -1011,6 +1011,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 @@ -1056,6 +1057,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 diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 220d6f6b62a..66955c6b20a 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -341,6 +341,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_MAKE_DATE_OPERATOR_FALLBACK_ENABLED: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.makeDate.operatorFallback.enabled") .category(CATEGORY_EXEC) diff --git a/spark/src/main/scala/org/apache/comet/local/CometLocalExec.scala b/spark/src/main/scala/org/apache/comet/local/CometLocalExec.scala new file mode 100644 index 00000000000..40297dee2be --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/local/CometLocalExec.scala @@ -0,0 +1,93 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.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.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, + 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 + * 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. + override def outputPartitioning: Partitioning = UnknownPartitioning(1) + + // 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) + + 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/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala b/spark/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala new file mode 100644 index 00000000000..6c0fa10bf52 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala @@ -0,0 +1,136 @@ +/* + * 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.{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 + * 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, 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. + */ +case class CometLocalResultExec(child: SparkPlan) extends UnaryExecNode { + override def output: Seq[Attribute] = child.output + override def outputPartitioning: Partitioning = ordinary.outputPartitioning + override def outputOrdering: Seq[SortOrder] = ordinary.outputOrdering + + override protected def doExecute(): RDD[InternalRow] = ordinary.execute() + + override def executeCollect(): Array[InternalRow] = + if (sparkContext.isLocal) LocalResultHandoff.collect(child.execute(), maxResultSize) + else super.executeCollect() + + // 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 => + EliminateRedundantTransitions(session)(ColumnarToRowExec(local.block)) + } + .getOrElse(child) + + // 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") + + 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 is running. */ + def active: Int = slots.size() + + /** + * 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). 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") + val id = nextId.incrementAndGet() + val slot = new Slot(maxResultSize) + slots.put(id, slot) + try { + 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, " + + 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]): Unit = { + val slot = slots.get(id) + if (slot == null) return + val buffer = new ArrayBuffer[InternalRow]() + var bytes = 0L + while (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/src/main/scala/org/apache/comet/local/CometLocalRule.scala b/spark/src/main/scala/org/apache/comet/local/CometLocalRule.scala new file mode 100644 index 00000000000..59180453f79 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/local/CometLocalRule.scala @@ -0,0 +1,150 @@ +/* + * 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.{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.{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. */ + def isLocal(plan: SparkPlan): Boolean = plan match { + case _: CometLocalExec | _: CometLocalResultExec => true + 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 not + * admitted: only `collect` runs locally, and a limit collects through `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) + } + 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.map(_.dataType)) + CometLocalResultExec( + CometLocalExec(block.output, spec, resultOrdering(block, partitions), block)) + case converted => converted + } + } + + /** + * 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 && + // 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 + .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/src/main/scala/org/apache/comet/local/LocalQueryIterator.scala b/spark/src/main/scala/org/apache/comet/local/LocalQueryIterator.scala new file mode 100644 index 00000000000..21f27a1d740 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/local/LocalQueryIterator.scala @@ -0,0 +1,149 @@ +/* + * 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 = _ + 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) + + // 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) + // 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(), + spec.partitions, + spec.columns, + taskMemory.manager, + SparkEnv.get.blockManager.getLocalDiskDirs, + CometConf.COMET_BATCH_SIZE.get(), + memory.offHeapMode, + memory.memoryPoolType, + memory.memoryLimit, + memory.offHeapSize, + context.taskAttemptId(), + context.cpus().toLong, + if (maxResultSize > 0) maxResultSize else Long.MaxValue) + } 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() + else if (checkTimestampLabels) util.checkTimestampLabels(pending, spec.outputTypes) + 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/src/main/scala/org/apache/comet/local/NativeLocal.scala b/spark/src/main/scala/org/apache/comet/local/NativeLocal.scala new file mode 100644 index 00000000000..7c7e11df839 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/local/NativeLocal.scala @@ -0,0 +1,61 @@ +/* + * 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. Its running partitions read ahead of delivery up to `readAheadBytes`. + * 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, + offHeapSize: Long, + taskAttemptId: 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. */ + @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/src/main/scala/org/apache/comet/rules/CometRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometRule.scala index fbde4a79090..3ab69d63cd1 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 { @@ -163,8 +164,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) } } diff --git a/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala b/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala new file mode 100644 index 00000000000..910cc79af37 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/local/CometLocalExecutionSuite.scala @@ -0,0 +1,591 @@ +/* + * 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.{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 +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 + +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, 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( + 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 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() + val relations = counts.queryExecution.withCachedData.collect { + case relation: InMemoryRelation => relation + } + 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( + counts.collect().map(r => r.getLong(0) -> r.getLong(1)).toMap == + (0L until 5L).map(_ -> 200L).toMap) + } finally cached.unpersist(blocking = true) + } + } + + 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) + + 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") { + // 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()).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())) + // 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) + // 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) + } + } + + 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(_)) + for (_ <- 0 until 3) { + assert(query.collect().toSeq == rows) + assert(new NativeLocal().activeQueries() == 0) + } + 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("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 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( + 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("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()).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") { + 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]) + 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(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) = + 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) + } + + 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) + localNodes(query).head.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, + memory.offHeapSize, + context.taskAttemptId(), + 1L, + Long.MaxValue) + } + + 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() + } + } +}