viirya commented on code in PR #6883: URL: https://github.com/apache/datafusion-comet/pull/6883#discussion_r4432795956
########## docs/source/contributor-guide/local-execution.md: ########## @@ -0,0 +1,144 @@ +<!-- +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. +--> + +# 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; +- the plan is a batch query and not a subquery being prepared; +- the converted plan is one native block, a tree of Comet native operators whose root carries + the block's serialized plan, with no JVM input, broadcast or shuffle inside it; +- the block's only leaf is a `CometRangeExec`, so the range must be converted to Comet + (`spark.comet.exec.range.enabled`); +- its expressions are deterministic, contain no subquery, and do not call back into the JVM, + which the serialized plan is also checked for; +- it has between 1 and 1024 partitions. + +A root `CollectLimitExec`, which Spark plans for `take`, `head` and `show`, is kept above the +local query, and reads it as described under result delivery. Such plans have no exchange, so AQE does not wrap them and they are admitted +whether or not AQE is enabled. + +## Execution + +The admitted plan becomes `CometLocalResultExec` over `CometLocalExec`, a leaf whose RDD has a +single partition. That partition's task creates the native graph: + +- The memory pool, spill directories and DataFusion session are those `Native.createPlan` creates + for any Comet plan of the task: `create_execution_context` builds them, with Comet's function + registrations and session settings, and the pool acquires memory from Spark through the task's + `CometTaskMemoryManager`. +- The block's operators are created once per partition, by the planner for that partition, as a + Spark task would create them. Partition `i` of the graph is the output of the operators created + for partition `i`, so per-partition state such as a range slice is the same as with ordinary + Comet execution. +- `native/local` runs the partitions on Comet's Tokio runtime and delivers them in partition + order. At most `spark.executor.cores` partitions (N for `local[N]`) run at once, and the next + starts when the first is delivered, as Spark runs at most that many tasks. A partition's + operators register their memory consumers when it starts, so this also bounds how many share + the task's memory pool. The result therefore has Spark's row order. A partition's error is + reported when its turn comes. +- 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 native side hands batches to the JVM through a queue that holds at most one batch. Exported +Arrow arrays use `prepare_output` and `NativeUtil`; an imported batch stays valid after its +native query is closed, until the JVM closes it. A pull returns "pending" after 50 ms, so that +the task can check whether it was killed. + +`CometLocalResultExec` sits above Spark's columnar-to-row transition. For `collect` it still +runs the single task as a Spark job, which keeps cancellation, job groups and SQL metrics, but +the task copies its rows into a driver-side slot instead of returning an encoded task result. Spark would otherwise encode and compress every row in that one task and decode +them on the driver, all on one thread. While copying, the task counts uncompressed `UnsafeRow` +bytes and stops reading, which closes the native query, once they exceed +`spark.driver.maxResultSize`; the driver then fails the action. Spark counts compressed bytes, Review Comment: Agreed. The result-size section and `LocalResultHandoff.collect` now describe both accounting rules without saying which is stricter (f99686ccd). ########## docs/source/contributor-guide/local-execution.md: ########## @@ -0,0 +1,144 @@ +<!-- +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. +--> + +# 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; +- the plan is a batch query and not a subquery being prepared; +- the converted plan is one native block, a tree of Comet native operators whose root carries + the block's serialized plan, with no JVM input, broadcast or shuffle inside it; +- the block's only leaf is a `CometRangeExec`, so the range must be converted to Comet + (`spark.comet.exec.range.enabled`); +- its expressions are deterministic, contain no subquery, and do not call back into the JVM, + which the serialized plan is also checked for; +- it has between 1 and 1024 partitions. + +A root `CollectLimitExec`, which Spark plans for `take`, `head` and `show`, is kept above the +local query, and reads it as described under result delivery. Such plans have no exchange, so AQE does not wrap them and they are admitted +whether or not AQE is enabled. + +## Execution + +The admitted plan becomes `CometLocalResultExec` over `CometLocalExec`, a leaf whose RDD has a +single partition. That partition's task creates the native graph: + +- The memory pool, spill directories and DataFusion session are those `Native.createPlan` creates + for any Comet plan of the task: `create_execution_context` builds them, with Comet's function + registrations and session settings, and the pool acquires memory from Spark through the task's + `CometTaskMemoryManager`. +- The block's operators are created once per partition, by the planner for that partition, as a + Spark task would create them. Partition `i` of the graph is the output of the operators created + for partition `i`, so per-partition state such as a range slice is the same as with ordinary + Comet execution. +- `native/local` runs the partitions on Comet's Tokio runtime and delivers them in partition + order. At most `spark.executor.cores` partitions (N for `local[N]`) run at once, and the next + starts when the first is delivered, as Spark runs at most that many tasks. A partition's + operators register their memory consumers when it starts, so this also bounds how many share + the task's memory pool. The result therefore has Spark's row order. A partition's error is + reported when its turn comes. +- 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 native side hands batches to the JVM through a queue that holds at most one batch. Exported +Arrow arrays use `prepare_output` and `NativeUtil`; an imported batch stays valid after its +native query is closed, until the JVM closes it. A pull returns "pending" after 50 ms, so that +the task can check whether it was killed. + +`CometLocalResultExec` sits above Spark's columnar-to-row transition. For `collect` it still +runs the single task as a Spark job, which keeps cancellation, job groups and SQL metrics, but +the task copies its rows into a driver-side slot instead of returning an encoded task result. Spark would otherwise encode and compress every row in that one task and decode +them on the driver, all on one thread. While copying, the task counts uncompressed `UnsafeRow` +bytes and stops reading, which closes the native query, once they exceed +`spark.driver.maxResultSize`; the driver then fails the action. Spark counts compressed bytes, +so this check is the stricter one. On a non-local master the node uses Spark's ordinary +collect. + +Every other action reads the result as an RDD of the block's partitions: a take requests +partitions in growing batches until it has enough rows, an iterator fetches one partition at a +time so that it holds at most one, a tail scans partitions from the last one, and a cache, +`df.rdd`, `foreachPartition`, `localCheckpoint` or PySpark's collect and `toLocalIterator` keep +the partitions for later reads. So `CometLocalResultExec` executes as the ordinary native block, +one Spark task per partition, for everything but `collect`. The partitions these actions +evaluate, and so the rows they return and the errors they raise, are Spark's, and a cached result +is scanned as many partitions as without local execution. + +## Lifecycle and cancellation + +The JNI registry holds numeric, never reused IDs for at most 1024 live queries. EOF, errors, +early termination and task completion remove an entry. A concurrent pull holds an `Arc`, so a +close cannot free an entry that is being read, and cancellation does not take the reader lock. +A query records when it has delivered its last batch; if the runtime stops it before that, the +reader reports an error instead of a truncated result. + +Closing a query drops its streams on the calling thread, as `releasePlan` does for an ordinary +plan without JVM input: cancelling a Tokio task only takes effect once a worker is free, and +every worker can be waiting, for instance in Spark's `acquireMemory` for memory those streams +hold. Tasks that DataFusion operators spawn themselves still end asynchronously. + +## Testing + +Build native code before the JVM suite, and never use Maven `-pl`: + +```shell +cd native +cargo test -p datafusion-comet-local Review Comment: Added `make core` before the Maven run in f99686ccd. ########## native/core/src/local.rs: ########## @@ -0,0 +1,263 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! JNI adapter for local execution, which runs a whole admitted Comet plan as one DataFusion +//! graph inside the result task of a local Spark application. The registry holds only live +//! queries, never plans for reuse. EOF, errors and explicit close remove entries; Java task +//! completion also closes. + +use std::collections::HashMap; +use std::sync::atomic::{AtomicI64, Ordering}; +use std::sync::{Arc, LazyLock}; +use std::time::{Duration, Instant}; + +use datafusion::physical_plan::union::UnionExec; +use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; +use datafusion_comet_jni_bridge::{jni_new_global_ref, JVMClasses}; +use datafusion_comet_local::handle::{QueryHandle, QueryPoll}; +use datafusion_comet_local::LocalQuery; +use jni::objects::{JByteArray, JLongArray, JObject, JObjectArray, JString, ReleaseMode}; +use jni::sys::{jboolean, jint, jlong, JNI_FALSE}; +use jni::{Env, EnvUnowned}; +use parking_lot::Mutex; + +use crate::errors::{try_unwrap_or_throw, CometError, CometResult}; +use crate::execution::jni_api::{ + create_execution_context, create_memory_pool, get_runtime, prepare_output, + prepare_plan_settings, ExecutionContext, PlanInputs, PlanJvmRefs, +}; +use crate::execution::spark_config::{SparkConfig, SPARK_EXECUTOR_CORES}; + +struct Entry { + query: QueryHandle, + columns: usize, + // Holds the plan's memory pool and session until the query is closed. Dropped after the + // 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<Box<ExecutionContext>>, +} + +// 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<Mutex<HashMap<i64, Arc<Entry>>>> = 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)?)?; + // Run as many partitions at once as Spark would run tasks, which also bounds how many + // operators share the memory pool. `local[N]` resolves to N here. + let max_running = settings.spark_config.get_usize(SPARK_EXECUTOR_CORES, 1); + JVMClasses::init(env); + let serialized_plan = env.convert_byte_array(serialized_query)?; + let task_memory_manager = + 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 plan = local_plan(&context, partitions)?; + if plan.schema().fields().len() != columns as usize { + return Err(CometError::Internal( + "Local output schema width mismatch".into(), + )); + } + let query = LocalQuery::new( + plan, + context.session_ctx.task_ctx(), + max_running, Review Comment: Fixed in f99686ccd. Underneath this is how Spark divides memory: a local collect is one Spark task, and Comet's pools give each task running a native plan an equal share, so two local collects each running 16 sorts give each sort half a task's share. Each time a partition would start, the query now divides its slots by the executor's tasks running native plans (the count the pools already use), this task included. With two concurrent local collects on `local[16]`, each runs up to 8 partitions; the same holds when ordinary Comet tasks are running. Partitions already running keep running, so a job that starts after a local query has started all its partitions still leaves them less than a task's share until they end. The guide lists that as a limitation. With your setup (`local[16]`, 256 MiB off-heap, batch size 17, 96,000 rows in 16 partitions sorted within partitions by `sha2(cast(id as string), 256)`), two concurrent local collects both failed to get memory without the divided limit, on two runs. With it, each returned all 96,000 rows on two runs, as with local execution disabled. ########## native/local/src/handle.rs: ########## @@ -0,0 +1,203 @@ +// 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<Option<LocalQueryStream>>, + /// The thread of the reader in a poll, which `cancel` wakes. + reader: Mutex<Option<Thread>>, + cancelled: AtomicBool, + runtime: Handle, +} + +/// Wakes a reader parked in `QueryHandle::poll`. +struct Unpark(Thread); + +impl Wake for Unpark { + fn wake(self: Arc<Self>) { + self.0.unpark(); + } +} + +impl QueryHandle { + pub fn start(query: LocalQuery, runtime: &Handle) -> Result<Self> { + 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(), + }) + } + + /// 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<QueryPoll> { + 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<LocalQueryStream>, + wait: Duration, + ) -> Result<QueryPoll> { + let deadline = Instant::now() + wait; + let waker = Waker::from(Arc::new(Unpark(std::thread::current()))); + let mut cx = Context::from_waker(&waker); + // Starting a partition executes its operators, which may spawn tasks. + let _guard = self.runtime.enter(); + loop { + // Checked after the reader is registered, so a cancel either sees the reader and + // wakes it, or is seen here. + if self.cancelled.load(Ordering::Acquire) { + return Err(cancelled()); + } + let Some(query) = stream.as_mut() else { + return Err(cancelled()); + }; + // A panic must not unwind into the JNI caller. + let next = + std::panic::catch_unwind(AssertUnwindSafe(|| query.poll_next_unpin(&mut cx))) Review Comment: Agreed, fixed in f99686ccd. A partition now runs `plan.execute` on its own task, so the reader only spawns tasks and receives from channels. A failure to start is reported in that partition's turn, like any other partition error. A new native test starts partitions behind a 500 ms blocking `execute` and checks that the first poll returns pending right away. ########## native/local/src/handle.rs: ########## @@ -0,0 +1,203 @@ +// 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<Option<LocalQueryStream>>, + /// The thread of the reader in a poll, which `cancel` wakes. + reader: Mutex<Option<Thread>>, + cancelled: AtomicBool, + runtime: Handle, +} + +/// Wakes a reader parked in `QueryHandle::poll`. +struct Unpark(Thread); + +impl Wake for Unpark { + fn wake(self: Arc<Self>) { + self.0.unpark(); + } +} + +impl QueryHandle { + pub fn start(query: LocalQuery, runtime: &Handle) -> Result<Self> { + 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(), + }) + } + + /// 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<QueryPoll> { + 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<LocalQueryStream>, + wait: Duration, + ) -> Result<QueryPoll> { + let deadline = Instant::now() + wait; + let waker = Waker::from(Arc::new(Unpark(std::thread::current()))); + let mut cx = Context::from_waker(&waker); + // Starting a partition executes its operators, which may spawn tasks. + let _guard = self.runtime.enter(); + loop { + // Checked after the reader is registered, so a cancel either sees the reader and + // wakes it, or is seen here. + if self.cancelled.load(Ordering::Acquire) { + return Err(cancelled()); + } + let Some(query) = stream.as_mut() else { + return Err(cancelled()); + }; + // A panic must not unwind into the JNI caller. + let next = + std::panic::catch_unwind(AssertUnwindSafe(|| query.poll_next_unpin(&mut cx))) + .map_err(|_| DataFusionError::Execution("Local query panicked".into()))?; + match next { + Poll::Ready(Some(batch)) => return batch.map(QueryPoll::Batch), + Poll::Ready(None) => return Ok(QueryPoll::Finished), + Poll::Pending => (), + } + let now = Instant::now(); + if now >= deadline { + return Ok(QueryPoll::Pending); + } + std::thread::park_timeout(deadline - now); + } + } + + /// Stops the query and drops its stream on the calling thread, rather than leaving that to + /// 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(); + }); + while lock(&query.reader).is_none() { Review Comment: Fixed in f99686ccd. The reader sets a test-only flag right before `park_timeout`, and the test cancels only after it is set. With `reader.unpark()` removed, the test now fails after the 60 s park. ########## docs/source/contributor-guide/local-execution.md: ########## @@ -0,0 +1,146 @@ +<!-- +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. +--> + +# 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; +- 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. +- `native/local` runs the partitions on Comet's Tokio runtime and delivers them in partition + order. At most `spark.executor.cores` partitions (N for `local[N]`) run at once, and the next + starts when the first is delivered, as Spark runs at most that many tasks. A partition's + operators register their memory consumers when it starts, so this also bounds how many share + the task's memory pool. The result therefore has Spark's row order. A partition's error is + reported when its turn comes. +- 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 counts compressed bytes, +so this check is the stricter one. On a non-local master the node uses Spark's ordinary +collect. + +Every other action reads the result as an RDD of the block's partitions: an iterator fetches +one partition at a time so that it holds at most one, a tail scans partitions from the last one, and a cache, +`df.rdd`, `foreachPartition`, `localCheckpoint` or PySpark's collect and `toLocalIterator` keep Review Comment: Thanks, you're right, and my earlier reply to andygrove had that wrong. The guide now describes classic `collect()` as going through `executeCollect` and the handoff, and `toLocalIterator()` as partitioned (f99686ccd). -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
