sunchao commented on code in PR #6883: URL: https://github.com/apache/datafusion-comet/pull/6883#discussion_r4433047932
########## native/core/src/local.rs: ########## @@ -0,0 +1,307 @@ +// 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 log::warn; +use parking_lot::Mutex; + +use crate::errors::{try_unwrap_or_throw, CometError, CometResult}; +use crate::execution::jni_api::{ + create_execution_context, create_memory_pool, get_runtime, prepare_output, + prepare_plan_settings, release_plan, running_native_tasks, ExecutionContext, PlanInputs, + PlanJvmRefs, +}; +use crate::execution::spark_config::{SparkConfig, SPARK_EXECUTOR_CORES}; + +struct Entry { + query: QueryHandle, + columns: usize, + /// Validates exported batches, as `spark.comet.debug.enabled` does for any Comet plan. + validate: bool, + // Holds the plan's memory pool and session until the query is closed. Released after the + // query, which cancels the graph that reserves from the pool. Only the creating thread + // uses it; the mutex makes the entry shareable with result pulls. + context: Mutex<Option<Box<ExecutionContext>>>, +} + +impl Drop for Entry { + /// Releases the context as `Native.releasePlan` does, waiting briefly for the reservations + /// that the cancelled graph's tasks still hold, so that they are returned before the Spark + /// task can end. + fn drop(&mut self) { + self.query.cancel(); + if let Some(context) = self.context.get_mut().take() { + if let Err(error) = release_plan(context, &mut |_| Ok(())) { + warn!("Releasing a local query failed: {error}"); + } + } + } +} + +// JNI carries numeric IDs instead of raw pointers, so close racing with a result pull +// cannot free an object the pull is using. Last reader releases the removed entry. +static QUERIES: LazyLock<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)?)?; + // `local[N]` resolves to N cores here. + let slots = task_slots( + settings.spark_config.get_usize(SPARK_EXECUTOR_CORES, 1), + task_cpus as usize, + ); + JVMClasses::init(env); + let serialized_plan = env.convert_byte_array(serialized_query)?; + let task_memory_manager = + Arc::new(jni_new_global_ref!(env, comet_task_memory_manager_obj)?); + let memory_pool_type = memory_pool_type.try_to_string(env)?; + let mut local_dirs_vec = vec![]; + for i in 0..local_dirs.len(env)? { + let dir = local_dirs.get_element(env, i)?; + let dir = unsafe { JString::from_raw(&*env, dir.into_raw()) }; + local_dirs_vec.push(dir.try_to_string(env)?); + } + let id = next_id()?; + let inputs = PlanInputs { + id, + settings, + serialized_plan, + partition_count: partitions as usize, + batch_size: batch_size as usize, + off_heap_mode: off_heap_mode != JNI_FALSE, + memory_pool_type, + memory_limit, + off_heap_size, + task_attempt_id, + task_cpus: task_cpus as usize, + local_dirs: local_dirs_vec, + metrics_update_interval: 0, + start, + jvm: PlanJvmRefs::default(), + }; + let context = create_execution_context(inputs, |config| { + create_memory_pool(config, task_memory_manager, task_attempt_id) + })?; + let plan = local_plan(&context, partitions)?; + if plan.schema().fields().len() != columns as usize { + return Err(CometError::Internal( + "Local output schema width mismatch".into(), + )); + } + // Run at most as many partitions at once as Spark runs tasks. Comet's pools split Spark's + // memory equally among the tasks running native plans, this one included, so with T of + // them the partitions each get what a task would only while at most slots / T run. + let query = LocalQuery::new( + plan, + context.session_ctx.task_ctx(), + slots, + read_ahead_bytes.max(0) as usize, + ) + .with_running_limit(Arc::new(move || slots / running_native_tasks().max(1))); + let query = QueryHandle::start(query, &get_runtime())?; + register( + id, + Entry { + query, + columns: columns as usize, + validate: context.debug_native, + context: Mutex::new(Some(context)), + }, + )?; + Ok(id) + }) +} + +const MAX_PARTITIONS: jint = 1024; + +/// The tasks Spark runs at once with `cores`, as `LocalSchedulerBackend` counts them. +fn task_slots(cores: usize, task_cpus: usize) -> usize { + (cores / task_cpus.max(1)).max(1) +} + +/// One graph for the whole query. 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( Review Comment: This loop retains a physical plan per partition before the first stream starts. Planning a string literal clones its payload into each plan. For an admitted 1,024-partition range with a column-dependent filter containing a 1 MiB string literal, that is roughly 1 GiB of literal payloads even on `local[2]`. That figure follows from the copies in the code. Could plans be created as execution slots open, or could the partition-independent literal storage be shared? -- 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]
