diff --git a/docs/source/contributor-guide/expression-audits/generator_funcs.md b/docs/source/contributor-guide/expression-audits/generator_funcs.md index 3023cfaa774..140830c346d 100644 --- a/docs/source/contributor-guide/expression-audits/generator_funcs.md +++ b/docs/source/contributor-guide/expression-audits/generator_funcs.md @@ -37,4 +37,16 @@ - Same `CometExplodeExec` path as `posexplode`. Compatible for array inputs; empty and NULL arrays both emit one row with null `pos` and null `value` per Spark's `outer` semantics, which the planner requests as DataFusion's `NullHandling::PreserveAndExpandEmpty`. +## inline + +- Same `CometExplodeExec` path as `explode`, with the unnested struct flattened into one column per field. Spark 3.4.3, 3.5.8, 4.0.1, and 4.1.1 share one implementation: a NULL struct element yields a row of NULLs, which the native flattening reproduces by folding the struct's validity into each field. Every generated column is reported nullable. + +## inline_outer + +- Same path as `inline`. Empty and NULL arrays both emit one row of NULLs, which the planner requests as DataFusion's `NullHandling::PreserveAndExpandEmpty`. + +## stack + +- Planned as `inline` over an array holding one struct per output row, so `stack(n, v1, ..., vk)` evaluates each argument once per input row, as Spark does, and pads the last row with NULLs. Spark 3.4.3, 3.5.8, 4.0.1, and 4.1.1 share one implementation. An argument that cannot run natively falls back with that argument's own reason. + [Spark Expression Support]: ../../user-guide/latest/expressions.md diff --git a/docs/source/user-guide/latest/expressions.md b/docs/source/user-guide/latest/expressions.md index cf8baac011d..22c35e6ef4c 100644 --- a/docs/source/user-guide/latest/expressions.md +++ b/docs/source/user-guide/latest/expressions.md @@ -336,20 +336,22 @@ The type-name conversion functions (`bigint`, `binary`, `boolean`, `date`, `deci ## generator_funcs -`explode`, `explode_outer`, `posexplode`, and `posexplode_outer` are supported via -`CometExplodeExec` (operator-level, not expression-level) for array input; map input falls back -to Spark ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enabled by default via -`spark.comet.exec.explode.enabled`. +`explode`, `explode_outer`, `posexplode`, `posexplode_outer`, `inline`, `inline_outer`, and +`stack` are supported via `CometExplodeExec` (operator-level, not expression-level). `explode` and +`posexplode` support array input; map input falls back to Spark +([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enabled by default via +`spark.comet.exec.explode.enabled`. Other deterministic generators run on the JVM inside +`CometExplodeExec`; see [Generators](scala_java_udfs.md#generators). | Function | Status | Implementation | Notes | | --- | --- | --- | --- | | `explode` | ✅ | — | via `CometExplodeExec` | | `explode_outer` | ✅ | — | via `CometExplodeExec` ([audit](../../contributor-guide/expression-audits/generator_funcs.md#explode_outer)) | -| `inline` | 🔜 | — | Operator-level generator (like `explode`) | -| `inline_outer` | 🔜 | — | Operator-level generator (like `explode`) | +| `inline` | ✅ | — | via `CometExplodeExec` ([audit](../../contributor-guide/expression-audits/generator_funcs.md#inline)) | +| `inline_outer` | ✅ | — | via `CometExplodeExec` ([audit](../../contributor-guide/expression-audits/generator_funcs.md#inline_outer)) | | `posexplode` | ✅ | — | via `CometExplodeExec` | | `posexplode_outer` | ✅ | — | via `CometExplodeExec` ([audit](../../contributor-guide/expression-audits/generator_funcs.md#posexplode_outer)) | -| `stack` | 🔜 | — | Operator-level generator | +| `stack` | ✅ | — | via `CometExplodeExec`, planned as `inline` ([audit](../../contributor-guide/expression-audits/generator_funcs.md#stack)) | --- @@ -375,7 +377,7 @@ to Spark ([#2837](https://github.com/apache/datafusion-comet/issues/2837)). Enab | `get_json_object` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default; the native path is opt-in via allowIncompatible ([audit](../../contributor-guide/expression-audits/json_funcs.md#get_json_object)) | | `json_array_length` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default; the native path, which differs for single-quoted JSON, unescaped control characters, and trailing content, is opt-in via allowIncompatible ([audit](../../contributor-guide/expression-audits/json_funcs.md#json_array_length)) | | `json_object_keys` | ✅ | Codegen dispatch | | -| `json_tuple` | 🔜 | — | [#3160](https://github.com/apache/datafusion-comet/issues/3160) | +| `json_tuple` | ✅ | — | Runs on the JVM inside `CometExplodeExec` ([Generators](scala_java_udfs.md#generators)); [#3160](https://github.com/apache/datafusion-comet/issues/3160) tracks a native implementation | | `schema_of_json` | ✅ | Codegen dispatch | | | `to_json` | ✅ | Hybrid | Routes through the JVM codegen dispatcher by default, including options and map/array inputs; the native path is opt-in via allowIncompatible ([audit](../../contributor-guide/expression-audits/json_funcs.md#to_json)) | diff --git a/docs/source/user-guide/latest/operators.md b/docs/source/user-guide/latest/operators.md index e96df28b126..13bb28fc4e9 100644 --- a/docs/source/user-guide/latest/operators.md +++ b/docs/source/user-guide/latest/operators.md @@ -125,12 +125,12 @@ combined with `OR`, which Spark plans as an existence `BroadcastNestedLoopJoinEx ## Generators and set operations -| Operator | Status | Notes | -| -------------- | ------ | ---------------------------------------------------------------------------------------------------------------- | -| `GenerateExec` | ✅ | Supports `explode`, `explode_outer`, `posexplode`, `posexplode_outer` over arrays. `inline` / `stack` fall back. | -| `ExpandExec` | ✅ | | -| `UnionExec` | ✅ | | -| `CoalesceExec` | ✅ | | +| Operator | Status | Notes | +| -------------- | ------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `GenerateExec` | ✅ | Supports `explode`, `explode_outer`, `posexplode`, `posexplode_outer` over arrays, `inline`, `inline_outer`, and `stack`. Other deterministic generators, including Hive UDTFs, run on the JVM inside the operator; see [Generators](scala_java_udfs.md#generators). | +| `ExpandExec` | ✅ | | +| `UnionExec` | ✅ | | +| `CoalesceExec` | ✅ | | ## Writes diff --git a/docs/source/user-guide/latest/scala_java_udfs.md b/docs/source/user-guide/latest/scala_java_udfs.md index 43a30374578..6ebf2e094bb 100644 --- a/docs/source/user-guide/latest/scala_java_udfs.md +++ b/docs/source/user-guide/latest/scala_java_udfs.md @@ -21,7 +21,7 @@ Comet executes Spark's Scala and Java [scalar user-defined functions (UDFs)](https://spark.apache.org/docs/latest/sql-ref-functions-udf-scalar.html) within the Comet pipeline. The presence of a UDF does not force the enclosing operator out of the Comet pipeline; surrounding Rust-implemented operators stay in the pipeline. -This page covers Spark's `ScalaUDF` (Scala `udf(...)`, `spark.udf.register(...)` over Scala or Java functional interfaces, and SQL `CREATE FUNCTION ... AS 'com.example.MyUDF'`). Other UDF kinds (Python / Pandas, Hive, aggregate) are out of scope and continue to fall back to Spark. +This page covers Spark's `ScalaUDF` (Scala `udf(...)`, `spark.udf.register(...)` over Scala or Java functional interfaces, and SQL `CREATE FUNCTION ... AS 'com.example.MyUDF'`). Other UDF kinds (Python / Pandas, Hive scalar UDFs, aggregate) are out of scope and continue to fall back to Spark. Hive UDTFs and other generators are covered in [Generators](#generators). This feature is enabled by default. Set `spark.comet.exec.scalaUDF.codegen.enabled` to `false` to route plans containing a `ScalaUDF` back to Spark for the enclosing operator. @@ -42,7 +42,7 @@ This feature is enabled by default. Set `spark.comet.exec.scalaUDF.codegen.enabl ## Not supported - Aggregate UDFs (`ScalaAggregator`, `TypedImperativeAggregate`, the legacy `UserDefinedAggregateFunction`). -- Table UDFs and generators. +- Python UDTFs (`@udtf`). - Python `@udf` and Pandas `@pandas_udf`. - Hive `GenericUDF` and `SimpleUDF`. - `NullType` and `UserDefinedType` arguments and return types. UDT-typed columns fall back to Spark; to keep execution in the Comet pipeline, store and read the underlying representation directly (e.g. write MLlib `Vector` outputs as `Struct, values: Array>` rather than `VectorUDT`). @@ -50,6 +50,16 @@ This feature is enabled by default. Set `spark.comet.exec.scalaUDF.codegen.enabl When a UDF is rejected, the reason surfaces through Comet's standard fallback diagnostics; the query still runs on Spark. +## Generators + +A deterministic generator that Comet has no native implementation of, such as a Hive `GenericUDTF` registered with `CREATE FUNCTION` or Spark's `json_tuple`, runs inside Comet's native generate operator. The operator calls the generator on the JVM for each batch and keeps one generator instance per partition, so state the generator holds carries across batches. Once a partition's input ends, the rows of the generator's `terminate()` (a Hive UDTF's `close()`) follow, carrying the partition's last input row, as in Spark. + +The operator works a batch at a time rather than a row at a time. It evaluates the generator for every row of a batch, and holds all the rows one input row generates, before it passes any of them on. A `LIMIT` above the generator does not stop evaluation partway through a batch, so a generator that raises an error on a later row can fail a query Spark would have satisfied from earlier rows, and a generator that emits a very large number of rows for one input row needs memory for all of them. + +Spark's `EXCEPT ALL` and `INTERSECT ALL` use an internal generator, `ReplicateRows`. Comet runs it natively as `inline` over `array_repeat`, without the JVM. + +Set `spark.comet.exec.jvmGenerator.enabled` to `false` to run such a generator in Spark instead. + ## Behavior - Non-deterministic expressions referenced from the argument tree (`rand`, `uuid`, `monotonically_increasing_id`) produce per-partition sequences consistent with Spark. diff --git a/docs/source/user-guide/latest/understanding-comet-plans.md b/docs/source/user-guide/latest/understanding-comet-plans.md index d86f8076bde..2b7bcda3b01 100644 --- a/docs/source/user-guide/latest/understanding-comet-plans.md +++ b/docs/source/user-guide/latest/understanding-comet-plans.md @@ -314,28 +314,28 @@ own Rust plan over its child's output. `CometTakeOrderedAndProject` is another: Rust plan, and when its input has several partitions, it first gathers the top rows of each partition through a single-partition Comet shuffle. -| Node | Spark equivalent | -| ------------------------------ | --------------------------------------------------------------------------------- | -| `CometProject` | `ProjectExec` | -| `CometFilter` | `FilterExec` | -| `CometSort` | `SortExec` | -| `CometLocalLimit` | `LocalLimitExec` | -| `CometGlobalLimit` | `GlobalLimitExec` | -| `CometExpand` | `ExpandExec` | -| `CometExplode` | `GenerateExec` (for `explode`, `explode_outer`, `posexplode`, `posexplode_outer`) | -| `CometHashAggregate` | `HashAggregateExec`, `ObjectHashAggregateExec` | -| `CometHashJoin` | `ShuffledHashJoinExec` | -| `CometBroadcastHashJoin` | `BroadcastHashJoinExec` | -| `CometBroadcastNestedLoopJoin` | `BroadcastNestedLoopJoinExec` | -| `CometSortMergeJoin` | `SortMergeJoinExec` | -| `CometMergeRows` | `MergeRowsExec` (Spark 3.5+; opt-in) | -| `CometWindowExec` | `WindowExec` | -| `CometWindowGroupLimitExec` | `WindowGroupLimitExec` (Spark 3.5 and later) | -| `CometSample` | `SampleExec` (sampling without replacement) | -| `CometTakeOrderedAndProject` | `TakeOrderedAndProjectExec` | -| `CometWriteFiles` | `WriteFilesExec` (Spark 4.0 and later, experimental Rust Parquet writer) | -| `CometNativeWrite` | `DataWritingCommandExec` (Spark 3.x, experimental Rust Parquet writer) | -| `CometIcebergWrite` | `IcebergWrite` (native Iceberg data-file writes, opt-in) | +| Node | Spark equivalent | +| ------------------------------ | ----------------------------------------------------------------------------------- | +| `CometProject` | `ProjectExec` | +| `CometFilter` | `FilterExec` | +| `CometSort` | `SortExec` | +| `CometLocalLimit` | `LocalLimitExec` | +| `CometGlobalLimit` | `GlobalLimitExec` | +| `CometExpand` | `ExpandExec` | +| `CometExplode` | `GenerateExec`; generators without a native implementation run on the JVM inside it | +| `CometHashAggregate` | `HashAggregateExec`, `ObjectHashAggregateExec` | +| `CometHashJoin` | `ShuffledHashJoinExec` | +| `CometBroadcastHashJoin` | `BroadcastHashJoinExec` | +| `CometBroadcastNestedLoopJoin` | `BroadcastNestedLoopJoinExec` | +| `CometSortMergeJoin` | `SortMergeJoinExec` | +| `CometMergeRows` | `MergeRowsExec` (Spark 3.5+; opt-in) | +| `CometWindowExec` | `WindowExec` | +| `CometWindowGroupLimitExec` | `WindowGroupLimitExec` (Spark 3.5 and later) | +| `CometSample` | `SampleExec` (sampling without replacement) | +| `CometTakeOrderedAndProject` | `TakeOrderedAndProjectExec` | +| `CometWriteFiles` | `WriteFilesExec` (Spark 4.0 and later, experimental Rust Parquet writer) | +| `CometNativeWrite` | `DataWritingCommandExec` (Spark 3.x, experimental Rust Parquet writer) | +| `CometIcebergWrite` | `IcebergWrite` (native Iceberg data-file writes, opt-in) | ### JVM-Side Operators diff --git a/native/core/src/execution/planner.rs b/native/core/src/execution/planner.rs index 603bba872bc..89d21ebd7cd 100644 --- a/native/core/src/execution/planner.rs +++ b/native/core/src/execution/planner.rs @@ -5317,6 +5317,8 @@ mod tests { outer, position, project_list: vec![bound], + inline: false, + terminate: None, })), ..Default::default() }; diff --git a/native/core/src/execution/planner/explode.rs b/native/core/src/execution/planner/explode.rs index 1159f242a53..953e9e534f9 100644 --- a/native/core/src/execution/planner/explode.rs +++ b/native/core/src/execution/planner/explode.rs @@ -15,8 +15,10 @@ // specific language governing permissions and limitations // under the License. -//! Builder for the native `Explode` operator, which plans Spark's `explode`, `posexplode`, and -//! their `_outer` variants as a projection followed by Comet's `ExplodeExec`. +//! Builder for the native `Explode` operator, which plans Spark's `explode`, `posexplode`, +//! `inline`, and their `_outer` variants as a projection followed by Comet's `ExplodeExec`. +//! Spark's `stack` arrives here as `inline` over an array of structs built by the JVM serde, and +//! any other generator as `inline` over the rows a JVM `CometUDF` produces for each input row. use std::sync::Arc; @@ -137,8 +139,16 @@ impl OperatorBuilder for ExplodeBuilder { let mut output_fields: Vec = Vec::new(); // Add all projection columns (non-array columns) + // The rows `terminate` yields carry the last input row's columns, which are all null when + // the input is empty, so those columns are nullable whatever the input says. + let carried_nullable = explode.terminate.is_some(); for i in 0..projections.len() { - output_fields.push(project_schema.field(i).clone()); + let field = project_schema.field(i).clone(); + output_fields.push(if carried_nullable { + field.with_nullable(true) + } else { + field + }); } let array_input_index = if explode.position { @@ -162,11 +172,30 @@ impl OperatorBuilder for ExplodeBuilder { } }; - output_fields.push(Field::new( - array_field.name(), - element_type, - true, // Element is nullable after unnesting - )); + // `inline` flattens each struct element into one column per field. Every field is + // nullable after unnesting: a null struct element, or the row `inline_outer` emits for a + // NULL or empty array, yields a row of nulls. + let struct_column_indices = if explode.inline { + let DataType::Struct(struct_fields) = &element_type else { + return Err(ExecutionError::GeneralError(format!( + "Expected List of Struct for inline, got {:?}", + array_field.data_type() + ))); + }; + output_fields.extend( + struct_fields + .iter() + .map(|f| f.as_ref().clone().with_nullable(true)), + ); + vec![array_input_index] + } else { + output_fields.push(Field::new( + array_field.name(), + element_type, + true, // Element is nullable after unnesting + )); + vec![] + }; let output_schema = Arc::new(Schema::new(output_fields)); @@ -193,13 +222,18 @@ impl OperatorBuilder for ExplodeBuilder { let unnest_options = UnnestOptions::new().with_null_handling(null_handling); // Comet's specialized fork of `UnnestExec`; see `operators::explode`. - let unnest_exec = Arc::new(ExplodeExec::new( + let mut unnest_exec = ExplodeExec::new( project_exec, list_unnests, - vec![], // No struct columns to unnest + struct_column_indices, output_schema, unnest_options, - )?); + )?; + if let Some(terminate) = &explode.terminate { + let terminate = planner.create_expr(terminate, Arc::new(Schema::empty()))?; + unnest_exec = unnest_exec.with_terminate(terminate)?; + } + let unnest_exec = Arc::new(unnest_exec); Ok(( scans, diff --git a/native/operators/src/explode.rs b/native/operators/src/explode.rs index 6aae0e2d349..dcddb37fd3e 100644 --- a/native/operators/src/explode.rs +++ b/native/operators/src/explode.rs @@ -51,7 +51,9 @@ //! * dropping upstream's `ListUnnest` declaration in favor of importing the public one; //! * `find_longest_length` applies the empty-list bump once after the row-wise maximum rather //! than once per array, which is equivalent because `max` is associative; -//! * `list_output_lens` and the contiguous-run fast path described above. +//! * `list_output_lens` and the contiguous-run fast path described above; +//! * `flatten_struct_cols` folds a struct's null mask into each field it flattens, since the +//! field values under a null struct are unspecified and Spark's `inline` emits nulls there. //! //! `ExplodeExec` and `ExplodeStream` were always Comet's own. @@ -62,16 +64,16 @@ use arrow::array::{ }; use arrow::compute::kernels::length::length; use arrow::compute::kernels::zip::zip; -use arrow::compute::{cast, is_not_null, kernels, sum}; -use arrow::datatypes::{DataType, Int64Type, SchemaRef}; -use arrow::record_batch::RecordBatch; +use arrow::compute::{cast, is_not_null, kernels, sum, take_record_batch}; +use arrow::datatypes::{DataType, Int64Type, Schema, SchemaRef, UInt32Type}; +use arrow::record_batch::{RecordBatch, RecordBatchOptions}; // Upstream imports these as `arrow_ord::cmp::{eq, lt}`; Comet reaches them through `arrow`, // which does not have `arrow_ord` as a direct dependency. use arrow::compute::kernels::cmp::{eq, lt}; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{ - exec_datafusion_err, exec_err, internal_err, Constraints, HashMap, HashSet, Result, - UnnestOptions, + exec_datafusion_err, exec_err, internal_err, Constraints, HashMap, HashSet, NullHandling, + Result, UnnestOptions, }; use datafusion::execution::TaskContext; use datafusion::physical_expr::equivalence::ProjectionMapping; @@ -82,6 +84,7 @@ use datafusion::physical_plan::metrics::{ SplitMetrics, }; use datafusion::physical_plan::stream::BatchSplitStream; +use datafusion_comet_common::children_with_parent_nulls; // `ListUnnest` is the one item the copied region below does NOT need to duplicate: unlike the // kernels, upstream exports it publicly. use datafusion::physical_plan::unnest::ListUnnest; @@ -104,6 +107,8 @@ pub struct ExplodeExec { list_column_indices: Vec, struct_column_indices: Vec, options: UnnestOptions, + /// See [`ExplodeExec::with_terminate`]. + terminate: Option>, metrics: ExecutionPlanMetricsSet, cache: Arc, } @@ -129,11 +134,24 @@ impl ExplodeExec { list_column_indices, struct_column_indices, options, + terminate: None, metrics: ExecutionPlanMetricsSet::new(), cache: Arc::new(cache), }) } + /// Unnest one more row once the input is exhausted: `terminate`, evaluated against a one-row + /// batch with no columns, takes the place of the list column in the last input row. This is + /// how a JVM generator's `terminate()` output joins the last row's other columns, as in + /// Spark's `GenerateExec`. Requires exactly one list column to unnest. + pub fn with_terminate(mut self, terminate: Arc) -> Result { + if self.list_column_indices.len() != 1 { + return internal_err!("ExplodeExec terminate requires exactly one list column"); + } + self.terminate = Some(terminate); + Ok(self) + } + /// Compute the plan properties, keeping whatever the child guarantees about the columns /// that unnesting passes through untouched. /// @@ -232,11 +250,13 @@ impl ExecutionPlan for ExplodeExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&Arc) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // ExplodeExec holds no physical expressions of its own; the columns to unnest are - // identified by index, so there is nothing to visit here. - Ok(TreeNodeRecursion::Continue) + // The columns to unnest are identified by index, so `terminate` is the only expression. + match &self.terminate { + Some(terminate) => f(terminate), + None => Ok(TreeNodeRecursion::Continue), + } } fn with_new_children( @@ -246,13 +266,17 @@ impl ExecutionPlan for ExplodeExec { if children.len() != 1 { return internal_err!("ExplodeExec expects exactly one child"); } - Ok(Arc::new(ExplodeExec::new( + let mut exec = ExplodeExec::new( Arc::clone(&children[0]), self.list_column_indices.clone(), self.struct_column_indices.clone(), Arc::clone(&self.schema), self.options.clone(), - )?)) + )?; + if let Some(terminate) = &self.terminate { + exec = exec.with_terminate(Arc::clone(terminate))?; + } + Ok(Arc::new(exec)) } fn execute( @@ -274,6 +298,8 @@ impl ExecutionPlan for ExplodeExec { input_rows: MetricBuilder::new(&self.metrics).counter("input_rows", partition), batch_size, pending_input: None, + terminate: self.terminate.clone(), + last_row: None, }); // Chunking the input bounds each build to roughly `batch_size` rows, but two cases can @@ -369,6 +395,10 @@ struct ExplodeStream { /// longer than `batch_size`, and recursive unnesting (where the expansion cannot be /// predicted up front), both still materialize their full expansion in one build. pending_input: Option, + /// See [`ExplodeExec::with_terminate`]. Taken once it has run. + terminate: Option>, + /// The last input row, copied out of its batch, kept while `terminate` is pending. + last_row: Option, } impl RecordBatchStream for ExplodeStream { @@ -437,6 +467,14 @@ impl ExplodeStream { self.input_rows.add(batch.num_rows()); if batch.num_rows() > 0 { let timer = self.baseline_metrics.elapsed_compute().timer(); + if self.terminate.is_some() { + // A copy rather than a slice, so the rest of the batch is not held + // until the input ends. + let last = PrimitiveArray::::from(vec![ + batch.num_rows() as u32 - 1, + ]); + self.last_row = Some(take_record_batch(&batch, &last)?); + } let output_lens = self.predict_output_lens(&batch); timer.done(); self.pending_input = Some(PendingInput { @@ -455,6 +493,18 @@ impl ExplodeStream { if other.is_none() { let input_schema = self.input.schema(); self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + if let Some(terminate) = self.terminate.take() { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = self.terminate_batch(&terminate); + timer.done(); + if let Some(batch) = result.transpose() { + if let Ok(batch) = &batch { + batch.record_output(&self.baseline_metrics); + } + return Poll::Ready(Some(batch)); + } + } } return Poll::Ready(other); } @@ -462,6 +512,56 @@ impl ExplodeStream { } } + /// Unnest the rows `terminate` yields in place of the last input row's list. With no input + /// rows, the other columns are null. Spark's `outer` handling does not apply to these rows, so + /// an empty result adds nothing. + fn terminate_batch( + &mut self, + terminate: &Arc, + ) -> Result> { + let options = RecordBatchOptions::new().with_row_count(Some(1)); + let no_columns = + RecordBatch::try_new_with_options(Arc::new(Schema::empty()), vec![], &options)?; + let rows = terminate.evaluate(&no_columns)?.into_array(1)?; + + // Nothing to emit, so there is no need to build the placeholder row below. That row's + // columns are null when no input row arrived, which a non-nullable column rejects. + let list = rows.as_list::(); + if list.is_null(0) || list.value_length(0) == 0 { + return Ok(None); + } + + // The output schema marks every column terminate may null as nullable (see the planner), + // so the placeholder batch has to as well. + let input_schema = Arc::new(Schema::new( + self.input + .schema() + .fields() + .iter() + .map(|f| f.as_ref().clone().with_nullable(true)) + .collect::>(), + )); + let mut columns = match self.last_row.take() { + Some(row) => row.columns().to_vec(), + None => input_schema + .fields() + .iter() + .map(|f| new_null_array(f.data_type(), 1)) + .collect(), + }; + columns[self.list_type_columns[0].index_in_input_schema] = rows; + let batch = RecordBatch::try_new(input_schema, columns)?; + + build_batch( + &batch, + &self.schema, + &self.list_type_columns, + &self.struct_column_indices, + &UnnestOptions::new().with_null_handling(NullHandling::Drop), + None, + ) + } + /// Compute how many output rows each input row of `batch` will expand into, so the input /// can be chunked to keep each build bounded. /// @@ -573,7 +673,9 @@ fn flatten_struct_cols( Some(_) => match column_data.data_type() { DataType::Struct(_) => { let struct_arr = column_data.as_any().downcast_ref::().unwrap(); - columns_expanded.extend_from_slice(struct_arr.columns()); + // Arrow leaves the field values under a null struct unspecified, so they + // must not leak out once the struct is flattened. + columns_expanded.extend(children_with_parent_nulls(struct_arr)?); } data_type => { return internal_err!( @@ -1569,6 +1671,8 @@ mod tests { input_rows: MetricBuilder::new(&metrics).counter("input_rows", 0), batch_size: 4, pending_input: None, + terminate: None, + last_row: None, }; let mut rows = 0; @@ -1872,4 +1976,200 @@ mod tests { ); assert_output_lens(with_nulls.slice(1, 3), [&[0, 0, 3], &[1, 0, 3], &[1, 1, 3]]); } + + /// Explode column "l" of `(id, l)` batches, ending with the rows `terminate` yields. + async fn explode_with_terminate( + input: Vec, + terminate: Vec, + null_handling: NullHandling, + ) -> Vec<(Option, Option)> { + let input_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, true), + Field::new("l", DataType::new_list(DataType::Int32, true), true), + ])); + let source = + MemorySourceConfig::try_new_exec(&[input], Arc::clone(&input_schema), None).unwrap(); + let output_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, true), + Field::new("l", DataType::Int32, true), + ])); + let rows = ListArray::from_iter_primitive::(vec![Some( + terminate.into_iter().map(Some).collect::>(), + )]); + let terminate = Arc::new(datafusion::physical_expr::expressions::Literal::new( + datafusion::common::ScalarValue::List(Arc::new(rows)), + )); + let explode = ExplodeExec::new( + source, + vec![ListUnnest { + index_in_input_schema: 1, + depth: 1, + }], + vec![], + output_schema, + UnnestOptions::new().with_null_handling(null_handling), + ) + .unwrap() + .with_terminate(terminate) + .unwrap(); + let batches = + datafusion::physical_plan::common::collect(explode.execute(0, Arc::default()).unwrap()) + .await + .unwrap(); + batches + .iter() + .flat_map(|b| { + let id = b.column(0).as_primitive::(); + let l = b.column(1).as_primitive::(); + id.iter().zip(l.iter()).collect::>() + }) + .collect() + } + + fn id_list_batch(rows: Vec<(i32, Vec)>) -> RecordBatch { + let (ids, lists): (Vec<_>, Vec<_>) = rows.into_iter().unzip(); + let lists = ListArray::from_iter_primitive::( + lists + .into_iter() + .map(|l| Some(l.into_iter().map(Some).collect::>())), + ); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, true), + Field::new("l", lists.data_type().clone(), true), + ])); + RecordBatch::try_new( + schema, + vec![Arc::new(Int32Array::from(ids)), Arc::new(lists)], + ) + .unwrap() + } + + #[tokio::test] + async fn terminate_rows_carry_the_last_input_row() { + let input = vec![ + id_list_batch(vec![(1, vec![10]), (2, vec![])]), + id_list_batch(vec![(3, vec![30, 31])]), + ]; + assert_eq!( + explode_with_terminate(input, vec![100, 101], NullHandling::Drop).await, + [ + (Some(1), Some(10)), + (Some(3), Some(30)), + (Some(3), Some(31)), + (Some(3), Some(100)), + (Some(3), Some(101)), + ] + ); + } + + #[tokio::test] + async fn terminate_without_input_rows_has_null_columns() { + let input = vec![id_list_batch(vec![])]; + assert_eq!( + explode_with_terminate(input, vec![100], NullHandling::Drop).await, + [(None, Some(100))] + ); + } + + #[tokio::test] + async fn empty_terminate_adds_no_outer_row() { + // `outer` keeps an input row whose list is empty, but not an empty `terminate()`. + let input = vec![id_list_batch(vec![(1, vec![])])]; + assert_eq!( + explode_with_terminate(input, vec![], NullHandling::PreserveAndExpandEmpty).await, + [(Some(1), None)] + ); + } + + /// `terminate` over an empty input whose carried column is non-nullable, with the nullable + /// output column the planner declares for it. + async fn terminate_over_empty_non_nullable_input(terminate: Vec) -> usize { + let input_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("l", DataType::new_list(DataType::Int32, true), true), + ])); + let source = MemorySourceConfig::try_new_exec( + &[vec![RecordBatch::new_empty(Arc::clone(&input_schema))]], + input_schema, + None, + ) + .unwrap(); + let output_schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, true), + Field::new("l", DataType::Int32, true), + ])); + let rows = ListArray::from_iter_primitive::(vec![Some( + terminate.into_iter().map(Some).collect::>(), + )]); + let terminate = Arc::new(datafusion::physical_expr::expressions::Literal::new( + datafusion::common::ScalarValue::List(Arc::new(rows)), + )); + let explode = ExplodeExec::new( + source, + vec![ListUnnest { + index_in_input_schema: 1, + depth: 1, + }], + vec![], + output_schema, + UnnestOptions::new().with_null_handling(NullHandling::Drop), + ) + .unwrap() + .with_terminate(terminate) + .unwrap(); + datafusion::physical_plan::common::collect(explode.execute(0, Arc::default()).unwrap()) + .await + .unwrap() + .iter() + .map(|b| b.num_rows()) + .sum() + } + + #[tokio::test] + async fn empty_terminate_over_empty_non_nullable_input_adds_nothing() { + assert_eq!(terminate_over_empty_non_nullable_input(vec![]).await, 0); + } + + #[tokio::test] + async fn terminate_over_empty_non_nullable_input_emits_null_carried_column() { + assert_eq!(terminate_over_empty_non_nullable_input(vec![7, 8]).await, 2); + } + + #[test] + fn flatten_struct_cols_nulls_fields_under_a_null_struct() { + // Row 1's struct is null but its fields hold values, which Arrow allows. Spark's `inline` + // turns a null struct into a row of nulls, so none of those values may surface. + let a: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(2), None, Some(4)])); + let b: ArrayRef = Arc::new(Int32Array::from(vec![10, 20, 30, 40])); + let fields = vec![ + Arc::new(Field::new("a", DataType::Int32, true)), + Arc::new(Field::new("b", DataType::Int32, false)), + ]; + let structs: ArrayRef = Arc::new(StructArray::new( + fields.clone().into(), + vec![a, b], + Some(NullBuffer::from(vec![true, false, true, true])), + )); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])); + let indices = HashSet::from([0]); + + let batch = flatten_struct_cols(&[Arc::clone(&structs)], &schema, &indices).unwrap(); + let a = batch.column(0).as_primitive::(); + let b = batch.column(1).as_primitive::(); + assert_eq!(a.iter().collect::>(), [Some(1), None, None, Some(4)]); + assert_eq!( + b.iter().collect::>(), + [Some(10), None, Some(30), Some(40)] + ); + + // A slice keeps the struct's mask aligned with its fields. + let batch = flatten_struct_cols(&[structs.slice(1, 2)], &schema, &indices).unwrap(); + let a = batch.column(0).as_primitive::(); + let b = batch.column(1).as_primitive::(); + assert_eq!(a.iter().collect::>(), [None, None]); + assert_eq!(b.iter().collect::>(), [None, Some(30)]); + } } diff --git a/native/proto/src/proto/operator.proto b/native/proto/src/proto/operator.proto index 28d94e78bf3..dc177beff91 100644 --- a/native/proto/src/proto/operator.proto +++ b/native/proto/src/proto/operator.proto @@ -1121,6 +1121,13 @@ message Explode { repeated spark.spark_expression.Expr project_list = 3; // Whether to emit a position column alongside the exploded values (posexplode) bool position = 4; + // Whether the array holds structs whose fields become separate output columns (inline). A + // null struct element yields a row of nulls. + bool inline = 5; + // Evaluated once after the input is exhausted, against a one-row batch with no columns. Yields + // an array of the same type as `child` holding the rows a JVM generator's `terminate()` emits, + // which carry the last input row's projected columns. + spark.spark_expression.Expr terminate = 6; } message HashJoin { diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 824118f3efe..754fd806110 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -812,6 +812,17 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(true) + val COMET_EXEC_JVM_GENERATOR_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.exec.jvmGenerator.enabled") + .category(CATEGORY_EXEC) + .doc( + "Whether to run a generator that Comet has no native implementation of, such as a " + + "Hive UDTF or `json_tuple`, inside Comet's native generate operator, which calls the " + + "generator on the JVM for each batch. When disabled, such a generator falls back to " + + "Spark.") + .booleanConf + .createWithDefault(true) + val COMET_SHUFFLE_NATIVE_HASH_PARTITIONING_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.shuffle.native.partitioning.hash.enabled") .withAlternative("spark.comet.native.shuffle.partitioning.hash.enabled") diff --git a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala index 7244875e2cb..98fb495427a 100644 --- a/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala +++ b/spark/src/main/scala/org/apache/comet/serde/CometScalaUDF.scala @@ -116,6 +116,18 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { case _ => false } + /** + * `expr` serialized with Spark's closure serializer, which respects the task context + * ClassLoader so classes from user jars resolve, and is how the JVM side of a `JvmScalarUdf` + * reads it back. Throws when the tree holds something the serializer refuses. + */ + def closureSerialize(expr: Expression): Array[Byte] = { + val buffer = SparkEnv.get.closureSerializer.newInstance().serialize(expr) + val serialized = new Array[Byte](buffer.remaining()) + buffer.get(serialized) + serialized + } + /** * Bind `expr`, closure-serialize it, and emit a `JvmScalarUdf` proto routed through * [[CometScalaUDFCodegen]] so that native execution evaluates the expression inside the @@ -202,11 +214,7 @@ object CometScalaUDF extends CometExpressionSerde[ScalaUDF] { // kernels and state on the executor (see `DispatchOccurrence`). val bytes = try { - val serializer = SparkEnv.get.closureSerializer.newInstance() - val buffer = serializer.serialize(DispatchOccurrence.tag(boundExpr)) - val serialized = new Array[Byte](buffer.remaining()) - buffer.get(serialized) - serialized + closureSerialize(DispatchOccurrence.tag(boundExpr)) } catch { // `NonFatal` rather than `NotSerializableException`: Java serialization reports an // unserializable object graph as one of several exception types depending on where in diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometGeneratorUDF.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometGeneratorUDF.scala new file mode 100644 index 00000000000..8057c58ad2a --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometGeneratorUDF.scala @@ -0,0 +1,254 @@ +/* + * 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.spark.sql.comet + +import java.nio.ByteBuffer +import java.util.concurrent.atomic.AtomicLong + +import scala.collection.mutable +import scala.collection.mutable.ArrayBuffer +import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal + +import org.apache.arrow.vector.{BitVector, ValueVector, VarBinaryVector} +import org.apache.arrow.vector.types.pojo.Field +import org.apache.spark.SparkEnv +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, AttributeSeq, BindReferences, Expression, Generator, GenericInternalRow, Literal, Unevaluable, UnsafeProjection} +import org.apache.spark.sql.catalyst.util.GenericArrayData +import org.apache.spark.sql.comet.execution.arrow.ArrowWriter +import org.apache.spark.sql.execution.ExecSubqueryExpression +import org.apache.spark.sql.types.ArrayType +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} + +import org.apache.comet.CometConf +import org.apache.comet.codegen.CometBatchKernelCodegen +import org.apache.comet.serde.{CometInvokeTargets, ExprOuterClass} +import org.apache.comet.serde.CometScalaUDF.closureSerialize +import org.apache.comet.serde.QueryPlanSerde.{exprToProto, serializeDataType} +import org.apache.comet.udf.{CometUDF, NativeUdfCall} +import org.apache.comet.udf.codegen.CometScalaUDFCodegen +import org.apache.comet.vector.CometVector + +/** + * Runs a Spark [[Generator]] for Comet's native generate operator. For each input row it returns + * the rows the generator emits as one `array` element, which the native operator unnests + * and flattens the way it does for `inline`. An input row that emits nothing gets an empty array, + * so `outer` keeps it as one row of nulls. + * + * Arg 0 is the digest of the closure-serialized bound generator, arg 1 the serialized bytes, and + * arg 2 a boolean scalar that is true for the call the native operator makes once its input is + * exhausted. That call returns one element holding the rows of `Generator.terminate`. Args 3..N + * are the columns the bound generator reads, in ordinal order. + * + * A generator can keep state across rows and emit rows from `terminate`, as a Hive UDTF does, so + * every native plan gets its own instance: the key is the digest and the plan, and `releasePlan` + * drops the plan's instances. The serde ships the bound generator rather than evaluating it on + * the driver, so this works in cluster mode without driver-side state. `evaluate` is synchronized + * for the reason `CometScalaUDFCodegen.evaluate` is. + */ +class CometGeneratorUDF extends CometUDF { + import CometGeneratorUDF._ + + private val generators = mutable.HashMap.empty[(ByteBuffer, Long), BoundGenerator] + + override def evaluate(inputs: Array[ValueVector], numRows: Int): ValueVector = + evaluate(inputs, numRows, partitionIndex = 0, planId = NoPlan) + + override def evaluate( + inputs: Array[ValueVector], + numRows: Int, + partitionIndex: Int, + planId: Long): ValueVector = this.synchronized { + require( + inputs.length >= 3, + "CometGeneratorUDF requires the generator digest, its serialized bytes and the terminate " + + s"flag, got ${inputs.length} inputs") + val digest = ByteBuffer.wrap(binaryScalar(inputs(0))) + val bound = generators.getOrElseUpdate( + (digest, planId), + new BoundGenerator(deserialize(binaryScalar(inputs(1))))) + + val output = CometBatchKernelCodegen.allocateOutput(bound.outputField, numRows, 0) + try { + val writer = ArrowWriter.createFieldWriter(output) + val outputRow = new GenericInternalRow(1) + def writeRows(rows: Any): Unit = { + // A generator may reuse a row's buffers from one `next` to the next, which + // `GenericInternalRow.copy` shares rather than copies. + val copies = new ArrayBuffer[Any] + iteratorOf(rows).foreach(row => copies += bound.snapshot(row)) + outputRow.update(0, new GenericArrayData(copies.toArray)) + writer.write(outputRow, 0) + } + + if (inputs(2).asInstanceOf[BitVector].get(0) == 1) { + writeRows(bound.generator.terminate()) + } else { + val columns: Array[ColumnVector] = + inputs.drop(3).map(v => CometVector.getVector(v, null): ColumnVector) + val batch = new ColumnarBatch(columns, numRows) + var i = 0 + while (i < numRows) { + writeRows(bound.generator.eval(batch.getRow(i))) + i += 1 + } + } + writer.finish() + output + } catch { + case t: Throwable => + output.close() + throw t + } + } + + override def releasePlan(planId: Long): Unit = this.synchronized { + generators.keys.filter(_._2 == planId).toList.foreach(generators.remove) + } + + private def deserialize(bytes: Array[Byte]): Generator = { + val loader = Option(Thread.currentThread().getContextClassLoader) + .getOrElse(classOf[Generator].getClassLoader) + SparkEnv.get.closureSerializer + .newInstance() + .deserialize[Generator](ByteBuffer.wrap(bytes), loader) + } +} + +object CometGeneratorUDF { + + /** The plan id of calls from outside a native plan, such as unit tests. */ + private val NoPlan = -1L + + /** Numbers each generator the serde ships; see [[calls]]. */ + private val nextOccurrence = new AtomicLong() + + /** Why `generator` cannot run through [[CometGeneratorUDF]], or `None` when it can. */ + def declineReason(generator: Generator): Option[String] = { + val name = generator.nodeName + if (!CometConf.COMET_EXEC_JVM_GENERATOR_ENABLED.get()) { + return Some( + s"Unsupported generator: $name. Set ${CometConf.COMET_EXEC_JVM_GENERATOR_ENABLED.key}" + + "=true to run it on the JVM inside Comet") + } + // The generator is serialized at plan time, before any subquery in it has run, and an + // unevaluable or native-only node has nothing to run on the JVM. Attributes are unevaluable + // only until they are bound, which `calls` does. + generator + .find { + case _: Attribute => false + case _: ExecSubqueryExpression | _: Unevaluable | _: NativeUdfCall => true + case _ => false + } + .map(node => s"$name: ${node.nodeName} cannot run inside a JVM generator") + .orElse(CometInvokeTargets.declineReason(generator)) + .orElse(Some(s"$name: unsupported output type ${outputType(generator).catalogString}") + .filter(_ => serializeDataType(outputType(generator)).isEmpty)) + } + + /** + * The `JvmScalarUdf` calls that run `generator` through [[CometGeneratorUDF]]: one evaluated + * for each input batch, and the one the native operator makes once its input is exhausted. + * Fails with the reason when `generator` cannot be shipped. + */ + def calls( + generator: Generator, + inputs: Seq[Attribute]): Either[String, (ExprOuterClass.Expr, ExprOuterClass.Expr)] = { + val name = generator.nodeName + // Bind against only the attributes the generator reads, which become the data args. + val attrs = generator.collect { case a: AttributeReference => a }.distinct + val bound = BindReferences.bindReference(generator, AttributeSeq(attrs)) + val bytes = + try { + closureSerialize(bound) + } catch { + case NonFatal(e) => + return Left( + s"$name could not be closure-serialized (${e.getClass.getSimpleName}: " + + s"${e.getMessage})") + } + // The digest also covers a number unique to this generator, so two identical generators in + // one plan get their own instances and state. + val occurrence = ByteBuffer.allocate(8).putLong(nextOccurrence.getAndIncrement()).array() + val digest = CometScalaUDFCodegen.digest(bytes ++ occurrence) + val returnType = serializeDataType(outputType(generator)).getOrElse { + return Left(s"$name: unsupported output type ${outputType(generator).catalogString}") + } + + def call( + terminate: Boolean, + dataArgs: Seq[Expression]): Either[String, ExprOuterClass.Expr] = { + val args = (Seq(Literal(digest), Literal(bytes), Literal(terminate)) ++ dataArgs) + .map(exprToProto(_, inputs)) + if (args.exists(_.isEmpty)) { + Left(s"$name: an argument could not be passed to the JVM generator") + } else { + val udf = ExprOuterClass.JvmScalarUdf + .newBuilder() + .setClassName(classOf[CometGeneratorUDF].getName) + .addAllArgs(args.flatten.asJava) + .setReturnType(returnType) + .setReturnNullable(false) + Right(ExprOuterClass.Expr.newBuilder().setJvmScalarUdf(udf).build()) + } + } + for { + process <- call(terminate = false, attrs) + terminate <- call(terminate = true, Nil) + } yield (process, terminate) + } + + /** + * A generator ready to run in one native plan, and the Arrow field its output is exported as. + */ + private class BoundGenerator(val generator: Generator) { + val outputField: Field = CometGeneratorUDF.outputField(generator) + + private val projection = UnsafeProjection.create(generator.elementSchema) + + /** A copy of `row` that owns its bytes, so it survives the generator reusing its buffers. */ + def snapshot(row: InternalRow): InternalRow = projection(row).copy() + } + + /** The type of each element of the output: the rows one input row generates. */ + def outputType(generator: Generator): ArrayType = + ArrayType(generator.elementSchema, containsNull = false) + + private def outputField(generator: Generator): Field = + CometBatchKernelCodegen.toFfiArrowField("generated", outputType(generator), nullable = false) + + private def binaryScalar(v: ValueVector): Array[Byte] = { + val vec = v.asInstanceOf[VarBinaryVector] + require( + vec.getValueCount >= 1 && !vec.isNull(0), + "CometGeneratorUDF requires a binary scalar") + vec.get(0) + } + + /** `Generator.eval` returns a `TraversableOnce` on Scala 2.12 and an `IterableOnce` on 2.13. */ + private def iteratorOf(rows: Any): Iterator[InternalRow] = rows match { + case it: Iterator[_] => it.asInstanceOf[Iterator[InternalRow]] + case it: Iterable[_] => it.iterator.asInstanceOf[Iterator[InternalRow]] + case opt: Option[_] => opt.iterator.asInstanceOf[Iterator[InternalRow]] + case other => + throw new IllegalStateException(s"Unexpected generator output ${other.getClass.getName}") + } +} diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala index 53b4ae7888d..c0e865302f8 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/arrow/ArrowWriters.scala @@ -43,7 +43,7 @@ import org.apache.spark.unsafe.Platform * potential conflicts with Spark's Arrow dependencies, hence we cannot reuse Spark's ArrowWriter * directly. */ -private[arrow] object ArrowWriter { +private[comet] object ArrowWriter { def create(root: VectorSchemaRoot, fixedWidthCapacity: Int): ArrowWriter = { require(fixedWidthCapacity >= 0, "Fixed-width capacity must be non-negative") val children = root.getFieldVectors().asScala.map { vector => @@ -668,7 +668,7 @@ private[arrow] final class DictionaryCache { } } -private[arrow] abstract class ArrowFieldWriter { +private[comet] abstract class ArrowFieldWriter { def valueVector: ValueVector diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala index dd7c59f6556..4ed16af9a13 100644 --- a/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala +++ b/spark/src/main/scala/org/apache/spark/sql/comet/operators.scala @@ -19,8 +19,6 @@ package org.apache.spark.sql.comet -import java.util.Locale - import scala.annotation.tailrec import scala.collection.mutable import scala.collection.mutable.ArrayBuffer @@ -31,7 +29,7 @@ import org.apache.spark.broadcast.Broadcast import org.apache.spark.internal.Logging import org.apache.spark.rdd.RDD import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.expressions.{Ascending, Attribute, AttributeSeq, AttributeSet, CodegenObjectFactoryMode, Expression, ExpressionSet, Generator, LeafExpression, Literal, NamedExpression, SortOrder, XXH64} +import org.apache.spark.sql.catalyst.expressions.{ArrayRepeat, Ascending, Attribute, AttributeSeq, AttributeSet, Cast, CodegenObjectFactoryMode, CreateArray, CreateNamedStruct, ExplodeBase, Expression, ExpressionSet, Generator, GreaterThan, If, Inline, LeafExpression, Literal, NamedExpression, ReplicateRows, SortOrder, Stack, XXH64} import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, AggregateMode, Average, CollectList, CollectSet, Final, ImperativeAggregate, Mode, Partial, PartialMerge, Percentile, Sum} import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback import org.apache.spark.sql.catalyst.optimizer.{BuildLeft, BuildRight, BuildSide} @@ -62,7 +60,7 @@ import org.apache.comet.{CometConf, CometExecIterator, CometRuntimeException, Co import org.apache.comet.CometSparkSessionExtensions.{isCometShuffleEnabled, isSpark35Plus, withFallbackReason} import org.apache.comet.parquet.{CometParquetUtils, HadoopFileOutput} import org.apache.comet.rules.CometExecRule -import org.apache.comet.serde.{CometCollectBuffer, CometOperatorSerde, Compatible, OperatorOuterClass, QueryContextInterner, SupportLevel, Unsupported} +import org.apache.comet.serde.{CometCollectBuffer, CometOperatorSerde, Compatible, ExprOuterClass, OperatorOuterClass, QueryContextInterner, SupportLevel, Unsupported} import org.apache.comet.serde.OperatorOuterClass.{AggregateMode => CometAggregateMode, Operator} import org.apache.comet.serde.QueryPlanSerde import org.apache.comet.serde.QueryPlanSerde.{aggExprToProto, exprToProto, hasNonDefaultStringCollation, isStringCollationType, supportedSortType} @@ -1938,22 +1936,41 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { if (!op.generator.deterministic) { return Unsupported(Some("Only deterministic generators are supported")) } - if (op.generator.children.length != 1) { - return Unsupported(Some("generators with multiple inputs are not supported")) - } - val nodeName = op.generator.nodeName.toLowerCase(Locale.ROOT) - if (nodeName != "explode" && nodeName != "posexplode") { - return Unsupported(Some(s"Unsupported generator: ${op.generator.nodeName}")) - } - op.generator.children.head.dataType match { - case _: ArrayType => - Compatible() - case _: MapType => - // TODO add support for map types - // https://github.com/apache/datafusion-comet/issues/2837 - Unsupported(Some("Comet only supports explode/explode_outer for arrays, not maps")) + op.generator match { + // Spark only accepts an array of structs for `inline`, and `stack` is planned as `inline`; + // see `stackAsStructArray`. + case inline: Inline => + inline.child.dataType match { + case ArrayType(struct: StructType, _) if struct.isEmpty => + // An empty struct carries no columns, so the unnested rows have no data to hold their + // count and the batch cannot be built. + Unsupported(Some("inline over an array of structs with no fields is not supported")) + case _ => Compatible() + } + case stack: Stack => + // The rewrite builds one struct expression per output row, so keep large row counts on + // Spark, which stops generating code for them too. + if (stack.supportCodegen) { + Compatible() + } else { + Unsupported(Some("stack with more than 50 rows is not supported")) + } + case explode: ExplodeBase => + explode.child.dataType match { + case _: ArrayType => + Compatible() + case _: MapType => + // TODO add support for map types + // https://github.com/apache/datafusion-comet/issues/2837 + Unsupported(Some("Comet only supports explode/explode_outer for arrays, not maps")) + case other => + Unsupported(Some(s"Unsupported data type: $other")) + } case other => - Unsupported(Some(s"Unsupported data type: $other")) + CometGeneratorUDF.declineReason(other) match { + case Some(reason) => Unsupported(Some(reason)) + case None => Compatible() + } } } @@ -1961,17 +1978,34 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { op: GenerateExec, builder: Operator.Builder, childOp: OperatorOuterClass.Operator*): Option[OperatorOuterClass.Operator] = { - val childExpr = op.generator.children.head - val childExprProto = exprToProto(childExpr, op.child.output) + val explodeBuilder = OperatorOuterClass.Explode.newBuilder().setOuter(op.outer) + val childExprProto = op.generator match { + case inline: Inline => + explodeBuilder.setInline(true) + exprToProto(inline.child, op.child.output) + case stack: Stack => + explodeBuilder.setInline(true) + exprToProto(stackAsStructArray(stack), op.child.output) + case explode: ExplodeBase => + explodeBuilder.setPosition(explode.position) + exprToProto(explode.child, op.child.output) + case rows: ReplicateRows => + explodeBuilder.setInline(true) + // The rewrite is native end to end, so it is preferred to the JVM generator. + exprToProto(replicateRowsAsStructArray(rows), op.child.output) + .orElse(jvmGeneratorCall(op, explodeBuilder)) + case _ => + explodeBuilder.setInline(true) + jvmGeneratorCall(op, explodeBuilder) + } if (childExprProto.isEmpty) { return None } - // Convert the projection expressions (columns to carry forward) - // These are the non-generator output columns - val generatorOutput = op.generatorOutput.toSet - val projectExprs = op.output.filterNot(generatorOutput.contains).map { attr => + // The child columns carried forward. Not derived from `op.output` minus the generator output: + // a generator may reuse its child's attribute ids, as `ReplicateRows` does for EXCEPT ALL. + val projectExprs = op.requiredChildOutput.map { attr => exprToProto(attr, op.child.output) } @@ -1979,18 +2013,63 @@ object CometExplodeExec extends CometOperatorSerde[GenerateExec] { return None } - val isPosExplode = op.generator.nodeName.toLowerCase(Locale.ROOT) == "posexplode" - - val explodeBuilder = OperatorOuterClass.Explode - .newBuilder() + explodeBuilder .setChild(childExprProto.get) - .setOuter(op.outer) - .setPosition(isPosExplode) .addAllProjectList(projectExprs.map(_.get).asJava) Some(builder.setExplode(explodeBuilder).build()) } + /** + * The expression that runs `op.generator` on the JVM, which yields for each input row an array + * of the structs it emits, and sets the one that runs its `terminate` on `explodeBuilder`. + */ + private def jvmGeneratorCall( + op: GenerateExec, + explodeBuilder: OperatorOuterClass.Explode.Builder): Option[ExprOuterClass.Expr] = + CometGeneratorUDF.calls(op.generator, op.child.output) match { + case Right((process, terminate)) => + explodeBuilder.setTerminate(terminate) + Some(process) + case Left(reason) => + withFallbackReason(op, reason) + None + } + + /** + * `ReplicateRows(n, v1, ..., vk)` emits `n` copies of the row `(v1, ..., vk)`, which is + * `inline` over `array_repeat` of that row as a struct. A NULL or negative `n` gives no rows in + * both. The count is a `Long` and `array_repeat` takes an `Int`, so a count beyond + * `Int.MaxValue` is clamped rather than wrapped; no array that long fits in memory. + */ + private def replicateRowsAsStructArray(rows: ReplicateRows): Expression = { + val count = rows.children.head + val row = CreateNamedStruct(rows.elementSchema.fields.toSeq.zip(rows.children.tail).flatMap { + case (field, value) => Seq(Literal(field.name), value) + }) + val maxInt = Literal(Int.MaxValue.toLong) + ArrayRepeat( + row, + If(GreaterThan(count, maxInt), Literal(Int.MaxValue), Cast(count, IntegerType))) + } + + /** + * `stack(n, v1, ..., vk)` emits `n` rows of `ceil(k / n)` columns, filled row by row and padded + * with nulls. That is `inline` over an array holding one struct per output row. + */ + private def stackAsStructArray(stack: Stack): Expression = { + val numRows = stack.children.head.eval().asInstanceOf[Int] + val values = stack.children.tail + val fields = stack.elementSchema.fields.toSeq + CreateArray(Seq.tabulate(numRows) { row => + CreateNamedStruct(fields.zipWithIndex.flatMap { case (field, col) => + val index = row * fields.length + col + val value = if (index < values.length) values(index) else Literal(null, field.dataType) + Seq(Literal(field.name), value) + }) + }) + } + override def createExec(nativeOp: Operator, op: GenerateExec): CometNativeExec = { CometExplodeExec( nativeOp, diff --git a/spark/src/test/resources/sql-tests/expressions/array/inline.sql b/spark/src/test/resources/sql-tests/expressions/array/inline.sql new file mode 100644 index 00000000000..95d7c7b0951 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/array/inline.sql @@ -0,0 +1,120 @@ +-- 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. + +-- inline and inline_outer run natively: the planner unnests the array of +-- structs and flattens each struct into one column per field. A NULL struct +-- element yields a row of NULLs. inline_outer emits one row of NULLs for a +-- NULL or empty array, and plain inline drops both. + +statement +CREATE TABLE test_inline(id int, arr array>) USING parquet + +statement +INSERT INTO test_inline VALUES + (1, array(named_struct('a', 1, 'b', 'x'), named_struct('a', 2, 'b', 'y'))), + (2, array(named_struct('a', 3, 'b', NULL))), + (3, array(named_struct('a', NULL, 'b', NULL))), + (4, array(NULL, named_struct('a', 4, 'b', 'z'), NULL)), + (5, array()), + (6, NULL), + (NULL, array(named_struct('a', 5, 'b', 'w'))) + +query +SELECT id, inline(arr) FROM test_inline + +query +SELECT id, inline_outer(arr) FROM test_inline + +query +SELECT id, a, b FROM test_inline LATERAL VIEW inline(arr) t AS a, b + +query +SELECT id, a, b FROM test_inline LATERAL VIEW OUTER inline(arr) t AS a, b + +-- only the generated columns +query +SELECT inline(arr) FROM test_inline + +-- a filter and an aggregate above the generator +query +SELECT id, a FROM test_inline LATERAL VIEW inline(arr) t AS a, b WHERE a > 1 + +query +SELECT b, count(*), sum(a) FROM test_inline LATERAL VIEW OUTER inline(arr) t AS a, b GROUP BY b + +-- a struct field named like a carried column +statement +CREATE TABLE test_inline_dup(id int, arr array>) USING parquet + +statement +INSERT INTO test_inline_dup VALUES + (1, array(named_struct('id', 10, 'v', 'a'), named_struct('id', 11, 'v', 'b'))), + (2, array(named_struct('id', 20, 'v', NULL))) + +query +SELECT * FROM test_inline_dup LATERAL VIEW inline(arr) t AS id2, v + +-- single-field structs +statement +CREATE TABLE test_inline_single(id int, arr array>) USING parquet + +statement +INSERT INTO test_inline_single VALUES + (1, array(named_struct('v', 1L), NULL, named_struct('v', NULL))), + (2, array()) + +query +SELECT id, inline_outer(arr) FROM test_inline_single + +-- nested and wide field types +statement +CREATE TABLE test_inline_nested( + id int, + arr array, + m: map, + s: struct>>) USING parquet + +statement +INSERT INTO test_inline_nested VALUES + (1, array( + named_struct('d', 1.5D, 'dec', 12345.678BD, 'dt', date'2024-01-01', + 'ts', timestamp'2024-01-01 12:34:56', 'bin', X'CAFE', 'ints', array(1, NULL, 3), + 'm', map('k', 1), 's', named_struct('x', 1, 'y', 'one')), + NULL, + named_struct('d', double('NaN'), 'dec', NULL, 'dt', NULL, 'ts', NULL, 'bin', NULL, + 'ints', array(), 'm', NULL, 's', NULL))), + (2, NULL) + +query +SELECT id, inline_outer(arr) FROM test_inline_nested + +-- an array built in the query rather than read from the table +query +SELECT id, inline(array(named_struct('k', id, 'v', 'p'), named_struct('k', id * 10, 'v', 'q'))) FROM test_inline WHERE id IS NOT NULL + +query +SELECT id, inline(arrays_zip(array(1, 2, 3), array('a', 'b'))) FROM test_inline WHERE id = 1 + +-- NULL struct elements produced by an expression rather than read from Parquet +query +SELECT id, inline_outer(array(IF(id % 2 = 0, NULL, named_struct('a', id, 'b', concat('v', id))), named_struct('a', -1, 'b', NULL))) FROM test_inline diff --git a/spark/src/test/resources/sql-tests/expressions/json/json_tuple.sql b/spark/src/test/resources/sql-tests/expressions/json/json_tuple.sql new file mode 100644 index 00000000000..e78c3cc4353 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/json/json_tuple.sql @@ -0,0 +1,59 @@ +-- 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. + +-- json_tuple has no native implementation. It runs on the JVM inside the native generate operator, +-- one row per input row, with NULL for a field the document lacks and a row of NULLs for a NULL or +-- invalid document. + +statement +CREATE TABLE test_json_tuple(id int, js string) USING parquet + +statement +INSERT INTO test_json_tuple VALUES + (1, '{"a": 1, "b": "x", "c": true}'), + (2, '{"a": null, "b": "y"}'), + (3, '{"b": "multi word", "d": {"e": 5}}'), + (4, '{"a": [1, 2, 3], "b": 1.5}'), + (5, 'not json'), + (6, ''), + (7, NULL), + (8, '{"a": "é日本", "b": ""}'), + (NULL, '{"a": 9}') + +query +SELECT id, json_tuple(js, 'a', 'b') FROM test_json_tuple + +query +SELECT json_tuple(js, 'a', 'b', 'c', 'd') FROM test_json_tuple + +-- a field the documents lack, and a repeated field +query +SELECT id, json_tuple(js, 'missing', 'a', 'a') FROM test_json_tuple + +query +SELECT id, a, b FROM test_json_tuple LATERAL VIEW json_tuple(js, 'a', 'b') j AS a, b + +-- the outer form keeps the rows that yield no values +query +SELECT id, a FROM test_json_tuple LATERAL VIEW OUTER json_tuple(js, 'a') j AS a + +-- a filter and an aggregate above the generator +query +SELECT id, a FROM test_json_tuple LATERAL VIEW json_tuple(js, 'a') j AS a WHERE a IS NOT NULL + +query +SELECT b, count(*) FROM test_json_tuple LATERAL VIEW json_tuple(js, 'b') j AS b GROUP BY b diff --git a/spark/src/test/resources/sql-tests/expressions/misc/stack.sql b/spark/src/test/resources/sql-tests/expressions/misc/stack.sql new file mode 100644 index 00000000000..98a5b792212 --- /dev/null +++ b/spark/src/test/resources/sql-tests/expressions/misc/stack.sql @@ -0,0 +1,91 @@ +-- 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. + +-- stack(n, v1, ..., vk) emits n rows of ceil(k / n) columns for every input +-- row, filling them row by row and padding the last row with NULLs. Comet plans +-- it as inline over an array holding one struct per output row. + +statement +CREATE TABLE test_stack(id int, x int, y int, s1 string, s2 string) USING parquet + +statement +INSERT INTO test_stack VALUES + (1, 10, 20, 'a', 'b'), + (2, NULL, 21, NULL, 'c'), + (3, 12, NULL, 'd', NULL), + (NULL, NULL, NULL, NULL, NULL) + +-- unpivot two columns into (name, value) pairs +query +SELECT id, stack(2, 'x', x, 'y', y) AS (name, value) FROM test_stack + +query +SELECT id, name, value FROM test_stack LATERAL VIEW stack(2, 'x', x, 'y', y) t AS name, value + +-- uneven argument count: the last row is padded with NULLs +query +SELECT id, stack(3, x, y, id, x, y) FROM test_stack + +-- more rows than values +query +SELECT id, stack(3, x) FROM test_stack + +-- a single row +query +SELECT id, stack(1, x, s1, y, s2) FROM test_stack + +-- one column per row +query +SELECT id, stack(4, s1, s2, s1, s2) FROM test_stack + +-- a NULL literal takes the column's type +query +SELECT id, stack(2, x, NULL) FROM test_stack + +query +SELECT id, stack(2, NULL, s1, 'k', NULL) FROM test_stack + +-- a row count given as a foldable expression rather than a literal +query +SELECT id, stack(1 + 1, x, y) FROM test_stack + +-- expressions as arguments +query +SELECT id, stack(2, x + 1, upper(s1), y * 2, concat(s1, s2)) FROM test_stack + +-- a filter and an aggregate above the generator +query +SELECT name, count(value), sum(value) FROM test_stack LATERAL VIEW stack(2, 'x', x, 'y', y) t AS name, value WHERE id IS NOT NULL GROUP BY name + +-- only the generated columns +query +SELECT stack(2, x, y) FROM test_stack + +-- complex values +statement +CREATE TABLE test_stack_nested(id int, a1 array, a2 array, st1 struct, st2 struct) USING parquet + +statement +INSERT INTO test_stack_nested VALUES + (1, array(1, 2), array(3), named_struct('p', 1, 'q', 'one'), named_struct('p', 2, 'q', NULL)), + (2, NULL, array(), NULL, named_struct('p', NULL, 'q', 'two')) + +query +SELECT id, stack(2, a1, st1, a2, st2) FROM test_stack_nested + +query +SELECT id, stack(3, a1, a2) FROM test_stack_nested diff --git a/spark/src/test/resources/sql-tests/operators/except_all.sql b/spark/src/test/resources/sql-tests/operators/except_all.sql new file mode 100644 index 00000000000..4425b8c780e --- /dev/null +++ b/spark/src/test/resources/sql-tests/operators/except_all.sql @@ -0,0 +1,108 @@ +-- 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. + +-- EXCEPT ALL and INTERSECT ALL are planned with the internal ReplicateRows generator, which +-- repeats each group of rows as many times as the counts left after the set operation. It has no +-- native implementation, so it runs on the JVM inside the native generate operator and carries +-- every column of the row through. + +statement +CREATE TABLE test_set_left( + id int, + d decimal(20, 2), + bin binary, + f float, + dbl double, + dt date, + ts timestamp, + ts_ntz timestamp_ntz, + arr array, + st struct, + s string +) USING parquet + +statement +CREATE TABLE test_set_right LIKE test_set_left + +statement +INSERT INTO test_set_left VALUES + (1, 1.50, X'01', 1.5, 1.5, DATE '2024-01-01', TIMESTAMP '2024-01-01 00:00:00', TIMESTAMP_NTZ '2024-01-01 00:00:00', array(1, 2), named_struct('a', 1, 'b', 'x'), 'a'), + (1, 1.50, X'01', 1.5, 1.5, DATE '2024-01-01', TIMESTAMP '2024-01-01 00:00:00', TIMESTAMP_NTZ '2024-01-01 00:00:00', array(1, 2), named_struct('a', 1, 'b', 'x'), 'a'), + (1, 1.50, X'01', 1.5, 1.5, DATE '2024-01-01', TIMESTAMP '2024-01-01 00:00:00', TIMESTAMP_NTZ '2024-01-01 00:00:00', array(1, 2), named_struct('a', 1, 'b', 'x'), 'a'), + (2, -2.25, X'00FF', CAST('NaN' AS FLOAT), CAST('NaN' AS DOUBLE), DATE '9999-12-31', TIMESTAMP '9999-12-31 23:59:59', TIMESTAMP_NTZ '9999-12-31 23:59:59', array(NULL, 3), named_struct('a', NULL, 'b', 'y'), 'é日本'), + (2, -2.25, X'00FF', CAST('NaN' AS FLOAT), CAST('NaN' AS DOUBLE), DATE '9999-12-31', TIMESTAMP '9999-12-31 23:59:59', TIMESTAMP_NTZ '9999-12-31 23:59:59', array(NULL, 3), named_struct('a', NULL, 'b', 'y'), 'é日本'), + (3, 0.00, X'', float('-0.0'), double('-0.0'), DATE '1970-01-01', TIMESTAMP '1970-01-01 00:00:00', TIMESTAMP_NTZ '1970-01-01 00:00:00', array(), named_struct('a', 3, 'b', NULL), ''), + (4, 99999999999999999.99, X'AB', CAST('Infinity' AS FLOAT), CAST('-Infinity' AS DOUBLE), DATE '2000-02-29', TIMESTAMP '2000-02-29 12:34:56', TIMESTAMP_NTZ '2000-02-29 12:34:56', array(7), named_struct('a', 4, 'b', 'z'), 'w'), + (NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL), + (NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL) + +statement +INSERT INTO test_set_right VALUES + (1, 1.50, X'01', 1.5, 1.5, DATE '2024-01-01', TIMESTAMP '2024-01-01 00:00:00', TIMESTAMP_NTZ '2024-01-01 00:00:00', array(1, 2), named_struct('a', 1, 'b', 'x'), 'a'), + (2, -2.25, X'00FF', CAST('NaN' AS FLOAT), CAST('NaN' AS DOUBLE), DATE '9999-12-31', TIMESTAMP '9999-12-31 23:59:59', TIMESTAMP_NTZ '9999-12-31 23:59:59', array(NULL, 3), named_struct('a', NULL, 'b', 'y'), 'é日本'), + (2, -2.25, X'00FF', CAST('NaN' AS FLOAT), CAST('NaN' AS DOUBLE), DATE '9999-12-31', TIMESTAMP '9999-12-31 23:59:59', TIMESTAMP_NTZ '9999-12-31 23:59:59', array(NULL, 3), named_struct('a', NULL, 'b', 'y'), 'é日本'), + (2, -2.25, X'00FF', CAST('NaN' AS FLOAT), CAST('NaN' AS DOUBLE), DATE '9999-12-31', TIMESTAMP '9999-12-31 23:59:59', TIMESTAMP_NTZ '9999-12-31 23:59:59', array(NULL, 3), named_struct('a', NULL, 'b', 'y'), 'é日本'), + (3, 0.00, X'', float('0.0'), double('0.0'), DATE '1970-01-01', TIMESTAMP '1970-01-01 00:00:00', TIMESTAMP_NTZ '1970-01-01 00:00:00', array(), named_struct('a', 3, 'b', NULL), ''), + (NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL) + +-- every column at once +query +SELECT * FROM test_set_left EXCEPT ALL SELECT * FROM test_set_right + +query +SELECT * FROM test_set_left INTERSECT ALL SELECT * FROM test_set_right + +-- one type at a time, so a type the carried-column path mishandles shows up on its own +query +SELECT d FROM test_set_left EXCEPT ALL SELECT d FROM test_set_right + +query +SELECT bin FROM test_set_left EXCEPT ALL SELECT bin FROM test_set_right + +query +SELECT f FROM test_set_left EXCEPT ALL SELECT f FROM test_set_right + +query +SELECT dbl FROM test_set_left EXCEPT ALL SELECT dbl FROM test_set_right + +query +SELECT dt FROM test_set_left EXCEPT ALL SELECT dt FROM test_set_right + +query +SELECT ts FROM test_set_left EXCEPT ALL SELECT ts FROM test_set_right + +query +SELECT ts_ntz FROM test_set_left EXCEPT ALL SELECT ts_ntz FROM test_set_right + +query +SELECT arr FROM test_set_left EXCEPT ALL SELECT arr FROM test_set_right + +query +SELECT st FROM test_set_left EXCEPT ALL SELECT st FROM test_set_right + +query +SELECT s FROM test_set_left EXCEPT ALL SELECT s FROM test_set_right + +query +SELECT s FROM test_set_left INTERSECT ALL SELECT s FROM test_set_right + +-- an empty result and an empty right side +query +SELECT s FROM test_set_left EXCEPT ALL SELECT s FROM test_set_left + +query +SELECT s FROM test_set_left EXCEPT ALL SELECT s FROM test_set_right WHERE id > 100 diff --git a/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala index 1ccfacdffe5..3e886c07f18 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometGenerateExecSuite.scala @@ -20,10 +20,11 @@ package org.apache.comet.exec import org.apache.spark.sql.{CometTestBase, DataFrame, Row} -import org.apache.spark.sql.comet.CometExplodeExec +import org.apache.spark.sql.comet.{CometExplodeExec, CometTestGenerators} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec import org.apache.spark.sql.execution.exchange.ReusedExchangeExec import org.apache.spark.sql.functions.col +import org.apache.spark.sql.hive.{CometHiveUDTFs, CountingUDTF} import org.apache.spark.sql.internal.SQLConf import org.apache.comet.CometConf @@ -237,6 +238,48 @@ class CometGenerateExecSuite extends CometTestBase { } } + test("stack with an argument Comet cannot run natively falls back") { + withSQLConf( + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "false") { + spark.udf.register("plus_one", (i: Int) => i + 1) + val df = Seq((1, 10), (2, 20)) + .toDF("id", "v") + .selectExpr("id", "stack(2, v, plus_one(v))") + checkSparkAnswerAndFallbackReason( + df, + s"${CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key}=false") + } + } + + test("stack with more than 50 rows falls back") { + withSQLConf(CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val values = (1 to 51).map(i => s"v + $i").mkString(", ") + val df = Seq((1, 10), (2, 20)).toDF("id", "v").selectExpr("id", s"stack(51, $values)") + checkSparkAnswerAndFallbackReason(df, "stack with more than 50 rows is not supported") + } + } + + test("stack with 50 rows runs natively") { + withSQLConf(CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val values = (1 to 50).map(i => s"v + $i").mkString(", ") + val df = Seq((1, 10), (2, 20)).toDF("id", "v").selectExpr("id", s"stack(50, $values)") + checkSparkAnswerAndOperator(df) + } + } + + test("inline over an array of structs with no fields falls back") { + withSQLConf(CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + Seq("inline", "inline_outer").foreach { fn => + val df = + Seq((1, 10), (2, 20)).toDF("id", "v").selectExpr(s"$fn(array(struct(), struct()))") + checkSparkAnswerAndFallbackReason( + df, + "inline over an array of structs with no fields is not supported") + } + } + } + test("explode with nullable projected column") { withSQLConf( CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true", @@ -722,4 +765,146 @@ class CometGenerateExecSuite extends CometTestBase { } } + /** Run `f` over a single-file Parquet table `t` of `(id, s, js)` rows. */ + private def withGeneratorTable(f: => Unit): Unit = { + val rows = Seq( + (1, "a", """{"a": 1, "b": "x"}"""), + (2, null, """{"a": null}"""), + (3, "c", "not json"), + (4, "d", null)) + withTempPath { path => + rows.toDF("id", "s", "js").repartition(1).write.parquet(path.getAbsolutePath) + withParquetTable(path.getAbsolutePath, "t")(f) + } + } + + /** Check `df` against Spark, and that its generator ran in Comet's native operator. */ + private def checkSparkAnswerWithNativeGenerate(df: DataFrame): Unit = { + val (_, cometPlan) = checkSparkAnswer(df) + assert( + collectFirst(cometPlan) { case generate: CometExplodeExec => generate }.isDefined, + s"Expected a native generator in\n$cometPlan") + } + + test("Hive UDTF runs natively and emits its terminate rows") { + CometHiveUDTFs.register(spark, "counting_udtf", classOf[CountingUDTF]) + // Two rows per batch, so the running count and the `close` total cross a batch boundary. + withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "2") { + withGeneratorTable { + // The row `close` emits carries the last input row's `id`, as in Spark. + checkSparkAnswerAndOperator(sql("SELECT id, counting_udtf(s) FROM t")) + checkSparkAnswerAndOperator( + sql("SELECT id, value, n FROM t LATERAL VIEW OUTER counting_udtf(s) u AS value, n")) + // With no input rows only the `close` row remains. + checkSparkAnswerAndOperator(sql("SELECT counting_udtf(s) FROM t WHERE id > 10")) + // Two identical UDTFs in one plan each keep their own count. + checkSparkAnswerAndOperator(sql(""" + |SELECT id, a.value, a.n, b.value, b.n FROM t + |LATERAL VIEW counting_udtf(s) a AS value, n + |LATERAL VIEW counting_udtf(s) b AS value, n""".stripMargin)) + } + } + } + + test("Hive UDTF keeps separate state per partition") { + CometHiveUDTFs.register(spark, "counting_udtf", classOf[CountingUDTF]) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.SHUFFLE_PARTITIONS.key -> "3", + CometConf.COMET_BATCH_SIZE.key -> "1") { + withGeneratorTable { + // Each partition ends with one `close` row holding its own count, so the counts add up + // to the input rows whichever partition each row lands in. + checkSparkAnswerWithNativeGenerate(sql(""" + |SELECT count(*), sum(n) FROM ( + | SELECT counting_udtf(s) FROM (SELECT * FROM t DISTRIBUTE BY id)) + |WHERE value = 'total'""".stripMargin)) + } + } + } + + test("generator rows that share a reused buffer are kept apart") { + CometTestGenerators.registerReusedBuffer(spark, "reused_buffer") + withGeneratorTable { + checkSparkAnswerWithNativeGenerate( + sql("SELECT id, reused_buffer(encode(s, 'UTF-8')) FROM t")) + checkSparkAnswerWithNativeGenerate( + sql("SELECT id, b FROM t LATERAL VIEW OUTER reused_buffer(encode(s, 'UTF-8')) r AS b")) + } + } + + test("terminate with no rows over a non-nullable carried column adds nothing") { + // `k` is never null, so the column that stands in for a last input row cannot hold the nulls + // a partition with no input would give it, and `json_tuple` has no `terminate` rows to emit. + withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + withGeneratorTable { + val df = sql(""" + |SELECT k, json_tuple(js, 'a') FROM ( + | SELECT 1 AS k, js FROM t WHERE id > 10)""".stripMargin) + checkSparkAnswerWithNativeGenerate(df) + assert(df.collect().isEmpty) + } + } + } + + test("json_tuple runs natively") { + withGeneratorTable { + checkSparkAnswerAndOperator(sql("SELECT id, json_tuple(js, 'a', 'b') FROM t")) + checkSparkAnswerAndOperator( + sql("SELECT id, a FROM t LATERAL VIEW OUTER json_tuple(js, 'a') j AS a WHERE a IS NULL")) + } + } + + test("EXCEPT ALL replicates rows with array_repeat, not a JVM generator") { + withGeneratorTable { + val (_, cometPlan) = + checkSparkAnswer(sql("SELECT s FROM t UNION ALL SELECT s FROM t EXCEPT ALL SELECT 'a'")) + val generate = collectFirst(cometPlan) { case generate: CometExplodeExec => generate } + assert(generate.isDefined, s"Expected a native generator in\n$cometPlan") + assert(!generate.get.nativeOp.getExplode.getChild.hasJvmScalarUdf) + } + } + + test("EXCEPT ALL repeats rows by a count of zero, one and many") { + withSQLConf(CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + withTempView("l", "r") { + Seq(("a", 1), ("b", 1), ("b", 1), ("b", 1), ("c", 1), (null, 1), (null, 1)) + .toDF("k", "v") + .createOrReplaceTempView("l") + Seq(("a", 1), ("b", 1), ("c", 1), ("c", 1), (null, 1)) + .toDF("k", "v") + .createOrReplaceTempView("r") + checkSparkAnswerWithNativeGenerate( + sql("SELECT k, v FROM l EXCEPT ALL SELECT k, v FROM r")) + checkSparkAnswerWithNativeGenerate( + sql("SELECT k, v FROM l INTERSECT ALL SELECT k, v FROM r")) + } + } + } + + test("a generator without a native implementation falls back when the JVM path is off") { + withSQLConf(CometConf.COMET_EXEC_JVM_GENERATOR_ENABLED.key -> "false") { + withGeneratorTable { + checkSparkAnswerAndFallbackReason( + sql("SELECT id, json_tuple(js, 'a') FROM t"), + "Unsupported generator: JsonTuple") + } + } + } + + test("EXCEPT ALL in a correlated subquery keeps the carried columns") { + // The decorrelated plan carries t0a and t0b past `ReplicateRows`, whose output reuses + // those attribute ids. + withTempView("t0", "t1", "t2") { + sql("CREATE TEMP VIEW t0(t0a, t0b) AS VALUES (1, 1), (2, 0)") + sql("CREATE TEMP VIEW t1(t1a, t1c) AS VALUES (1, 3)") + sql("CREATE TEMP VIEW t2(t2b, t2c) AS VALUES (1, 5), (2, 7)") + checkSparkAnswerWithNativeGenerate(sql(""" + |SELECT t0a, (SELECT count(t1c) FROM + | (SELECT t1c FROM t1 WHERE t1a = t0a + | EXCEPT ALL + | SELECT t2c FROM t2 WHERE t2b = t0b)) + |FROM t0""".stripMargin)) + } + } } diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/CometTestGenerators.scala b/spark/src/test/scala/org/apache/spark/sql/comet/CometTestGenerators.scala new file mode 100644 index 00000000000..81eb52909a4 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/comet/CometTestGenerators.scala @@ -0,0 +1,68 @@ +/* + * 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.spark.sql.comet + +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.FunctionIdentifier +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionInfo, Generator, GenericInternalRow} +import org.apache.spark.sql.catalyst.expressions.codegen.CodegenFallback +import org.apache.spark.sql.types.{BinaryType, StructField, StructType} + +/** Catalyst generators for tests, registered as SQL functions. */ +object CometTestGenerators { + def register(spark: SparkSession, name: String, build: Seq[Expression] => Generator): Unit = + spark.sessionState.functionRegistry.registerFunction( + FunctionIdentifier(name), + new ExpressionInfo(name, name), + (children: Seq[Expression]) => build(children)) + + def registerReusedBuffer(spark: SparkSession, name: String): Unit = + register(spark, name, ReusedBufferGenerator(_)) +} + +/** + * Emits the bytes of its argument three times, `x1`, `x2` and `x3`, through one row and one byte + * array it overwrites between rows, as a generator that avoids allocating a row per output does. + * Only a copy of the bytes made before the iterator advances keeps the three rows apart. + */ +case class ReusedBufferGenerator(children: Seq[Expression]) + extends Generator + with CodegenFallback { + + override def elementSchema: StructType = StructType(Seq(StructField("b", BinaryType))) + + override def eval(input: InternalRow): Iterator[InternalRow] = { + val prefix = children.head.eval(input) match { + case null => Array.empty[Byte] + case bytes: Array[Byte] => bytes + } + val buffer = new Array[Byte](prefix.length + 1) + prefix.copyToArray(buffer) + val row = new GenericInternalRow(Array[Any](buffer)) + Iterator.tabulate(3) { i => + buffer(buffer.length - 1) = ('1' + i).toByte + row + } + } + + override protected def withNewChildrenInternal( + newChildren: IndexedSeq[Expression]): ReusedBufferGenerator = copy(children = newChildren) +} diff --git a/spark/src/test/scala/org/apache/spark/sql/hive/CometHiveUDTFs.scala b/spark/src/test/scala/org/apache/spark/sql/hive/CometHiveUDTFs.scala new file mode 100644 index 00000000000..8c9b7684230 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/hive/CometHiveUDTFs.scala @@ -0,0 +1,63 @@ +/* + * 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.spark.sql.hive + +import java.util.Arrays + +import org.apache.hadoop.hive.ql.udf.generic.GenericUDTF +import org.apache.hadoop.hive.serde2.objectinspector.{ObjectInspector, ObjectInspectorFactory, StructObjectInspector} +import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.FunctionIdentifier +import org.apache.spark.sql.catalyst.expressions.{Expression, ExpressionInfo} +import org.apache.spark.sql.hive.HiveShim.HiveFunctionWrapper + +/** Registers Hive UDTFs in a session that has no Hive support, as `HiveSessionCatalog` would. */ +object CometHiveUDTFs { + def register(spark: SparkSession, name: String, udtfClass: Class[_ <: GenericUDTF]): Unit = + spark.sessionState.functionRegistry.registerFunction( + FunctionIdentifier(name), + new ExpressionInfo(udtfClass.getName, name), + (children: Seq[Expression]) => + HiveGenericUDTF(name, new HiveFunctionWrapper(udtfClass.getName), children)) +} + +/** + * Emits each input value with a running count of the rows it has processed, and from `close` one + * more row holding the total. The count is per instance, so it shows whether the rows of a + * partition reach one instance, and the `close` row whether `terminate` runs. + */ +class CountingUDTF extends GenericUDTF { + private var count = 0L + + override def initialize(argOIs: StructObjectInspector): StructObjectInspector = + ObjectInspectorFactory.getStandardStructObjectInspector( + Arrays.asList("value", "n"), + Arrays.asList[ObjectInspector]( + PrimitiveObjectInspectorFactory.javaStringObjectInspector, + PrimitiveObjectInspectorFactory.javaLongObjectInspector)) + + override def process(args: Array[AnyRef]): Unit = { + count += 1 + forward(Array[AnyRef](Option(args(0)).map(_.toString).orNull, Long.box(count))) + } + + override def close(): Unit = forward(Array[AnyRef]("total", Long.box(count))) +}