viirya commented on code in PR #6883: URL: https://github.com/apache/datafusion-comet/pull/6883#discussion_r4255353170
########## 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: Fixed in 49dbdf859. You're right that Spark's `outputOrdering` only holds within each partition, so concatenating the partitions can't keep it in general. `CometLocalExec` now reports the block's ordering only when it holds across the concatenated partitions: with one partition, or when the block's partitioning is a `RangePartitioning` whose ordering is a prefix of the block's ordering, as a range's slices are. `sortWithinPartitions` over a multi-slice range now reports no ordering. Your reproduction is now a test: the window over the cached result returns Spark's ranks. ########## 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: Agreed, reporting any partition's failure at once broke `executeTake`'s semantics. Each partition's error now travels through its own channel and is reported when that partition's turn comes, so a limit satisfied by earlier partitions never sees a later partition's failure, while a full collect still fails once it reaches it. That also removed the shared failure slot. Your ANSI `DIV` example is now a test, together with the full collect that still fails with `DIVIDE_BY_ZERO`, and the native tests check that a later failure is reported only in its turn. ########## 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: Replaced with `fetch_add` and a check that the counter has not wrapped, which works on the 1.95 minimum. `cargo clippy --all-targets --workspace -- -D warnings` is clean with Rust 1.99.0, the toolchain CI uses. -- 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]
