| // 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. |
| |
| //! Define JNI APIs which can be called from Java/Scala. |
| |
| use super::{serde, utils::SparkArrowConvert}; |
| use crate::{ |
| errors::{try_unwrap_or_throw, CometError, CometResult}, |
| execution::{ |
| metrics::utils::update_comet_metric, planner::PhysicalPlanner, serde::to_arrow_datatype, |
| shuffle::spark_unsafe::row::process_sorted_row_partition, sort::RdxSort, |
| }, |
| jvm_bridge::{jni_new_global_ref, JVMClasses}, |
| }; |
| use arrow::array::{Array, RecordBatch, UInt32Array}; |
| use arrow::compute::{take, TakeOptions}; |
| use arrow::datatypes::DataType as ArrowDataType; |
| use datafusion::common::{Result as DataFusionResult, ScalarValue}; |
| use datafusion::execution::disk_manager::DiskManagerMode; |
| use datafusion::execution::memory_pool::MemoryPool; |
| use datafusion::execution::runtime_env::RuntimeEnvBuilder; |
| use datafusion::logical_expr::ScalarUDF; |
| use datafusion::{ |
| execution::disk_manager::DiskManagerBuilder, |
| physical_plan::{display::DisplayableExecutionPlan, SendableRecordBatchStream}, |
| prelude::{SessionConfig, SessionContext}, |
| }; |
| use datafusion_comet_proto::spark_operator::Operator; |
| use datafusion_spark::function::bitwise::bit_get::SparkBitGet; |
| use datafusion_spark::function::bitwise::bitwise_not::SparkBitwiseNot; |
| use datafusion_spark::function::datetime::date_add::SparkDateAdd; |
| use datafusion_spark::function::datetime::date_sub::SparkDateSub; |
| use datafusion_spark::function::datetime::last_day::SparkLastDay; |
| use datafusion_spark::function::datetime::next_day::SparkNextDay; |
| use datafusion_spark::function::hash::crc32::SparkCrc32; |
| use datafusion_spark::function::hash::sha1::SparkSha1; |
| use datafusion_spark::function::hash::sha2::SparkSha2; |
| use datafusion_spark::function::map::map_from_entries::MapFromEntries; |
| use datafusion_spark::function::math::expm1::SparkExpm1; |
| use datafusion_spark::function::math::hex::SparkHex; |
| use datafusion_spark::function::math::width_bucket::SparkWidthBucket; |
| use datafusion_spark::function::string::char::CharFunc; |
| use datafusion_spark::function::string::concat::SparkConcat; |
| use futures::poll; |
| use futures::stream::StreamExt; |
| use jni::objects::JByteBuffer; |
| use jni::sys::{jlongArray, JNI_FALSE}; |
| use jni::{ |
| errors::Result as JNIResult, |
| objects::{ |
| GlobalRef, JByteArray, JClass, JIntArray, JLongArray, JObject, JObjectArray, JString, |
| ReleaseMode, |
| }, |
| sys::{jboolean, jdouble, jint, jlong}, |
| JNIEnv, |
| }; |
| use std::collections::HashMap; |
| use std::path::PathBuf; |
| use std::time::{Duration, Instant}; |
| use std::{sync::Arc, task::Poll}; |
| use tokio::runtime::Runtime; |
| use tokio::sync::mpsc; |
| |
| use crate::execution::memory_pools::{ |
| create_memory_pool, handle_task_shared_pool_release, parse_memory_pool_config, MemoryPoolConfig, |
| }; |
| use crate::execution::operators::ScanExec; |
| use crate::execution::shuffle::{read_ipc_compressed, CompressionCodec}; |
| use crate::execution::spark_plan::SparkPlan; |
| |
| use crate::execution::tracing::{log_memory_usage, trace_begin, trace_end, with_trace}; |
| |
| use crate::execution::memory_pools::logging_pool::LoggingMemoryPool; |
| use crate::execution::spark_config::{ |
| SparkConfig, COMET_DEBUG_ENABLED, COMET_DEBUG_MEMORY, COMET_EXPLAIN_NATIVE_ENABLED, |
| COMET_MAX_TEMP_DIRECTORY_SIZE, COMET_TRACING_ENABLED, SPARK_EXECUTOR_CORES, |
| }; |
| use crate::parquet::encryption_support::{CometEncryptionFactory, ENCRYPTION_FACTORY_ID}; |
| use datafusion_comet_proto::spark_operator::operator::OpStruct; |
| use log::info; |
| use std::sync::OnceLock; |
| #[cfg(feature = "jemalloc")] |
| use tikv_jemalloc_ctl::{epoch, stats}; |
| |
| static TOKIO_RUNTIME: OnceLock<Runtime> = OnceLock::new(); |
| |
| fn parse_usize_env_var(name: &str) -> Option<usize> { |
| std::env::var_os(name).and_then(|n| n.to_str().and_then(|s| s.parse::<usize>().ok())) |
| } |
| |
| fn build_runtime(default_worker_threads: Option<usize>) -> Runtime { |
| let mut builder = tokio::runtime::Builder::new_multi_thread(); |
| if let Some(n) = parse_usize_env_var("COMET_WORKER_THREADS") { |
| info!("Comet tokio runtime: using COMET_WORKER_THREADS={n}"); |
| builder.worker_threads(n); |
| } else if let Some(n) = default_worker_threads { |
| info!("Comet tokio runtime: using spark.executor.cores={n} worker threads"); |
| builder.worker_threads(n); |
| } else { |
| info!("Comet tokio runtime: using default thread count"); |
| } |
| if let Some(n) = parse_usize_env_var("COMET_MAX_BLOCKING_THREADS") { |
| builder.max_blocking_threads(n); |
| } |
| builder |
| .enable_all() |
| .build() |
| .expect("Failed to create Tokio runtime") |
| } |
| |
| /// Initialize the global Tokio runtime with the given default worker thread count. |
| /// If the runtime is already initialized, this is a no-op. |
| pub fn init_runtime(default_worker_threads: usize) { |
| TOKIO_RUNTIME.get_or_init(|| build_runtime(Some(default_worker_threads))); |
| } |
| |
| /// Function to get a handle to the global Tokio runtime |
| pub fn get_runtime() -> &'static Runtime { |
| TOKIO_RUNTIME.get_or_init(|| build_runtime(None)) |
| } |
| |
| /// Comet native execution context. Kept alive across JNI calls. |
| struct ExecutionContext { |
| /// The id of the execution context. |
| pub id: i64, |
| /// Task attempt id |
| pub task_attempt_id: i64, |
| /// The deserialized Spark plan |
| pub spark_plan: Operator, |
| /// The number of partitions |
| pub partition_count: usize, |
| /// The DataFusion root operator converted from the `spark_plan` |
| pub root_op: Option<Arc<SparkPlan>>, |
| /// The input sources for the DataFusion plan |
| pub scans: Vec<ScanExec>, |
| /// The global reference of input sources for the DataFusion plan |
| pub input_sources: Vec<Arc<GlobalRef>>, |
| /// The record batch stream to pull results from |
| pub stream: Option<SendableRecordBatchStream>, |
| /// Receives batches from a spawned tokio task (async I/O path) |
| pub batch_receiver: Option<mpsc::Receiver<DataFusionResult<RecordBatch>>>, |
| /// Native metrics |
| pub metrics: Arc<GlobalRef>, |
| // The interval in milliseconds to update metrics |
| pub metrics_update_interval: Option<Duration>, |
| // The last update time of metrics |
| pub metrics_last_update_time: Instant, |
| /// Counter to avoid checking time on every poll iteration (reduces syscalls) |
| pub poll_count_since_metrics_check: u32, |
| /// The time it took to create the native plan and configure the context |
| pub plan_creation_time: Duration, |
| /// DataFusion SessionContext |
| pub session_ctx: Arc<SessionContext>, |
| /// Whether to enable additional debugging checks & messages |
| pub debug_native: bool, |
| /// Whether to write native plans with metrics to stdout |
| pub explain_native: bool, |
| /// Memory pool config |
| pub memory_pool_config: MemoryPoolConfig, |
| /// Whether to log memory usage on each call to execute_plan |
| pub tracing_enabled: bool, |
| } |
| |
| /// Accept serialized query plan and return the address of the native query plan. |
| /// # Safety |
| /// This function is inherently unsafe since it deals with raw pointers passed from JNI. |
| #[no_mangle] |
| pub unsafe extern "system" fn Java_org_apache_comet_Native_createPlan( |
| e: JNIEnv, |
| _class: JClass, |
| id: jlong, |
| iterators: JObjectArray, |
| serialized_query: JByteArray, |
| serialized_spark_configs: JByteArray, |
| partition_count: jint, |
| metrics_node: JObject, |
| metrics_update_interval: jlong, |
| comet_task_memory_manager_obj: JObject, |
| local_dirs: JObjectArray, |
| batch_size: jint, |
| off_heap_mode: jboolean, |
| memory_pool_type: JString, |
| memory_limit: jlong, |
| memory_limit_per_task: jlong, |
| task_attempt_id: jlong, |
| task_cpus: jlong, |
| key_unwrapper_obj: JObject, |
| ) -> jlong { |
| try_unwrap_or_throw(&e, |mut env| { |
| // Deserialize Spark configs |
| let bytes = env.convert_byte_array(serialized_spark_configs)?; |
| let spark_configs = serde::deserialize_config(bytes.as_slice())?; |
| let spark_config: HashMap<String, String> = spark_configs.entries.into_iter().collect(); |
| |
| // Initialize the tokio runtime with spark.executor.cores as the default |
| // worker thread count, falling back to 1 if not set. |
| let executor_cores = spark_config.get_usize(SPARK_EXECUTOR_CORES, 1); |
| init_runtime(executor_cores); |
| |
| // Access Comet configs |
| let debug_native = spark_config.get_bool(COMET_DEBUG_ENABLED); |
| let explain_native = spark_config.get_bool(COMET_EXPLAIN_NATIVE_ENABLED); |
| let tracing_enabled = spark_config.get_bool(COMET_TRACING_ENABLED); |
| let max_temp_directory_size = |
| spark_config.get_u64(COMET_MAX_TEMP_DIRECTORY_SIZE, 100 * 1024 * 1024 * 1024); |
| let logging_memory_pool = spark_config.get_bool(COMET_DEBUG_MEMORY); |
| |
| with_trace("createPlan", tracing_enabled, || { |
| // Init JVM classes |
| JVMClasses::init(&mut env); |
| |
| let start = Instant::now(); |
| |
| // Deserialize query plan |
| let bytes = env.convert_byte_array(serialized_query)?; |
| let spark_plan = serde::deserialize_op(bytes.as_slice())?; |
| |
| let metrics = Arc::new(jni_new_global_ref!(env, metrics_node)?); |
| |
| // Get the global references of input sources |
| let mut input_sources = vec![]; |
| let num_inputs = env.get_array_length(&iterators)?; |
| for i in 0..num_inputs { |
| let input_source = env.get_object_array_element(&iterators, i)?; |
| let input_source = Arc::new(jni_new_global_ref!(env, input_source)?); |
| input_sources.push(input_source); |
| } |
| |
| // Create DataFusion memory pool |
| let task_memory_manager = |
| Arc::new(jni_new_global_ref!(env, comet_task_memory_manager_obj)?); |
| |
| let memory_pool_type = env.get_string(&memory_pool_type)?.into(); |
| let memory_pool_config = parse_memory_pool_config( |
| off_heap_mode != JNI_FALSE, |
| memory_pool_type, |
| memory_limit, |
| memory_limit_per_task, |
| )?; |
| let memory_pool = |
| create_memory_pool(&memory_pool_config, task_memory_manager, task_attempt_id); |
| |
| let memory_pool = if logging_memory_pool { |
| Arc::new(LoggingMemoryPool::new(task_attempt_id as u64, memory_pool)) |
| } else { |
| memory_pool |
| }; |
| |
| // Get local directories for storing spill files |
| let num_local_dirs = env.get_array_length(&local_dirs)?; |
| let mut local_dirs_vec = vec![]; |
| for i in 0..num_local_dirs { |
| let local_dir: JString = env.get_object_array_element(&local_dirs, i)?.into(); |
| let local_dir = env.get_string(&local_dir)?; |
| local_dirs_vec.push(local_dir.into()); |
| } |
| |
| // We need to keep the session context alive. Some session state like temporary |
| // dictionaries are stored in session context. If it is dropped, the temporary |
| // dictionaries will be dropped as well. |
| let session = prepare_datafusion_session_context( |
| batch_size as usize, |
| memory_pool, |
| local_dirs_vec, |
| max_temp_directory_size, |
| task_cpus as usize, |
| &spark_config, |
| )?; |
| |
| let plan_creation_time = start.elapsed(); |
| |
| let metrics_update_interval = if metrics_update_interval > 0 { |
| Some(Duration::from_millis(metrics_update_interval as u64)) |
| } else { |
| None |
| }; |
| |
| // Handle key unwrapper for encrypted files |
| if !key_unwrapper_obj.is_null() { |
| let encryption_factory = CometEncryptionFactory { |
| key_unwrapper: jni_new_global_ref!(env, key_unwrapper_obj)?, |
| }; |
| session.runtime_env().register_parquet_encryption_factory( |
| ENCRYPTION_FACTORY_ID, |
| Arc::new(encryption_factory), |
| ); |
| } |
| |
| let exec_context = Box::new(ExecutionContext { |
| id, |
| task_attempt_id, |
| spark_plan, |
| partition_count: partition_count as usize, |
| root_op: None, |
| scans: vec![], |
| input_sources, |
| stream: None, |
| batch_receiver: None, |
| metrics, |
| metrics_update_interval, |
| metrics_last_update_time: Instant::now(), |
| poll_count_since_metrics_check: 0, |
| plan_creation_time, |
| session_ctx: Arc::new(session), |
| debug_native, |
| explain_native, |
| memory_pool_config, |
| tracing_enabled, |
| }); |
| |
| Ok(Box::into_raw(exec_context) as i64) |
| }) |
| }) |
| } |
| |
| /// Configure DataFusion session context. |
| fn prepare_datafusion_session_context( |
| batch_size: usize, |
| memory_pool: Arc<dyn MemoryPool>, |
| local_dirs: Vec<String>, |
| max_temp_directory_size: u64, |
| task_cpus: usize, |
| spark_config: &HashMap<String, String>, |
| ) -> CometResult<SessionContext> { |
| let paths = local_dirs.into_iter().map(PathBuf::from).collect(); |
| let disk_manager = DiskManagerBuilder::default() |
| .with_mode(DiskManagerMode::Directories(paths)) |
| .with_max_temp_directory_size(max_temp_directory_size); |
| let mut rt_config = RuntimeEnvBuilder::new().with_disk_manager_builder(disk_manager); |
| rt_config = rt_config.with_memory_pool(memory_pool); |
| |
| let mut session_config = SessionConfig::new() |
| .with_target_partitions(task_cpus) |
| // This DataFusion context is within the scope of an executing Spark Task. We want to set |
| // its internal parallelism to the number of CPUs allocated to Spark Tasks. This can be |
| // modified by changing spark.task.cpus in the Spark config. |
| .with_batch_size(batch_size) |
| // DataFusion partial aggregates can emit duplicate rows so we disable the |
| // skip partial aggregation feature because this is not compatible with Spark's |
| // use of partial aggregates. |
| .set( |
| "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", |
| // this is the threshold of number of groups / number of rows and the |
| // maximum value is 1.0, so we set the threshold a little higher just |
| // to be safe |
| &ScalarValue::Float64(Some(1.1)), |
| ); |
| |
| // Pass through DataFusion configs from Spark. |
| // e.g: spark-shell --conf spark.comet.datafusion.sql_parser.parse_float_as_decimal=true |
| // becomes datafusion.sql_parser.parse_float_as_decimal=true |
| const SPARK_COMET_DF_PREFIX: &str = "spark.comet.datafusion."; |
| for (key, value) in spark_config { |
| if let Some(df_key) = key.strip_prefix(SPARK_COMET_DF_PREFIX) { |
| let df_key = format!("datafusion.{df_key}"); |
| session_config = session_config.set_str(&df_key, value); |
| } |
| } |
| |
| let runtime = rt_config.build()?; |
| |
| let mut session_ctx = SessionContext::new_with_config_rt(session_config, Arc::new(runtime)); |
| |
| datafusion::functions_nested::register_all(&mut session_ctx)?; |
| register_datafusion_spark_function(&session_ctx); |
| // Must be the last one to override existing functions with the same name |
| datafusion_comet_spark_expr::register_all_comet_functions(&mut session_ctx)?; |
| |
| Ok(session_ctx) |
| } |
| |
| // register UDFs from datafusion-spark crate |
| fn register_datafusion_spark_function(session_ctx: &SessionContext) { |
| session_ctx.register_udf(ScalarUDF::new_from_impl(SparkExpm1::default())); |
| session_ctx.register_udf(ScalarUDF::new_from_impl(SparkSha2::default())); |
| session_ctx.register_udf(ScalarUDF::new_from_impl(CharFunc::default())); |
| session_ctx.register_udf(ScalarUDF::new_from_impl(SparkBitGet::default())); |
| session_ctx.register_udf(ScalarUDF::new_from_impl(SparkDateAdd::default())); |
| session_ctx.register_udf(ScalarUDF::new_from_impl(SparkDateSub::default())); |
| session_ctx.register_udf(ScalarUDF::new_from_impl(SparkLastDay::default())); |
| session_ctx.register_udf(ScalarUDF::new_from_impl(SparkNextDay::default())); |
| session_ctx.register_udf(ScalarUDF::new_from_impl(SparkSha1::default())); |
| session_ctx.register_udf(ScalarUDF::new_from_impl(SparkConcat::default())); |
| session_ctx.register_udf(ScalarUDF::new_from_impl(SparkBitwiseNot::default())); |
| session_ctx.register_udf(ScalarUDF::new_from_impl(SparkHex::default())); |
| session_ctx.register_udf(ScalarUDF::new_from_impl(SparkWidthBucket::default())); |
| session_ctx.register_udf(ScalarUDF::new_from_impl(MapFromEntries::default())); |
| session_ctx.register_udf(ScalarUDF::new_from_impl(SparkCrc32::default())); |
| session_ctx.register_udf(ScalarUDF::new_from_impl(SparkSpace::default())); |
| } |
| |
| /// Prepares arrow arrays for output. |
| fn prepare_output( |
| env: &mut JNIEnv, |
| array_addrs: JLongArray, |
| schema_addrs: JLongArray, |
| output_batch: RecordBatch, |
| validate: bool, |
| ) -> CometResult<jlong> { |
| let num_cols = env.get_array_length(&array_addrs)? as usize; |
| |
| let array_addrs = unsafe { env.get_array_elements(&array_addrs, ReleaseMode::NoCopyBack)? }; |
| let array_addrs = &*array_addrs; |
| |
| let schema_addrs = unsafe { env.get_array_elements(&schema_addrs, ReleaseMode::NoCopyBack)? }; |
| let schema_addrs = &*schema_addrs; |
| |
| let results = output_batch.columns(); |
| let num_rows = output_batch.num_rows(); |
| |
| // there are edge cases where num_cols can be zero due to Spark optimizations |
| // when the results of a query are not used |
| if num_cols > 0 { |
| if results.len() != num_cols { |
| return Err(CometError::Internal(format!( |
| "Output column count mismatch: expected {num_cols}, got {}", |
| results.len() |
| ))); |
| } |
| |
| if validate { |
| // Validate the output arrays. |
| for array in results.iter() { |
| let array_data = array.to_data(); |
| array_data |
| .validate_full() |
| .expect("Invalid output array data"); |
| } |
| } |
| |
| let mut i = 0; |
| while i < results.len() { |
| let array_ref = results.get(i).ok_or(CometError::IndexOutOfBounds(i))?; |
| |
| if array_ref.offset() != 0 { |
| // https://github.com/apache/datafusion-comet/issues/2051 |
| // Bug with non-zero offset FFI, so take to a new array which will have an offset of 0. |
| // We expect this to be a cold code path, hence the check_bounds: true and assert_eq. |
| let indices = UInt32Array::from((0..num_rows as u32).collect::<Vec<u32>>()); |
| let new_array = take( |
| array_ref, |
| &indices, |
| Some(TakeOptions { check_bounds: true }), |
| )?; |
| |
| assert_eq!(new_array.offset(), 0); |
| |
| new_array |
| .to_data() |
| .move_to_spark(array_addrs[i], schema_addrs[i])?; |
| } else { |
| array_ref |
| .to_data() |
| .move_to_spark(array_addrs[i], schema_addrs[i])?; |
| } |
| i += 1; |
| } |
| } |
| |
| Ok(num_rows as jlong) |
| } |
| |
| /// Pull the next input from JVM. Note that we cannot pull input batches in |
| /// `ScanStream.poll_next` when the execution stream is polled for output. |
| /// Because the input source could be another native execution stream, which |
| /// will be executed in another tokio blocking thread. It causes JNI throw |
| /// Java exception. So we pull input batches here and insert them into scan |
| /// operators before polling the stream, |
| #[inline] |
| fn pull_input_batches(exec_context: &mut ExecutionContext) -> Result<(), CometError> { |
| exec_context.scans.iter_mut().try_for_each(|scan| { |
| scan.get_next_batch()?; |
| Ok::<(), CometError>(()) |
| }) |
| } |
| |
| /// Accept serialized query plan and the addresses of Arrow Arrays from Spark, |
| /// then execute the query. Return addresses of arrow vector. |
| /// # Safety |
| /// This function is inherently unsafe since it deals with raw pointers passed from JNI. |
| #[no_mangle] |
| pub unsafe extern "system" fn Java_org_apache_comet_Native_executePlan( |
| e: JNIEnv, |
| _class: JClass, |
| stage_id: jint, |
| partition: jint, |
| exec_context: jlong, |
| array_addrs: JLongArray, |
| schema_addrs: JLongArray, |
| ) -> jlong { |
| try_unwrap_or_throw(&e, |mut env| { |
| // Retrieve the query |
| let exec_context = get_execution_context(exec_context); |
| |
| let tracing_event_name = match &exec_context.spark_plan.op_struct { |
| Some(OpStruct::ShuffleWriter(_)) => "executePlan(ShuffleWriter)", |
| _ => "executePlan", |
| }; |
| |
| if exec_context.tracing_enabled { |
| #[cfg(feature = "jemalloc")] |
| { |
| let e = epoch::mib().unwrap(); |
| let allocated = stats::allocated::mib().unwrap(); |
| e.advance().unwrap(); |
| use crate::execution::tracing::log_memory_usage; |
| log_memory_usage("jemalloc_allocated", allocated.read().unwrap() as u64); |
| } |
| } |
| |
| with_trace(tracing_event_name, exec_context.tracing_enabled, || { |
| let exec_context_id = exec_context.id; |
| |
| // Initialize the execution stream. |
| // Because we don't know if input arrays are dictionary-encoded when we create |
| // query plan, we need to defer stream initialization to first time execution. |
| if exec_context.root_op.is_none() { |
| let start = Instant::now(); |
| let planner = |
| PhysicalPlanner::new(Arc::clone(&exec_context.session_ctx), partition) |
| .with_exec_id(exec_context_id); |
| let (scans, root_op) = planner.create_plan( |
| &exec_context.spark_plan, |
| &mut exec_context.input_sources.clone(), |
| exec_context.partition_count, |
| )?; |
| let physical_plan_time = start.elapsed(); |
| |
| exec_context.plan_creation_time += physical_plan_time; |
| exec_context.root_op = Some(Arc::clone(&root_op)); |
| exec_context.scans = scans; |
| |
| if exec_context.explain_native { |
| let formatted_plan_str = |
| DisplayableExecutionPlan::new(root_op.native_plan.as_ref()).indent(true); |
| info!("Comet native query plan:\n{formatted_plan_str:}"); |
| } |
| |
| let task_ctx = exec_context.session_ctx.task_ctx(); |
| // Each Comet native execution corresponds to a single Spark partition, |
| // so we should always execute partition 0. |
| let stream = root_op.native_plan.execute(0, task_ctx)?; |
| |
| if exec_context.scans.is_empty() { |
| // No JVM data sources — spawn onto tokio so the executor |
| // thread parks in blocking_recv instead of busy-polling. |
| // |
| // Channel capacity of 2 allows the producer to work one batch |
| // ahead while the consumer processes the current one via JNI, |
| // without buffering excessive memory. Increasing this would |
| // trade memory for latency hiding if JNI/FFI overhead dominates; |
| // decreasing to 1 would serialize production and consumption. |
| let (tx, rx) = mpsc::channel(2); |
| let mut stream = stream; |
| get_runtime().spawn(async move { |
| while let Some(batch) = stream.next().await { |
| if tx.send(batch).await.is_err() { |
| break; |
| } |
| } |
| }); |
| exec_context.batch_receiver = Some(rx); |
| } else { |
| exec_context.stream = Some(stream); |
| } |
| } else { |
| // Pull input batches |
| pull_input_batches(exec_context)?; |
| } |
| |
| if let Some(rx) = &mut exec_context.batch_receiver { |
| match rx.blocking_recv() { |
| Some(Ok(batch)) => { |
| update_metrics(&mut env, exec_context)?; |
| return prepare_output( |
| &mut env, |
| array_addrs, |
| schema_addrs, |
| batch, |
| exec_context.debug_native, |
| ); |
| } |
| Some(Err(e)) => { |
| return Err(e.into()); |
| } |
| None => { |
| log_plan_metrics(exec_context, stage_id, partition); |
| return Ok(-1); |
| } |
| } |
| } |
| |
| // ScanExec path: busy-poll to interleave JVM batch pulls with stream polling |
| get_runtime().block_on(async { |
| loop { |
| let next_item = exec_context.stream.as_mut().unwrap().next(); |
| let poll_output = poll!(next_item); |
| |
| // Only check time every 100 polls to reduce syscall overhead |
| if let Some(interval) = exec_context.metrics_update_interval { |
| exec_context.poll_count_since_metrics_check += 1; |
| if exec_context.poll_count_since_metrics_check >= 100 { |
| let now = Instant::now(); |
| if now - exec_context.metrics_last_update_time >= interval { |
| update_metrics(&mut env, exec_context)?; |
| exec_context.metrics_last_update_time = now; |
| } |
| exec_context.poll_count_since_metrics_check = 0; |
| } |
| } |
| |
| match poll_output { |
| Poll::Ready(Some(output)) => { |
| return prepare_output( |
| &mut env, |
| array_addrs, |
| schema_addrs, |
| output?, |
| exec_context.debug_native, |
| ); |
| } |
| Poll::Ready(None) => { |
| log_plan_metrics(exec_context, stage_id, partition); |
| return Ok(-1); |
| } |
| Poll::Pending => { |
| // JNI call to pull batches from JVM into ScanExec operators. |
| // block_in_place lets tokio move other tasks off this worker |
| // while we wait for JVM data. |
| tokio::task::block_in_place(|| pull_input_batches(exec_context))?; |
| } |
| } |
| } |
| }) |
| }) |
| }) |
| } |
| |
| #[no_mangle] |
| /// Drop the native query plan object and context object. |
| pub extern "system" fn Java_org_apache_comet_Native_releasePlan( |
| e: JNIEnv, |
| _class: JClass, |
| exec_context: jlong, |
| ) { |
| try_unwrap_or_throw(&e, |mut env| unsafe { |
| let execution_context = get_execution_context(exec_context); |
| |
| // Update metrics |
| update_metrics(&mut env, execution_context)?; |
| |
| handle_task_shared_pool_release( |
| execution_context.memory_pool_config.pool_type, |
| execution_context.task_attempt_id, |
| ); |
| |
| let _: Box<ExecutionContext> = Box::from_raw(execution_context); |
| Ok(()) |
| }) |
| } |
| |
| /// Updates the metrics of the query plan. |
| fn update_metrics(env: &mut JNIEnv, exec_context: &mut ExecutionContext) -> CometResult<()> { |
| if let Some(native_query) = &exec_context.root_op { |
| let metrics = exec_context.metrics.as_obj(); |
| update_comet_metric(env, metrics, native_query) |
| } else { |
| Ok(()) |
| } |
| } |
| |
| fn log_plan_metrics(exec_context: &ExecutionContext, stage_id: jint, partition: jint) { |
| if exec_context.explain_native { |
| if let Some(plan) = &exec_context.root_op { |
| let formatted_plan_str = |
| DisplayableExecutionPlan::with_metrics(plan.native_plan.as_ref()).indent(true); |
| info!( |
| "Comet native query plan with metrics (Plan #{} Stage {} Partition {}):\ |
| \n plan creation took {:?}:\ |
| \n{formatted_plan_str:}", |
| plan.plan_id, stage_id, partition, exec_context.plan_creation_time |
| ); |
| } |
| } |
| } |
| |
| fn convert_datatype_arrays( |
| env: &'_ mut JNIEnv<'_>, |
| serialized_datatypes: JObjectArray, |
| ) -> JNIResult<Vec<ArrowDataType>> { |
| let array_len = env.get_array_length(&serialized_datatypes)?; |
| let mut res: Vec<ArrowDataType> = Vec::new(); |
| |
| for i in 0..array_len { |
| let inner_array = env.get_object_array_element(&serialized_datatypes, i)?; |
| let inner_array: JByteArray = inner_array.into(); |
| let bytes = env.convert_byte_array(inner_array)?; |
| let data_type = serde::deserialize_data_type(bytes.as_slice()).unwrap(); |
| let arrow_dt = to_arrow_datatype(&data_type); |
| res.push(arrow_dt); |
| } |
| |
| Ok(res) |
| } |
| |
| fn get_execution_context<'a>(id: i64) -> &'a mut ExecutionContext { |
| unsafe { |
| (id as *mut ExecutionContext) |
| .as_mut() |
| .expect("Comet execution context shouldn't be null!") |
| } |
| } |
| |
| /// Used by Comet shuffle external sorter to write sorted records to disk. |
| /// # Safety |
| /// This function is inherently unsafe since it deals with raw pointers passed from JNI. |
| #[no_mangle] |
| pub unsafe extern "system" fn Java_org_apache_comet_Native_writeSortedFileNative( |
| e: JNIEnv, |
| _class: JClass, |
| row_addresses: JLongArray, |
| row_sizes: JIntArray, |
| serialized_datatypes: JObjectArray, |
| file_path: JString, |
| prefer_dictionary_ratio: jdouble, |
| batch_size: jlong, |
| checksum_enabled: jboolean, |
| checksum_algo: jint, |
| current_checksum: jlong, |
| compression_codec: JString, |
| compression_level: jint, |
| tracing_enabled: jboolean, |
| ) -> jlongArray { |
| try_unwrap_or_throw(&e, |mut env| unsafe { |
| with_trace( |
| "writeSortedFileNative", |
| tracing_enabled != JNI_FALSE, |
| || { |
| let data_types = convert_datatype_arrays(&mut env, serialized_datatypes)?; |
| |
| let row_num = env.get_array_length(&row_addresses)? as usize; |
| let row_addresses = |
| env.get_array_elements(&row_addresses, ReleaseMode::NoCopyBack)?; |
| |
| let row_sizes = env.get_array_elements(&row_sizes, ReleaseMode::NoCopyBack)?; |
| |
| let row_addresses_ptr = row_addresses.as_ptr(); |
| let row_sizes_ptr = row_sizes.as_ptr(); |
| |
| let output_path: String = env.get_string(&file_path).unwrap().into(); |
| |
| let checksum_enabled = checksum_enabled == 1; |
| let current_checksum = if current_checksum == i64::MIN { |
| // Initial checksum is not available. |
| None |
| } else { |
| Some(current_checksum as u32) |
| }; |
| |
| let compression_codec: String = env.get_string(&compression_codec).unwrap().into(); |
| |
| let compression_codec = match compression_codec.as_str() { |
| "zstd" => CompressionCodec::Zstd(compression_level), |
| "lz4" => CompressionCodec::Lz4Frame, |
| "snappy" => CompressionCodec::Snappy, |
| _ => CompressionCodec::Lz4Frame, |
| }; |
| |
| let (written_bytes, checksum) = process_sorted_row_partition( |
| row_num, |
| batch_size as usize, |
| row_addresses_ptr, |
| row_sizes_ptr, |
| &data_types, |
| output_path, |
| prefer_dictionary_ratio, |
| checksum_enabled, |
| checksum_algo, |
| current_checksum, |
| &compression_codec, |
| )?; |
| |
| let checksum = if let Some(checksum) = checksum { |
| checksum as i64 |
| } else { |
| // Spark checksums (CRC32 or Adler32) are both u32, so we use i64::MIN to indicate |
| // checksum is not available. |
| i64::MIN |
| }; |
| |
| let long_array = env.new_long_array(2)?; |
| env.set_long_array_region(&long_array, 0, &[written_bytes, checksum])?; |
| |
| Ok(long_array.into_raw()) |
| }, |
| ) |
| }) |
| } |
| |
| #[no_mangle] |
| /// Used by Comet shuffle external sorter to sort in-memory row partition ids. |
| pub extern "system" fn Java_org_apache_comet_Native_sortRowPartitionsNative( |
| e: JNIEnv, |
| _class: JClass, |
| address: jlong, |
| size: jlong, |
| tracing_enabled: jboolean, |
| ) { |
| try_unwrap_or_throw(&e, |_| { |
| with_trace( |
| "sortRowPartitionsNative", |
| tracing_enabled != JNI_FALSE, |
| || { |
| // SAFETY: JVM unsafe memory allocation is aligned with long. |
| let array = |
| unsafe { std::slice::from_raw_parts_mut(address as *mut i64, size as usize) }; |
| array.rdxsort(); |
| Ok(()) |
| }, |
| ) |
| }) |
| } |
| |
| #[no_mangle] |
| /// Used by Comet native shuffle reader |
| /// # Safety |
| /// This function is inherently unsafe since it deals with raw pointers passed from JNI. |
| pub unsafe extern "system" fn Java_org_apache_comet_Native_decodeShuffleBlock( |
| e: JNIEnv, |
| _class: JClass, |
| byte_buffer: JByteBuffer, |
| length: jint, |
| array_addrs: JLongArray, |
| schema_addrs: JLongArray, |
| tracing_enabled: jboolean, |
| ) -> jlong { |
| try_unwrap_or_throw(&e, |mut env| { |
| with_trace("decodeShuffleBlock", tracing_enabled != JNI_FALSE, || { |
| let raw_pointer = env.get_direct_buffer_address(&byte_buffer)?; |
| let length = length as usize; |
| let slice: &[u8] = unsafe { std::slice::from_raw_parts(raw_pointer, length) }; |
| let batch = read_ipc_compressed(slice)?; |
| prepare_output(&mut env, array_addrs, schema_addrs, batch, false) |
| }) |
| }) |
| } |
| |
| #[no_mangle] |
| /// # Safety |
| /// This function is inherently unsafe since it deals with raw pointers passed from JNI. |
| pub unsafe extern "system" fn Java_org_apache_comet_Native_traceBegin( |
| e: JNIEnv, |
| _class: JClass, |
| event: JString, |
| ) { |
| try_unwrap_or_throw(&e, |mut env| { |
| let name: String = env.get_string(&event).unwrap().into(); |
| trace_begin(&name); |
| Ok(()) |
| }) |
| } |
| |
| #[no_mangle] |
| /// # Safety |
| /// This function is inherently unsafe since it deals with raw pointers passed from JNI. |
| pub unsafe extern "system" fn Java_org_apache_comet_Native_traceEnd( |
| e: JNIEnv, |
| _class: JClass, |
| event: JString, |
| ) { |
| try_unwrap_or_throw(&e, |mut env| { |
| let name: String = env.get_string(&event).unwrap().into(); |
| trace_end(&name); |
| Ok(()) |
| }) |
| } |
| |
| #[no_mangle] |
| /// # Safety |
| /// This function is inherently unsafe since it deals with raw pointers passed from JNI. |
| pub unsafe extern "system" fn Java_org_apache_comet_Native_logMemoryUsage( |
| e: JNIEnv, |
| _class: JClass, |
| name: JString, |
| value: jlong, |
| ) { |
| try_unwrap_or_throw(&e, |mut env| { |
| let name: String = env.get_string(&name).unwrap().into(); |
| log_memory_usage(&name, value as u64); |
| Ok(()) |
| }) |
| } |
| |
| // ============================================================================ |
| // Native Columnar to Row Conversion |
| // ============================================================================ |
| |
| use crate::execution::columnar_to_row::ColumnarToRowContext; |
| use arrow::ffi::{from_ffi, FFI_ArrowArray, FFI_ArrowSchema}; |
| use datafusion_spark::function::string::space::SparkSpace; |
| |
| /// Initialize a native columnar to row converter. |
| /// |
| /// # Safety |
| /// This function is inherently unsafe since it deals with raw pointers passed from JNI. |
| #[no_mangle] |
| pub unsafe extern "system" fn Java_org_apache_comet_Native_columnarToRowInit( |
| e: JNIEnv, |
| _class: JClass, |
| serialized_schema: JObjectArray, |
| batch_size: jint, |
| ) -> jlong { |
| try_unwrap_or_throw(&e, |mut env| { |
| // Deserialize the schema |
| let schema = convert_datatype_arrays(&mut env, serialized_schema)?; |
| |
| // Create the context |
| let ctx = Box::new(ColumnarToRowContext::new(schema, batch_size as usize)); |
| |
| Ok(Box::into_raw(ctx) as jlong) |
| }) |
| } |
| |
| /// Convert Arrow columnar data to Spark UnsafeRow format. |
| /// |
| /// # Safety |
| /// This function is inherently unsafe since it deals with raw pointers passed from JNI. |
| #[no_mangle] |
| pub unsafe extern "system" fn Java_org_apache_comet_Native_columnarToRowConvert( |
| e: JNIEnv, |
| _class: JClass, |
| c2r_handle: jlong, |
| array_addrs: JLongArray, |
| schema_addrs: JLongArray, |
| num_rows: jint, |
| ) -> jni::sys::jobject { |
| try_unwrap_or_throw(&e, |mut env| { |
| // Get the context |
| let ctx = (c2r_handle as *mut ColumnarToRowContext) |
| .as_mut() |
| .ok_or_else(|| CometError::Internal("Null columnar to row context".to_string()))?; |
| |
| let num_cols = env.get_array_length(&array_addrs)? as usize; |
| |
| // Get array and schema addresses |
| let array_addrs_elements = env.get_array_elements(&array_addrs, ReleaseMode::NoCopyBack)?; |
| let schema_addrs_elements = |
| env.get_array_elements(&schema_addrs, ReleaseMode::NoCopyBack)?; |
| |
| // Import Arrow arrays from FFI |
| let mut arrays = Vec::with_capacity(num_cols); |
| for i in 0..num_cols { |
| let array_ptr = array_addrs_elements[i] as *mut FFI_ArrowArray; |
| let schema_ptr = schema_addrs_elements[i] as *mut FFI_ArrowSchema; |
| |
| // Take ownership of the FFI structures |
| let ffi_array = std::ptr::read(array_ptr); |
| let ffi_schema = std::ptr::read(schema_ptr); |
| |
| // Convert to Arrow ArrayData |
| let array_data = from_ffi(ffi_array, &ffi_schema) |
| .map_err(|e| CometError::Internal(format!("Failed to import array: {}", e)))?; |
| |
| arrays.push(arrow::array::make_array(array_data)); |
| } |
| |
| // Convert columnar to row |
| let (buffer_ptr, offsets, lengths) = ctx.convert(&arrays, num_rows as usize)?; |
| |
| // Create Java int arrays for offsets and lengths |
| let offsets_array = env.new_int_array(offsets.len() as i32)?; |
| env.set_int_array_region(&offsets_array, 0, offsets)?; |
| |
| let lengths_array = env.new_int_array(lengths.len() as i32)?; |
| env.set_int_array_region(&lengths_array, 0, lengths)?; |
| |
| // Create the NativeColumnarToRowInfo object |
| let info_class = env.find_class("org/apache/comet/NativeColumnarToRowInfo")?; |
| let info_obj = env.new_object( |
| info_class, |
| "(J[I[I)V", |
| &[ |
| jni::objects::JValue::Long(buffer_ptr as jlong), |
| jni::objects::JValue::Object(&offsets_array), |
| jni::objects::JValue::Object(&lengths_array), |
| ], |
| )?; |
| |
| Ok(info_obj.into_raw()) |
| }) |
| } |
| |
| /// Close and release the native columnar to row converter. |
| /// |
| /// # Safety |
| /// This function is inherently unsafe since it deals with raw pointers passed from JNI. |
| #[no_mangle] |
| pub unsafe extern "system" fn Java_org_apache_comet_Native_columnarToRowClose( |
| e: JNIEnv, |
| _class: JClass, |
| c2r_handle: jlong, |
| ) { |
| try_unwrap_or_throw(&e, |_env| { |
| if c2r_handle != 0 { |
| let _ctx: Box<ColumnarToRowContext> = |
| Box::from_raw(c2r_handle as *mut ColumnarToRowContext); |
| // ctx is dropped here, freeing the buffer |
| } |
| Ok(()) |
| }) |
| } |
