andygrove commented on code in PR #7299: URL: https://github.com/apache/datafusion-comet/pull/7299#discussion_r4414426001
########## native/native-executors/src/assembly.rs: ########## @@ -0,0 +1,190 @@ +// 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. + +//! Turns a driver's stage graph into a Ballista physical plan. + +use std::collections::HashMap; +use std::sync::Arc; + +use comet::execution::native_executors::{ + block_schema, partition_descriptor, CometFragmentExec, COMET_VERSION, +}; +use datafusion::common::{plan_datafusion_err, plan_err, Result}; +use datafusion::physical_expr::expressions::Column; +use datafusion::physical_expr::PhysicalExpr; +use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion::physical_plan::repartition::RepartitionExec; +use datafusion::physical_plan::ExecutionPlanProperties; +use datafusion::physical_plan::{ExecutionPlan, Partitioning}; +use datafusion_comet_proto::spark_expression::expr::ExprStruct; +use datafusion_comet_proto::spark_native_executors::{Fragment, QueryRequest, Stage, StageInput}; +use datafusion_comet_proto::spark_partitioning::partitioning::PartitioningStruct; +use datafusion_comet_proto::PROTO_FINGERPRINT; +use prost::Message; + +/// Checks that `request` comes from a driver running this Comet build. The plan format is not +/// yet stable across versions, and an executor would ignore any field it doesn't know. +pub fn check_version(request: &QueryRequest) -> Result<()> { + if request.comet_version != COMET_VERSION || request.proto_fingerprint != PROTO_FINGERPRINT { + return plan_err!( + "This Comet Native Executor cluster runs Comet {COMET_VERSION} with plan format \ + {PROTO_FINGERPRINT}, but the query comes from Comet {} with plan format {}. The \ + driver and the cluster must run the same Comet build.", + request.comet_version, + request.proto_fingerprint + ); + } + Ok(()) +} + +/// Builds the physical plan of `request`: a [`CometFragmentExec`] per stage, with a node that +/// Ballista's planner turns into a stage boundary on each edge. The root is the last stage. +/// +/// Each stage's block is planned with Comet's planner to check it before anything runs. That can +/// block, so don't call this on an async runtime's thread. A stage that several stages read runs +/// once for each of them, since Ballista's planner doesn't share a stage between consumers yet. +pub fn build_plan(request: QueryRequest) -> Result<Arc<dyn ExecutionPlan>> { + let mut stages: Vec<Arc<dyn ExecutionPlan>> = Vec::with_capacity(request.stages.len()); + for (index, stage) in request.stages.into_iter().enumerate() { + let fragment = build_stage(stage, &stages, &request.conf) + .map_err(|e| e.context(format!("Stage {index} of the query")))?; + stages.push(fragment); + } + stages + .pop() + .ok_or_else(|| plan_datafusion_err!("The query has no stages")) +} + +fn build_stage( + stage: Stage, + producers: &[Arc<dyn ExecutionPlan>], + conf: &HashMap<String, String>, +) -> Result<Arc<dyn ExecutionPlan>> { + let block = stage + .plan + .as_ref() + .ok_or_else(|| plan_datafusion_err!("The stage has no plan"))? + .encode_to_vec(); + let mut inputs = Vec::with_capacity(stage.inputs.len()); + for input in &stage.inputs { + let producer = producers.get(input.producer as usize).ok_or_else(|| { + plan_datafusion_err!( + "The stage reads stage {}, which doesn't come before it", + input.producer + ) + })?; + inputs.push(exchange(input, Arc::clone(producer))?); + } + // A stage that reads exchanges has as many partitions as they deliver, which + // `CometFragmentExec` checks they agree on. + let num_partitions = stage + .inputs + .iter() + .zip(&inputs) + .find(|(input, _)| !input.broadcast) + .map(|(_, exchange)| exchange.output_partitioning().partition_count()) + .unwrap_or(stage.num_partitions as usize); + if num_partitions == 0 { + return plan_err!("The stage has no partitions"); + } + if !stage.partition_plans.is_empty() && stage.partition_plans.len() != num_partitions { Review Comment: Agreed, that shape returns wrong rows without an error. e38d3d1b6f makes `build_stage` reject a stage that has partition plans and a non-broadcast hash input, and the doc on `exchange` now points at that check. A one-partition file stage zipped with a `SinglePartition` input still assembles. `file_stage_cannot_read_a_hash_exchange` and `file_stage_can_read_a_single_partition_exchange` in `tests/assembly.rs` cover both. ########## native/core/src/execution/native_executors/fragment.rs: ########## @@ -0,0 +1,440 @@ +// 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. + +use std::collections::HashMap; +use std::fmt; +use std::sync::Arc; + +use arrow::array::{Array, AsArray, BinaryArray, RecordBatch, RecordBatchOptions, UInt32Array}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef, UInt32Type}; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{internal_datafusion_err, DataFusionError, Result}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::execution::disk_manager::DiskManager; +use datafusion::execution::memory_pool::{MemoryPool, UnboundedMemoryPool}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::{Distribution, EquivalenceProperties, PhysicalExpr}; +use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; +use datafusion::physical_plan::streaming::{PartitionStream, StreamingTableExec}; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, + InputDistributionRequirements, Partitioning, PlanProperties, SendableRecordBatchStream, +}; +use datafusion::prelude::SessionContext; +use datafusion_comet_proto::spark_native_executors::Fragment; +use datafusion_comet_proto::spark_operator::{Operator, SparkStructField}; +use futures::stream::BoxStream; +use futures::{StreamExt, TryStreamExt}; + +use crate::execution::jni_api::{comet_runtime_env, comet_session_context}; +use crate::execution::planner::{ + align_shuffle_writer_input, convert_spark_types_to_arrow_schema, PhysicalPlanner, +}; +use crate::execution::serde::deserialize_op; +use crate::execution::spark_config::{ + SparkConfig, COMET_BATCH_SIZE, COMET_STRICT_TIMESTAMP_LABELS, +}; + +const DEFAULT_BATCH_SIZE: usize = 8192; + +fn descriptor_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("partition", DataType::UInt32, false), + Field::new("plan", DataType::Binary, true), + ])) +} + +/// The partition descriptor of a fragment with `num_partitions` partitions. +/// +/// Partition `i` is a single row holding `i`, the partition's index in its stage, and +/// `plans[i]`, the partition's own copy of the block, when `plans` is not empty. A stage that +/// reads files has one, naming the partition's files. A distributed engine that splits a stage's +/// partitions between tasks slices the descriptor along with the stage's other leaves, so a task +/// still learns each partition's index in the stage: Spark seeds `rand`, +/// `monotonically_increasing_id` and similar expressions with it. +pub fn partition_descriptor( + num_partitions: usize, + plans: &[Vec<u8>], +) -> Result<Arc<dyn ExecutionPlan>> { + let schema = descriptor_schema(); + let partitions = (0..num_partitions) + .map(|index| { + let plan = plans.get(index).map(Vec::as_slice); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![index as u32])), + Arc::new(BinaryArray::from(vec![plan])), + ], + )?; + Ok(vec![batch]) + }) + .collect::<Result<Vec<_>>>()?; + Ok(MemorySourceConfig::try_new_exec(&partitions, schema, None)?) +} + +/// Runs a Comet native block for each partition of a stage. +/// +/// The first child is the partition descriptor (see [`partition_descriptor`]), and each other +/// child is an input of the block, matched to the block's input leaves in the order Comet's +/// planner meets them. Partition `p` of the fragment reads partition `p` of each partitioned +/// input, and the single partition of each broadcast input. +#[derive(Debug)] +pub struct CometFragmentExec { + fragment: Arc<Fragment>, + /// The partition descriptor, then the inputs. + children: Vec<Arc<dyn ExecutionPlan>>, + props: Arc<PlanProperties>, +} + +impl CometFragmentExec { + pub fn try_new(fragment: Arc<Fragment>, children: Vec<Arc<dyn ExecutionPlan>>) -> Result<Self> { + if fragment + .output_schema + .iter() + .any(|field| field.data_type.is_none()) + { + return Err(internal_datafusion_err!( + "CometFragmentExec's output schema has a field without a type" + )); + } + let Some((descriptor, inputs)) = children.split_first() else { + return Err(internal_datafusion_err!( + "CometFragmentExec needs a partition descriptor" + )); + }; + if inputs.len() != fragment.broadcast.len() { + return Err(internal_datafusion_err!( + "CometFragmentExec has {} inputs but {} broadcast flags", + inputs.len(), + fragment.broadcast.len() + )); + } + let partitions = descriptor.output_partitioning().partition_count(); + for (index, (input, broadcast)) in inputs.iter().zip(&fragment.broadcast).enumerate() { + let expected = if *broadcast { 1 } else { partitions }; + let actual = input.output_partitioning().partition_count(); + if actual != expected { + return Err(internal_datafusion_err!( + "CometFragmentExec input {index} has {actual} partitions, expected {expected}" + )); + } + } + let props = Arc::new(PlanProperties::new( + EquivalenceProperties::new(fragment_schema(&fragment.output_schema)), + Partitioning::UnknownPartitioning(partitions), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Ok(Self { + fragment, + children, + props, + }) + } + + pub fn fragment(&self) -> &Arc<Fragment> { + &self.fragment + } +} + +/// The schema of a fragment's output: the types Spark declared for it, every column nullable. +/// The cast to those types keeps a column nullable wherever the block's output is, even where +/// Spark knows it has no nulls, such as a join key a filter has checked, so a fragment can't +/// promise more. Nothing between stages relies on a column having no nulls. +fn fragment_schema(output_schema: &[SparkStructField]) -> SchemaRef { + let declared = convert_spark_types_to_arrow_schema(output_schema); + let fields: Vec<Field> = declared + .fields() + .iter() + .map(|field| field.as_ref().clone().with_nullable(true)) + .collect(); + Arc::new(Schema::new(fields)) +} + +impl DisplayAs for CometFragmentExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!( + f, + "CometFragmentExec: block={} bytes, broadcast={:?}", + self.fragment.block.len(), + self.fragment.broadcast + ) + } +} + +impl ExecutionPlan for CometFragmentExec { + fn name(&self) -> &str { + "CometFragmentExec" + } + + fn properties(&self) -> &Arc<PlanProperties> { + &self.props + } + + fn children(&self) -> Vec<&Arc<dyn ExecutionPlan>> { + self.children.iter().collect() + } + + /// The descriptor's and the partitioned inputs' partitions line up with this plan's, while a + /// broadcast input is read whole by every partition, so a distributed engine must not split + /// it between tasks. + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + let inputs = self.fragment.broadcast.iter().map(|broadcast| { + if *broadcast { + Distribution::SinglePartition + } else { + Distribution::UnspecifiedDistribution + } + }); + InputDistributionRequirements::new( + std::iter::once(Distribution::UnspecifiedDistribution) + .chain(inputs) + .collect(), + ) + } + + /// Repartitioning an input for parallelism would break its alignment with the descriptor. + fn benefits_from_input_partitioning(&self) -> Vec<bool> { + vec![false; self.children.len()] + } + + /// The block's expressions are in its serialized plan, so there are none to visit here. + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc<dyn PhysicalExpr>) -> Result<TreeNodeRecursion>, + ) -> Result<TreeNodeRecursion> { + Ok(TreeNodeRecursion::Continue) + } + + fn with_new_children( + self: Arc<Self>, + children: Vec<Arc<dyn ExecutionPlan>>, + ) -> Result<Arc<dyn ExecutionPlan>> { + Ok(Arc::new(Self::try_new( + Arc::clone(&self.fragment), + children, + )?)) + } + + fn execute( + &self, + partition: usize, + context: Arc<TaskContext>, + ) -> Result<SendableRecordBatchStream> { + let descriptor = Arc::clone(&self.children[0]); + let inputs = self.children[1..] + .iter() + .zip(&self.fragment.broadcast) + .map(|(input, broadcast)| { + let input_partition = if *broadcast { 0 } else { partition }; + fragment_input(Arc::clone(input), input_partition, Arc::clone(&context)) + }) + .collect::<Result<Vec<_>>>()?; + let fragment = Arc::clone(&self.fragment); + let schema = self.schema(); + let stream = futures::stream::once(async move { + let (index, own_block) = read_descriptor(&descriptor, partition, &context).await?; + // Planning can block on Comet's own runtime, which panics on a runtime thread. + tokio::task::spawn_blocking(move || { + let block = own_block.as_deref().unwrap_or(&fragment.block); + run_block(block, index, inputs, &fragment, schema, &context) + }) + .await + .map_err(|e| DataFusionError::External(Box::new(e)))? + }) + .try_flatten(); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream, + ))) + } +} + +/// Reads partition `partition` of a fragment's descriptor: the partition's index in its stage, +/// and the partition's own copy of the block, if it has one. +async fn read_descriptor( + descriptor: &Arc<dyn ExecutionPlan>, + partition: usize, + context: &Arc<TaskContext>, +) -> Result<(u32, Option<Vec<u8>>)> { + let batches: Vec<RecordBatch> = descriptor + .execute(partition, Arc::clone(context))? + .try_collect() + .await?; + let batch = match batches.as_slice() { + [batch] if batch.num_rows() == 1 => batch, + _ => { + return Err(internal_datafusion_err!( + "A fragment's partition descriptor must have one row per partition" + )) + } + }; + let index = batch.column(0).as_primitive::<UInt32Type>().value(0); + let plans = batch.column(1).as_binary::<i32>(); + let plan = (!plans.is_null(0)).then(|| plans.value(0).to_vec()); + Ok((index, plan)) +} + +/// Plans `block` for the partition with index `index` in its stage, feeding its input leaves from +/// `inputs`, and starts it. +/// +/// Its output is cast to the fragment's output schema, the types Spark declared for it, before it +/// leaves the stage, as a Spark task's shuffle writer casts its input: a hash exchange places a +/// row by the types of its key values as well as by the values. A wrong timezone label on a +/// timestamp fails instead in strict test runs, as it does there. Each batch then carries +/// `schema`, the fragment's, which a distributed engine checks it against. +fn run_block( + block: &[u8], + index: u32, + inputs: Vec<Arc<dyn ExecutionPlan>>, + fragment: &Fragment, + schema: SchemaRef, + context: &TaskContext, +) -> Result<BoxStream<'static, Result<RecordBatch>>> { + let spark_plan = deserialize_op(block)?; + let session = fragment_session( + &spark_plan, + &fragment.conf, + Arc::clone(context.memory_pool()), + Some(Arc::clone(&context.runtime_env().disk_manager)), + )?; + let planner = + PhysicalPlanner::new(Arc::clone(&session), index as i32).with_native_inputs(inputs); + let root = plan_block(&spark_plan, planner)?; + let root = align_shuffle_writer_input( + root, + &fragment.output_schema, + fragment.conf.get_bool(COMET_STRICT_TIMESTAMP_LABELS), + )?; + let stream = root.execute(0, session.task_ctx())?; + // The session holds state that the block's operators may use while they run. + Ok(stream + .map(move |batch| { + let _session = &session; + let batch = batch?; + // The cast made the names and types the declared ones, so this changes only + // nullability and metadata. + let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows())); + Ok(RecordBatch::try_new_with_options( + Arc::clone(&schema), + batch.columns().to_vec(), + &options, + )?) + }) + .boxed()) +} + +/// The Comet session a block runs in. It reserves memory from `memory_pool` and spills through +/// `disk_manager`, the host's, so that spill files go where the host keeps them. +fn fragment_session( + spark_plan: &Operator, + conf: &HashMap<String, String>, + memory_pool: Arc<dyn MemoryPool>, + disk_manager: Option<Arc<DiskManager>>, +) -> Result<Arc<SessionContext>> { + let mut runtime = comet_runtime_env(memory_pool); + runtime.disk_manager = disk_manager; + let session = comet_session_context( + runtime, + conf.get_usize(COMET_BATCH_SIZE, DEFAULT_BATCH_SIZE), + 1, + conf, + spark_plan, + ) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + Ok(Arc::new(session)) +} + +fn plan_block(spark_plan: &Operator, planner: PhysicalPlanner) -> Result<Arc<dyn ExecutionPlan>> { + let planner = planner.with_sql_text_pool(spark_plan); + let (_, _, root) = planner.create_plan(spark_plan, &mut vec![], 1)?; + Ok(Arc::clone(&root.native_plan)) +} + +/// The output schema of `block`, found by planning it without its data: each input leaf is empty, +/// and each file scan builds no object store for its files, since nothing reads them, so this +/// uses no storage credentials. +/// +/// Fails where running the block would fail at planning, such as for an operator or expression +/// this build doesn't support, or one that calls into the JVM. A problem with the files or their +/// store shows only when the block runs. Planning can block on Comet's own runtime, so don't call +/// this on an async runtime's thread. +pub fn block_schema(block: &[u8], conf: &HashMap<String, String>) -> Result<SchemaRef> { Review Comment: Good catch. fcb23120db makes a planner that runs without a JVM (the scheduler check and the executors) refuse `NativeScalarUdf`, the way `require_jvm` refuses JVM UDFs, so the library path in a request is never loaded there. I preferred that to a process-owned policy for the first PR. Native UDFs on these processes can come back with an allowlist that comes from the process environment. `native_scalar_udf_is_refused_on_native_executors` in the planner tests covers both the checking and the running planner. ########## native/core/src/execution/planner.rs: ########## @@ -312,6 +331,92 @@ impl PhysicalPlanner { self } + /// Feeds the block's input leaves from `inputs` instead of the JVM, one per leaf in the order + /// the planner meets them. Each input's output is cast to the types Spark declared for its + /// leaf. Expressions that call into the JVM fail to plan, since there is no JVM. + pub fn with_native_inputs(mut self, inputs: Vec<Arc<dyn ExecutionPlan>>) -> Self { + self.native_inputs = Some(Arc::new(NativeInputs::Plans(Mutex::new(inputs.into())))); + self + } + + /// Like [`Self::with_native_inputs`], but plans a block without its data, to check it and + /// learn its schema: each input leaf is empty, and each file scan is planned with its files + /// but builds no object store for them, since nothing reads them. Planning then uses no + /// storage credentials. + pub fn without_data(mut self) -> Self { + self.native_inputs = Some(Arc::new(NativeInputs::Empty)); + self + } + + /// Whether the planner plans the block without its data; see [`Self::without_data`]. + fn plans_without_data(&self) -> bool { + matches!(self.native_inputs.as_deref(), Some(NativeInputs::Empty)) + } + + /// Registers the object store a scan reads `file` through, configured by `options`, and + /// returns its URL and backend. When planning without data, the scan gets the local file + /// system instead, which is always registered, and no store is built. + fn scan_store( + &self, + file: &str, + options: &HashMap<String, String>, + ) -> Result<(ObjectStoreUrl, ObjectStoreBackend), ExecutionError> { + if self.plans_without_data() { + return Ok(( + ObjectStoreUrl::local_filesystem(), + ObjectStoreBackend::Local, + )); + } + let (url, _, backend) = prepare_object_store_with_configs( Review Comment: Done in fcb23120db: `JVMClasses::with_env` takes `JAVA_VM.get()` and returns a `CometError::Internal` when no JVM was started, so the S3 credential bridge, `encryption_support` and anything else missed here fail with an error instead of undefined behavior. `with_env_fails_without_a_jvm` in `jni-bridge` covers it. ########## native/native-executors/Cargo.toml: ########## @@ -0,0 +1,61 @@ +# 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. + +# Comet Native Executors: the comet-scheduler and comet-executor binaries, built on Apache +# DataFusion Ballista. See https://github.com/apache/datafusion-comet/issues/7288. +# +# This is its own workspace, with its own lockfile, so that the main Comet build never depends +# on Ballista. It builds Comet's native crates by path together with the Ballista release on the +# same DataFusion major version. +[workspace] Review Comment: Added in 53ea7d4f39: a `cargo` entry for `/native/native-executors` with the same ignore and grouping rules as `/native`. I left the lock as it is, so it picks up the arrow 59.3.x patch that `native/Cargo.lock` has on the first Dependabot run, and the README now says the two lockfiles can differ between updates. -- 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]
