sunchao commented on code in PR #6883: URL: https://github.com/apache/datafusion-comet/pull/6883#discussion_r4246764539
########## spark-local/src/main/scala/org/apache/comet/local/CometLocalExec.scala: ########## @@ -0,0 +1,88 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.local + +import org.apache.spark.{Partition, SparkContext, TaskContext} +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder} +import org.apache.spark.sql.catalyst.plans.QueryPlan +import org.apache.spark.sql.catalyst.plans.physical.{Partitioning, UnknownPartitioning} +import org.apache.spark.sql.execution.{ColumnarToRowExec, LeafExecNode, SparkPlan} +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.vectorized.ColumnarBatch + +/** What the result task needs to create the native graph: no plan objects. */ +private[local] case class LocalQuerySpec(plan: Array[Byte], partitions: Int, columns: Int) + +/** + * A whole query run as one native graph by the single Spark task of its result. `block` is the + * Comet native block the graph is created from, kept for explain only. + */ +case class CometLocalExec private[local] ( + override val output: Seq[Attribute], + spec: LocalQuerySpec, + ordering: Seq[SortOrder], + block: SparkPlan) + extends LeafExecNode { + override def supportsColumnar: Boolean = true + + // The result is one partition, but its rows are not clustered by anything. A query over a + // cached local result must still exchange to get its required distribution. + override def outputPartitioning: Partitioning = UnknownPartitioning(1) + + // Partitions are delivered in order, so the block's ordering holds across the result. + override def outputOrdering: Seq[SortOrder] = ordering Review Comment: [P2] Clear partition-local ordering when concatenating multiple partitions. With local and range execution enabled, caching `spark.range(0, 20, 1, 2).sortWithinPartitions(desc("id"))` stores `9..0, 19..10` in one partition, but this method advertises that partition as descending-sorted. `InMemoryTableScanExec` preserves that ordering, and `UnknownPartitioning(1)` satisfies an unpartitioned window’s `AllTuples` requirement. Consequently, `row_number().over(Window.orderBy(desc("id")))` skips the necessary sort and assigns `id=19` rank 11 instead of 1. This reproduces even with Comet disabled for the consuming window. Return no ordering for multi-partition concatenation unless ordering across partition boundaries is proven, or merge the streams by their sort keys. Evidence: Reproduced with exact-head JVM classes and verified native CI artifact on Spark 4.1.3/JDK 17, local[2], AQE disabled. Materialized the descending range cache, disabled Comet, then selected `id` and `row_number().over(Window.orderBy(desc("id")))`. Spark and ordinary Comet returned `(19,1)` and `(9,11)`. Local execution returned `(19,11)` and `(9,1)`. The failing consumer plan was `Window -> InMemoryTableScan`, with no sort. Probe and output: `/tmp/comet6883-review-current-r4qtAM/Review6883Probe.scala` and `probe.log`. ########## native/local/src/lib.rs: ########## @@ -0,0 +1,297 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Query-scoped execution, independent of Spark tasks and JNI. +//! +//! A caller supplies a fresh, fully planned graph for each execution and receives one +//! result stream. The caller owns the Tokio runtime and must keep it running until +//! cleanup finishes. Spark admission and JNI live outside this crate. + +pub mod handle; + +use std::collections::VecDeque; +use std::panic::AssertUnwindSafe; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex, MutexGuard}; +use std::task::{Context, Poll}; + +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::{DataFusionError, Result}; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::{ + execute_stream, ExecutionPlan, ExecutionPlanProperties, RecordBatchStream, + SendableRecordBatchStream, +}; +use futures::future::poll_fn; +use futures::task::AtomicWaker; +use futures::{FutureExt, Stream, StreamExt}; +use tokio::task::JoinHandle; + +/// Batches each root partition may run ahead of the one being delivered, when results are +/// delivered in partition order. +const PARTITION_BUFFER: usize = 2; + +/// One execution's graph and resources. Consumed when execution starts; not a plan cache. +pub struct LocalQuery { + plan: Arc<dyn ExecutionPlan>, + context: Arc<TaskContext>, + in_partition_order: bool, +} + +impl LocalQuery { + /// Takes ownership of a fresh graph and its query-scoped context. + /// + /// The planner must satisfy operator distribution and ordering requirements. It must not + /// reuse stateful operator instances from another query or attempt. + pub fn new(plan: Arc<dyn ExecutionPlan>, context: Arc<TaskContext>) -> Self { + Self { + plan, + context, + in_partition_order: false, + } + } + + /// Delivers the root partitions one after another, in partition order, as Spark returns + /// the partitions of a result. All partitions still run concurrently, each up to a few + /// batches ahead. Only for roots whose partitions do not share an exchange: a partition + /// that waits to be delivered would otherwise leave its exchange buffering for the others. + pub fn in_partition_order(mut self) -> Self { + self.in_partition_order = true; + self + } + + /// Starts all root partitions and exposes one streaming result. + /// + /// By default the partitions are coalesced as they produce batches, which does not + /// preserve their order. A single-partition root preserves its ordering either way. No + /// result-wide collection is performed here. + pub fn execute(self) -> Result<LocalQueryStream> { + tokio::runtime::Handle::try_current().map_err(|_| { + DataFusionError::Execution("LocalQuery requires an active Tokio runtime".into()) + })?; + let schema = self.plan.schema(); + let stream = if self.in_partition_order { + in_partition_order(&self.plan, &self.context)? + } else { + execute_stream(Arc::clone(&self.plan), Arc::clone(&self.context))? + }; + Ok(LocalQueryStream { + stream: Some(stream), + query: Some(self), + schema, + }) + } +} + +fn in_partition_order( + plan: &Arc<dyn ExecutionPlan>, + context: &Arc<TaskContext>, +) -> Result<SendableRecordBatchStream> { + let failure = Arc::new(Failure::default()); + let runtime = tokio::runtime::Handle::current(); + let partitions = (0..plan.output_partitioning().partition_count()) + .map(|partition| { + let stream = plan.execute(partition, Arc::clone(context))?; + Ok(Partition::spawn(&runtime, stream, Arc::clone(&failure))) + }) + .collect::<Result<VecDeque<_>>>()?; + Ok(Box::pin(RecordBatchStreamAdapter::new( + plan.schema(), + PartitionOrder { + partitions, + failure, + }, + ))) +} + +/// One root partition running on its own Tokio task, a few batches ahead of delivery. +struct Partition { + batches: tokio::sync::mpsc::Receiver<RecordBatch>, + /// 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<AtomicBool>, + /// 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<Mutex<Option<SendableRecordBatchStream>>>, + task: JoinHandle<()>, +} + +impl Partition { + fn spawn( + runtime: &tokio::runtime::Handle, + stream: SendableRecordBatchStream, + failure: Arc<Failure>, + ) -> Self { + let (sender, batches) = tokio::sync::mpsc::channel(PARTITION_BUFFER); + let ended = Arc::new(AtomicBool::new(false)); + let stream = Arc::new(Mutex::new(Some(stream))); + let (task_ended, polled) = (Arc::clone(&ended), Arc::clone(&stream)); + let task = runtime.spawn(async move { + let result = AssertUnwindSafe(async { + // Reserve before polling, so at most `PARTITION_BUFFER` batches are held. + while let Ok(permit) = sender.reserve().await { + let next = poll_fn(|cx| match lock(&polled).as_mut() { + Some(stream) => stream.poll_next_unpin(cx), + None => Poll::Ready(None), + }) + .await; + match next { + Some(Ok(batch)) => permit.send(batch), + Some(Err(error)) => return failure.record(error), + None => { + // A taken stream also ends here, but then nobody reads the flag. + task_ended.store(true, Ordering::Release); + return; + } + } + } + }) + .catch_unwind() + .await; + if result.is_err() { + failure.record(DataFusionError::Execution( + "Local query partition panicked".into(), + )); + } + }); + Self { + batches, + ended, + stream, + task, + } + } +} + +impl Drop for Partition { + /// Drops the partition's stream on the calling thread, since cancelling the task only takes + /// effect once a Tokio worker is free. See `QueryHandle::cancel`. + fn drop(&mut self) { + self.task.abort(); + let stream = lock(&self.stream).take(); + let _ = std::panic::catch_unwind(AssertUnwindSafe(|| drop(stream))); + } +} + +/// The first error of any partition, reported as soon as it happens rather than when the +/// failed partition's turn comes. +#[derive(Default)] +struct Failure { + error: Mutex<Option<DataFusionError>>, + waker: AtomicWaker, +} + +impl Failure { + /// Keeps the first error and wakes the consumer. + fn record(&self, error: DataFusionError) { + lock(&self.error).get_or_insert(error); + self.waker.wake(); + } + + fn take(&self) -> Option<DataFusionError> { + lock(&self.error).take() + } +} + +struct PartitionOrder { + partitions: VecDeque<Partition>, + failure: Arc<Failure>, +} + +impl Stream for PartitionOrder { + type Item = Result<RecordBatch>; + + fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> { + self.failure.waker.register(cx.waker()); + loop { + if let Some(error) = self.failure.take() { Review Comment: [P2] Do not surface speculative later-partition failures before a limited query needs those partitions. With ANSI mode, range/local execution enabled and batch size 17, `spark.range(0, 200000, 1, 2).filter("id >= 99999").selectExpr("100L DIV (100000L - id) AS v").limit(1).collect()` returns `[100]` in Spark and ordinary Comet. Local execution starts partition 1 eagerly, encounters division by zero at `id=100000`, and this shared failure check aborts the query before partition 0 delivers its valid row. Spark’s `executeTake` stops after partition 0 and never evaluates the failing partition. Make error delivery aware of limited consumption, or retain ordinary execution for these plans until that behavior can be preserved. Evidence: The exact-head reproduction returned `Right(ArraySeq(100))` with Spark, the same result with ordinary Comet, and `SparkArithmeticException [DIVIDE_BY_ZERO]` with local execution. The local plan was `CollectLimit 1 -> CometLocalResult -> ColumnarToRow -> CometLocal`, containing the projection, filter and two-slice range. Spark’s supported-version `SparkPlan.executeTake` sources request partitions incrementally. Reproduction and output are recorded in `/tmp/comet6883-review-current-r4qtAM/Review6883Probe.scala` and `probe.log`. ########## native/core/src/local.rs: ########## @@ -0,0 +1,249 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! JNI adapter for local execution, which runs a whole admitted Comet plan as one DataFusion +//! graph inside the result task of a local Spark application. The registry holds only live +//! queries, never plans for reuse. EOF, errors and explicit close remove entries; Java task +//! completion also closes. + +use std::collections::HashMap; +use std::sync::atomic::{AtomicI64, Ordering}; +use std::sync::{Arc, LazyLock}; +use std::time::{Duration, Instant}; + +use datafusion::physical_plan::union::UnionExec; +use datafusion::physical_plan::{ExecutionPlan, ExecutionPlanProperties}; +use datafusion_comet_jni_bridge::{jni_new_global_ref, JVMClasses}; +use datafusion_comet_local::handle::{QueryHandle, QueryPoll}; +use datafusion_comet_local::LocalQuery; +use jni::objects::{JByteArray, JLongArray, JObject, JObjectArray, JString, ReleaseMode}; +use jni::sys::{jboolean, jint, jlong, JNI_FALSE}; +use jni::{Env, EnvUnowned}; +use parking_lot::Mutex; + +use crate::errors::{try_unwrap_or_throw, CometError, CometResult}; +use crate::execution::jni_api::{ + create_execution_context, get_runtime, prepare_output, prepare_plan_settings, ExecutionContext, + PlanInputs, PlanJvmRefs, +}; +use crate::execution::memory_pools::create_memory_pool; + +struct Entry { + query: QueryHandle, + columns: usize, + // Holds the plan's memory pool and session until the query is closed. Dropped after the + // query, which cancels the graph that reserves from the pool. Only the creating thread + // uses it; the mutex makes the entry shareable with result pulls. + _context: Mutex<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, + task_attempt_id: jlong, + task_cpus: jlong, +) -> jlong { + try_unwrap_or_throw(&e, |env| { + if !(1..=MAX_PARTITIONS).contains(&partitions) || columns < 1 { + return Err(CometError::Internal( + "Invalid local query parameters".into(), + )); + } + let start = Instant::now(); + let settings = prepare_plan_settings(&env.convert_byte_array(serialized_spark_configs)?)?; + JVMClasses::init(env); + let serialized_plan = env.convert_byte_array(serialized_query)?; + let task_memory_manager = + Arc::new(jni_new_global_ref!(env, comet_task_memory_manager_obj)?); + let memory_pool_type = memory_pool_type.try_to_string(env)?; + let mut local_dirs_vec = vec![]; + for i in 0..local_dirs.len(env)? { + let dir = local_dirs.get_element(env, i)?; + let dir = unsafe { JString::from_raw(&*env, dir.into_raw()) }; + local_dirs_vec.push(dir.try_to_string(env)?); + } + let id = next_id()?; + let inputs = PlanInputs { + id, + settings, + serialized_plan, + partition_count: partitions as usize, + batch_size: batch_size as usize, + off_heap_mode: off_heap_mode != JNI_FALSE, + memory_pool_type, + memory_limit, + task_attempt_id, + task_cpus: task_cpus as usize, + local_dirs: local_dirs_vec, + metrics_update_interval: 0, + start, + jvm: PlanJvmRefs::default(), + }; + let context = create_execution_context(inputs, |config| { + create_memory_pool(config, task_memory_manager, task_attempt_id) + })?; + let plan = local_plan(&context, partitions)?; + if plan.schema().fields().len() != columns as usize { + return Err(CometError::Internal( + "Local output schema width mismatch".into(), + )); + } + let query = LocalQuery::new(plan, context.session_ctx.task_ctx()).in_partition_order(); + let query = QueryHandle::start(query, &get_runtime())?; + register( + id, + Entry { + query, + columns: columns as usize, + _context: Mutex::new(context), + }, + )?; + Ok(id) + }) +} + +const MAX_PARTITIONS: jint = 1024; + +/// One graph for the whole query. Each partition's operators are created by the planner for +/// that partition, as a Spark task would create them, so per-partition state such as a range +/// slice or a partition-dependent expression is correct. Partition `i` of the graph is the +/// single output partition of the operators created for partition `i`. +fn local_plan(context: &ExecutionContext, partitions: jint) -> CometResult<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> { + NEXT_ID + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| id.checked_add(1)) Review Comment: [P2] Avoid the deprecated atomic operation in the required Rust build. The stable toolchain used by CI deprecates `AtomicI64::fetch_update`, so the new call fails `cargo clippy --color=never --all-targets --workspace -- -D warnings`. This prevents the required Rust job from reaching its tests even when local execution is disabled. Use an equivalent checked compare/exchange loop or a narrowly scoped compatibility allowance that respects the declared Rust 1.95 minimum. Evidence: Exact-head job https://github.com/apache/datafusion-comet/actions/runs/38217614661/job/114708727490 fails at `core/src/local.rs:175:10`: `use of deprecated method ... fetch_update: renamed to try_update for consistency`. The log identifies `-D deprecated` as implied by `-D warnings`, exits 101, and skips subsequent Rust test steps. Aggregate `Required Checks` is also failed. -- 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]
