sunchao commented on code in PR #6883: URL: https://github.com/apache/datafusion-comet/pull/6883#discussion_r4419020951
########## spark/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala: ########## @@ -0,0 +1,128 @@ +/* + * 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} + +/** + * Row result boundary of a local query. Spark's collect encodes and compresses every row of each + * result partition in its task, then decodes them on the driver. A local collect instead runs the + * whole query as one native graph in a single Spark task (keeping cancellation, job groups and + * SQL metrics), which hands copied rows to the driver in this JVM, since admission requires an + * in-process local master. + * + * Only `collect` runs that way. Everything else reads the result as an RDD: a take requests + * partitions in growing batches, an iterator fetches one partition at a time, a tail scans from + * the last one, and a cache, `df.rdd` or PySpark's collect keep its partitions. So the node + * executes as the ordinary native block otherwise, one Spark task per partition, and reports the + * block's partitioning and ordering. + */ +case class CometLocalResultExec(child: SparkPlan) extends UnaryExecNode { + override def output: Seq[Attribute] = child.output + override def outputPartitioning: Partitioning = 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() + + private def ordinary: SparkPlan = + child + .collectFirst { case local: CometLocalExec => ColumnarToRowExec(local.block) } + .getOrElse(child) + + // Spark enforces this on serialized task results; this boundary enforces it while copying. + private def maxResultSize: Long = + sparkContext.getConf.getSizeAsBytes("spark.driver.maxResultSize", "1g") + + override protected def withNewChildInternal(newChild: SparkPlan): CometLocalResultExec = + copy(child = newChild) +} + +/** Driver-owned result slots, keyed by numeric IDs that the task closure carries. */ +private[local] object LocalResultHandoff { + private final class Slot(val maxResultSize: Long) { + @volatile var rows: Array[InternalRow] = _ + @volatile var exceededBytes: Long = -1L + } + + private val slots = new ConcurrentHashMap[Long, Slot]() + private val nextId = new AtomicLong() + + /** Number of live slots; zero whenever no collect 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). Unlike Spark's compressed size, this counts uncompressed UnsafeRow bytes, so it is + * the more conservative of the two. + */ + 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 Review Comment: **Handle separate driver and executor classloaders.** `slots` is a Scala singleton. A task in the same JVM can still load a second copy when the plugin JAR is shipped via `spark.jars` and `spark.executor.userClassPathFirst=true`. I reproduced this on Spark 4.1.3 with the handoff from this commit and one row. Ordinary collection succeeded under both classloading modes. The handoff succeeded with parent-first loading, but under child-first loading this lookup missed, the task returned successfully, and the driver threw `Local result not delivered`. Could the handoff use a registry shared across those loaders, or fall back to ordinary collection before scheduling in this configuration? ########## native/core/src/local.rs: ########## @@ -0,0 +1,259 @@ +// 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); Review Comment: **Use Spark’s concurrent task count here.** With `local[32]` and `spark.task.cpus=16`, Spark schedules two tasks at once (`LocalSchedulerBackend.maxNumConcurrentTasks` divides total cores by CPUs per task). This line starts 32 partition streams. In a native probe using this scheduler, DataFusion 55.2.0 sorts and a 256 MiB `CometFairMemoryPool`, 32 concurrent partitions failed a 10 MiB reservation. Changing only the limit to two returned all 557,056 rows. Could we divide by the `task_cpus` already passed to `createQuery`, clamp to at least one, and cover a non-default `spark.task.cpus`? ########## 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: **Either path can reach the result-size limit first.** Spark counts serialized task results, which include per-task overhead. With an 8 KiB cap on Spark 4.1.3, ordinary collection of 128 one-long rows split across 128 partitions failed after 26 task results. Coalescing the same Dataset’s RDD into this commit’s single-partition handoff returned all 128 (2,048 bytes of `UnsafeRow` payload). Please describe the two accounting rules here and in `LocalResultHandoff.collect`’s Scaladoc without promising the local limit is always stricter. ########## native/core/src/local.rs: ########## @@ -0,0 +1,259 @@ +// 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()) + .in_partition_order(max_running, read_ahead_bytes.max(0) as usize); + let query = QueryHandle::start(query, &get_runtime())?; + register( + id, + Entry { + query, + columns: columns as usize, + _context: Mutex::new(context), + }, + )?; + Ok(id) + }) +} + +const MAX_PARTITIONS: jint = 1024; + +/// One graph for the whole query. Each partition's operators are created by the planner for +/// that partition, as a Spark task would create them, so per-partition state such as a range +/// slice or a partition-dependent expression is correct. Partition `i` of the graph is the +/// single output partition of the operators created for partition `i`. +fn local_plan(context: &ExecutionContext, partitions: jint) -> CometResult<Arc<dyn ExecutionPlan>> { + let mut plans = Vec::with_capacity(partitions as usize); + for partition in 0..partitions { + let (scans, shuffle_scans, root) = context.planner(partition).create_plan( + &context.spark_plan, + &mut vec![], + partitions as usize, + )?; + if !scans.is_empty() + || !shuffle_scans.is_empty() + || root.native_plan.output_partitioning().partition_count() != 1 + { + return Err(CometError::Internal( + "Local query plans must have no JVM input and one output partition".into(), + )); + } + plans.push(Arc::clone(&root.native_plan)); + } + Ok(UnionExec::try_new(plans)?) +} + +fn next_id() -> CometResult<i64> { + // 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<jlong> { + let entry = QUERIES + .lock() + .get(&id) + .cloned() + .ok_or_else(|| CometError::Internal("Local query is closed or unknown".into()))?; + if arrays.len(env)? != entry.columns || schemas.len(env)? != entry.columns { + return Err(CometError::Internal( + "Local output column count mismatch".into(), + )); + } + // The caller supplies writable Arrow C structs. Reject null pointers before exporting. + for addresses in [&arrays, &schemas] { + let elements = unsafe { addresses.get_elements(env, ReleaseMode::NoCopyBack)? }; + if elements.contains(&0) { + return Err(CometError::Internal("Null local output pointer".into())); + } + } + match entry.query.poll(Duration::from_millis(50))? { + QueryPoll::Pending => Ok(-2), + QueryPoll::Finished => Ok(-1), + QueryPoll::Batch(batch) => prepare_output(env, arrays, schemas, batch, false), Review Comment: **Keep the output validation switches effective for local collects.** Ordinary Comet passes `exec_context.debug_native` here so `spark.comet.debug.enabled` runs Arrow’s `validate_full` before export. The hardcoded `false` skips it. On import, `LocalQueryIterator` also omits `NativeUtil.checkTimestampLabels`, so `spark.comet.testing.strictTimestampLabels.enabled` misses local output, including nested timestamps. Could the local path forward the debug setting and retain the expected Spark output types needed for the JVM label check? ########## 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: **Build the JNI library in the test recipe.** `cargo test -p datafusion-comet-local` builds the Rust scheduling crate and its tests. The `NativeLocal` JNI entry points live in `native/core`, so this sequence leaves a clean checkout without the `libcomet` required by the JVM suite. Adding `make core` from the repository root before the Maven command would make these steps reproducible. ########## native/core/src/local.rs: ########## @@ -0,0 +1,259 @@ +// 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(); Review Comment: **Wait for outstanding plan reservations when closing.** Cancelling the query drops its streams, but a sort can own a worker that releases its reservation only when it next yields. Ordinary `release_plan` drops the context and then waits on `plan_memory` before Spark can reclaim the task’s memory. This path skips that wait. In a controlled probe using this `QueryHandle` and `PlanMemoryPool`, a worker still held 4,096 bytes after cancellation and handle drop, then released them when it resumed. Could both teardown paths share the bounded reservation wait, including early exit on cancellation or the result-size limit? ########## spark/src/main/scala/org/apache/comet/local/CometLocalResultExec.scala: ########## @@ -0,0 +1,128 @@ +/* + * 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} + +/** + * Row result boundary of a local query. Spark's collect encodes and compresses every row of each + * result partition in its task, then decodes them on the driver. A local collect instead runs the + * whole query as one native graph in a single Spark task (keeping cancellation, job groups and + * SQL metrics), which hands copied rows to the driver in this JVM, since admission requires an + * in-process local master. + * + * Only `collect` runs that way. Everything else reads the result as an RDD: a take requests + * partitions in growing batches, an iterator fetches one partition at a time, a tail scans from + * the last one, and a cache, `df.rdd` or PySpark's collect keep its partitions. So the node + * executes as the ordinary native block otherwise, one Spark task per partition, and reports the + * block's partitioning and ordering. + */ +case class CometLocalResultExec(child: SparkPlan) extends UnaryExecNode { + override def output: Seq[Attribute] = child.output + override def outputPartitioning: Partitioning = 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() + + private def ordinary: SparkPlan = + child + .collectFirst { case local: CometLocalExec => ColumnarToRowExec(local.block) } Review Comment: **Preserve ordinary plan preparation for RDD actions.** Constructing `ColumnarToRowExec` here happens after `EliminateRedundantTransitions` has run. For a bare range, that rule normally restores Spark’s `RangeExec` and avoids a native round trip. On Spark 4.1.3 with this commit’s Scala and the same ordinary Comet JNI library in both runs, `spark.range(0, 16000000, 1, 8).queryExecution.toRdd.count()` returned the same count in both modes, but created zero native iterators with the local flag off and eight with it on. This also bypasses the rule’s columnar-to-row selection for projections. Could the fallback execute the normally prepared block and extend the RDD test to check the bare-range behavior covered by `CometRangeExecSuite`? -- 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]
