diff --git a/changelog.d/746.added.md b/changelog.d/746.added.md new file mode 100644 index 000000000..05d2d20d6 --- /dev/null +++ b/changelog.d/746.added.md @@ -0,0 +1 @@ +The engine can plan and run grouped, cogrouped and map-style Python UDFs and Python aggregate UDFs (the `mapInPandas`, `applyInPandas`, `cogroup` and grouped-aggregate kinds) over a UDF runtime; plans that contain them carry the new plan-codec feature `python_udf_relations`. The runtimes that execute them follow in later releases. diff --git a/crates/krishiv-common/krishiv-error-classes.json b/crates/krishiv-common/krishiv-error-classes.json index 7be19cbd0..eadc63875 100644 --- a/crates/krishiv-common/krishiv-error-classes.json +++ b/crates/krishiv-common/krishiv-error-classes.json @@ -130,6 +130,12 @@ ], "sqlState" : "53400" }, + "KRISHIV_PYTHON_UDF_RESULT_MISMATCH" : { + "message" : [ + "The Python UDF returned a result that does not match what it declares: . Make the function return the declared columns, in their declared types, and for a grouped aggregate one value per group." + ], + "sqlState" : "39000" + }, "KRISHIV_PYTHON_WORKER_EXITED" : { "message" : [ "Python worker exited unexpectedly (crashed): . The next query starts a new worker. If a UDF ran out of memory, raise \"krishiv.python.worker.memoryLimit\"." diff --git a/crates/krishiv-exec/CLAUDE.md b/crates/krishiv-exec/CLAUDE.md index 35be37ff5..80abdafff 100644 --- a/crates/krishiv-exec/CLAUDE.md +++ b/crates/krishiv-exec/CLAUDE.md @@ -130,6 +130,7 @@ Paths are relative to this crate. - `src/window/mod.rs` window operator and plan types; `src/window/eval.rs` frames, ranking and frame aggregates; `src/planner/window.rs` plans it over sorted input - `src/operators/file_scan.rs` and `file_write.rs` file scan and write operators over `krishiv-io` - `src/planner/write.rs` (P3-T023a) `WriteMode` (append, truncate, overwrite by filter, dynamic overwrite), `WriteRequest` (mode, table name, input schema, options, write id) and `TableCapabilities` (Spark's `TableCapability` write values, checked as Spark v4.0.0's `TableCapabilityCheck`: every mode needs `BATCH_WRITE` or `V1_BATCH_WRITE`, a truncate also `TRUNCATE` or `OVERWRITE_BY_FILTER`, dynamic overwrite `BATCH_WRITE` and `OVERWRITE_DYNAMIC`); `TableProvider::capabilities` / `writer(&WriteRequest)` default to no write modes, and `TableCapabilities::check` raises `UNSUPPORTED_FEATURE.TABLE_OPERATION` ("... in batch mode") before anything is written +- `src/python_relation.rs` (design P3-T085 §3.2) the Python relation seam: `RelationCallSpec`, `PreparedRelationUdf::open`, `UdfTask` fed by `Frame`s (`UdfRuntime::prepare_relation` defaults to `KRISHIV_UDF_NOT_AVAILABLE`); `python_relation/sink.rs` `RelationSink` (map, grouped, cogrouped, aggregate shapes: one task per driver, groups cut where the sorted key changes, frames lead with the keys, a one-sided cogroup key gets a zero-row frame for the other side, the window checked per frame with a reserved backlog, `Blocked` while a lane's frame window is full or cuts wait, its future framing the backlog so no input is taken before it drained; a frame's computed columns reserved, results in the shared plan's list as they arrive), `python_relation/lane.rs` (a lane's frames, its reserved results and the pump; a refused result stops the task's reads), `python_relation/pair.rs` (an aggregate's pending keys joined with the values as they arrive); `planner/python_relation.rs` plans `MapPartitions`/`GroupMap`/`CoGroupMap`/Python-UDAF `Aggregate` onto `PhysNode::PythonMap`/`PythonGroupMap`/`PythonCoGroupMap`/`PythonAggregate` (hash exchange + spillable sort on the keys; a cogroup is one union stream sorted by key then side), refuses barrier mode, Python window aggregates and UDTFs (PR-26c); `physical/instantiate/python_relation.rs` builds them - `src/planner/ndv_sample.rs` (P2-T164 U6) `FileTable::sample_distinct`: plan-time distinct-value samples of the file columns an equality or `IN` compares, and of the column group keys while `krishiv.optimizer.groupKeyCode` is on (the session picks them, `catalog/ndv_sampling.rs`) (first rows in file order through `ScanStream`, `krishiv.optimizer.ndvSampling.*`), exact for a table the sample read whole, else GEE (`estimate`); cached in `krishiv_io::DistinctSamples`; `FileTable::statistics` reports them as `StatSource::Sampled`; `FileTable::sample_outcomes` says per table what each sample did (`SampleOutcome`: sampled, cached, timed out, no time left), which `EXPLAIN COST` prints as "Distinct-Value Samples" and `EXPLAIN VERBOSE` as a `Statistics` step ## Extension points diff --git a/crates/krishiv-exec/src/lib.rs b/crates/krishiv-exec/src/lib.rs index b44f91563..808007941 100644 --- a/crates/krishiv-exec/src/lib.rs +++ b/crates/krishiv-exec/src/lib.rs @@ -97,6 +97,7 @@ mod pipeline; pub mod planner; pub mod profile; pub mod python_eval; +pub mod python_relation; // Row layout (P2-T157): the join's key index (P2-T158) and the aggregate's // row tables (P2-T159); the key-page prototypes inside stay behind the // `row-table-prototype` feature. diff --git a/crates/krishiv-exec/src/physical/data.rs b/crates/krishiv-exec/src/physical/data.rs index 18ff66019..f118ed808 100644 --- a/crates/krishiv-exec/src/physical/data.rs +++ b/crates/krishiv-exec/src/physical/data.rs @@ -123,6 +123,7 @@ value_types!( krishiv_plan::Scan, krishiv_plan::AttrId, krishiv_plan::GeneratorKind, + krishiv_plan::PythonUdfRef, // Scan and runtime-filter descriptions. krishiv_io::ColumnPath, krishiv_io::FilterId, diff --git a/crates/krishiv-exec/src/physical/instantiate.rs b/crates/krishiv-exec/src/physical/instantiate.rs index 23dfca902..40fcaeb8e 100644 --- a/crates/krishiv-exec/src/physical/instantiate.rs +++ b/crates/krishiv-exec/src/physical/instantiate.rs @@ -89,6 +89,7 @@ mod detached; mod distributed; mod order; mod python_eval; +mod python_relation; pub(crate) use detached::{Detached, LoopRun}; mod join; use distributed::planned_schema_holds; @@ -927,6 +928,10 @@ impl Instantiation<'_> { fn untagged(&mut self, n: &PhysNode) -> Result { match n { PhysNode::PythonEval { .. } => self.python_eval(n), + PhysNode::PythonMap { .. } + | PhysNode::PythonGroupMap { .. } + | PhysNode::PythonCoGroupMap { .. } + | PhysNode::PythonAggregate { .. } => self.python_relation(n), PhysNode::Scan(s) => self.scan(s), PhysNode::Range(r) => { let source = Arc::new(crate::operators::Range::new( diff --git a/crates/krishiv-exec/src/physical/instantiate/python_relation.rs b/crates/krishiv-exec/src/physical/instantiate/python_relation.rs new file mode 100644 index 000000000..480c3c021 --- /dev/null +++ b/crates/krishiv-exec/src/physical/instantiate/python_relation.rs @@ -0,0 +1,300 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Instantiation of the Python relation nodes (design P3-T085 §3.1-§3.2): +//! each node's UDFs prepared by the worker's [`UdfRuntime`] as relation call +//! sites, and a [`RelationSink`] breaker over the input. +//! +//! [`UdfRuntime`]: crate::python_eval::UdfRuntime + +use super::*; +use arrow_schema::{DataType as ArrowType, Field}; +use krishiv_common::config::ARROW_MAX_RECORDS_PER_BATCH; +use krishiv_common::error::{KrishivError, KrishivUdfNotAvailable}; +use krishiv_plan::{ExprNode, PythonUdfRef, ResolvedExpr}; +use krishiv_types::arrow::client_arrow_type; + +use crate::operators::ScalarExpr; +use crate::physical::{PythonCallPlan, PythonRelationPlan}; +use crate::python_relation::{ + DEFAULT_WINDOW, InputColumn, PreparedRelationUdf, RelationCallSpec, RelationLane, + RelationLimits, RelationShape, RelationSink, RunnerConf, +}; + +/// Spark's `spark.sql.execution.arrow.maxBytesPerBatch`, read as the +/// session wrote it (Krishiv registers no key for it yet): no limit unless +/// set to other than Spark's default. +const MAX_BYTES_PER_BATCH: &str = "spark.sql.execution.arrow.maxBytesPerBatch"; + +impl Instantiation<'_> { + /// A Python relation node: a breaker whose sink feeds the node's tasks + /// and whose source yields what they returned. + pub(super) fn python_relation(&mut self, node: &PhysNode) -> Result { + let (input, shape) = match node { + PhysNode::PythonMap { input, .. } => (input, RelationShape::Map), + PhysNode::PythonGroupMap { input, keys, .. } => { + (input, RelationShape::Grouped { keys: keys.clone() }) + } + PhysNode::PythonCoGroupMap { + input, keys, side, .. + } => ( + input, + RelationShape::CoGrouped { + keys: keys.clone(), + side: *side, + }, + ), + PhysNode::PythonAggregate { + input, + keys, + key_values, + .. + } => ( + input, + RelationShape::Aggregate { + keys: keys.clone(), + key_values: key_values.clone(), + }, + ), + _ => return Err(internal("not a Python relation node".to_owned())), + }; + let input = self.node(input)?; + let conf = self.runner_conf(); + let (lanes, schema) = match node { + PhysNode::PythonMap { udf, .. } + | PhysNode::PythonGroupMap { udf, .. } + | PhysNode::PythonCoGroupMap { udf, .. } => { + let keys = match &shape { + RelationShape::Grouped { keys } | RelationShape::CoGrouped { keys, .. } => { + keys.len() + } + RelationShape::Map | RelationShape::Aggregate { .. } => 0, + }; + self.relation_lane(&input.schema, udf, keys, &conf)? + } + PhysNode::PythonAggregate { + layout, + key_values, + calls, + .. + } => self.aggregate_lanes(&input.schema, layout, key_values, calls, &conf)?, + _ => return Err(internal("not a Python relation node".to_owned())), + }; + let sink = RelationSink::try_new( + Arc::clone(&input.schema), + shape, + lanes, + schema, + RelationLimits { + max_rows: conf.max_records_per_batch, + max_bytes: conf.max_bytes_per_batch, + window: DEFAULT_WINDOW, + }, + self.spill.clone(), + )?; + let source = SharedSource::new(sink.schema()); + Ok(Built { + schema: sink.schema(), + node: Arc::new(BreakerNode { + input: input.node, + sink, + source, + ordered: false, + }), + }) + } + + /// The runner configuration of this worker's session. + fn runner_conf(&self) -> RunnerConf { + let config = self.worker.config.as_deref(); + let max_rows = config + .and_then(|c| c.get(&ARROW_MAX_RECORDS_PER_BATCH)) + .and_then(krishiv_common::config::ConfigValue::as_int) + .unwrap_or(10_000); + let max_bytes = config + .and_then(|c| c.get_unregistered(MAX_BYTES_PER_BATCH)) + .and_then(|v| v.trim().parse::().ok()) + // Spark's default, `Int.MaxValue`, is no limit + // (PythonArrowInput.scala:227 at v4.2.0). + .filter(|&v| v != i64::from(i32::MAX)) + .unwrap_or(0); + let mut settings = std::collections::BTreeMap::new(); + if let Some(c) = config { + for (k, v) in c.explicit_registered() { + settings.insert(k.to_owned(), v); + } + for (k, v) in c.explicit_unregistered() { + settings.insert(k.to_owned(), v.to_owned()); + } + } + RunnerConf { + max_records_per_batch: usize::try_from(max_rows).unwrap_or(0), + max_bytes_per_batch: usize::try_from(max_bytes).unwrap_or(0), + settings, + } + } + + /// The prepared relation call site of `spec`. + fn prepare_relation( + &self, + udf: &PythonUdfRef, + spec: &RelationCallSpec, + ) -> Result> { + let runtime = self.worker.python.as_ref().ok_or_else(|| { + KrishivError::from(KrishivUdfNotAvailable { + name: udf.name.clone(), + version: udf.version.to_string(), + }) + })?; + runtime.prepare_relation(spec) + } + + /// The lane of a map, grouped or cogrouped node, and the node's schema. + fn relation_lane( + &self, + input: &SchemaRef, + plan: &PythonRelationPlan, + key_count: usize, + conf: &RunnerConf, + ) -> Result<(Vec, SchemaRef)> { + let canonical: Vec = plan + .output + .iter() + .map(|c| Field::new(&c.name, c.data_type.to_arrow(), c.nullable)) + .collect(); + let client: Vec = canonical + .iter() + .map(|f| client_arrow_type(f.data_type())) + .collect(); + let client_fields: Vec = canonical + .iter() + .zip(&client) + .map(|(f, t)| Field::new(f.name(), t.clone(), f.is_nullable())) + .collect(); + let mut sides = Vec::with_capacity(plan.sides.len()); + for side in &plan.sides { + let columns = side + .columns + .iter() + .map(|&c| { + let field = input + .fields() + .get(c) + .ok_or_else(|| internal(format!("frame column {c} out of range")))?; + Ok(Arc::new(InputColumn { + index: c, + data_type: field.data_type().clone(), + nullable: field.is_nullable(), + }) as Arc) + }) + .collect::>>()?; + sides.push(columns); + } + let spec = RelationCallSpec { + name: plan.udf.name.clone(), + version: plan.udf.version, + eval_type: plan.udf.eval_type, + thread_safe: plan.udf.thread_safe, + input_types: plan.sides.iter().map(|s| s.types.clone()).collect(), + input_names: plan.sides.iter().map(|s| s.names.clone()).collect(), + key_count, + arity: plan + .sides + .first() + .map_or(0, |s| s.columns.len().saturating_sub(key_count)), + output_names: plan.output.iter().map(|c| c.name.clone()).collect(), + output_types: plan.output.iter().map(|c| c.data_type.clone()).collect(), + output_arrow: Arc::new(Schema::new(client_fields)), + conf: conf.clone(), + }; + let udf = self.prepare_relation(&plan.udf, &spec)?; + let lane = RelationLane { + name: plan.udf.name.clone(), + udf, + sides, + names: plan.sides.iter().map(|s| s.names.clone()).collect(), + client, + }; + Ok((vec![lane], Arc::new(Schema::new(canonical)))) + } + + /// One lane per call of a Python aggregate, and the node's schema: the + /// key columns, then one column per call. + fn aggregate_lanes( + &self, + input: &SchemaRef, + layout: &[AttrId], + keys: &[usize], + calls: &[PythonCallPlan], + conf: &RunnerConf, + ) -> Result<(Vec, SchemaRef)> { + let values = Values { + env: &self.env, + subqueries: &[], + }; + let mut fields: Vec = keys + .iter() + .map(|&k| { + input + .fields() + .get(k) + .map(|f| (**f).clone()) + .ok_or_else(|| internal(format!("key column {k} out of range"))) + }) + .collect::>()?; + let mut lanes = Vec::with_capacity(calls.len()); + for (i, c) in calls.iter().enumerate() { + let ExprNode::Resolved(ResolvedExpr::Function(f)) = &c.call.node else { + return Err(internal( + "a Python aggregate call is not a function".to_owned(), + )); + }; + let Some(udf) = PythonUdfRef::of(&f.id, f.deterministic) else { + return Err(internal( + "a Python aggregate call is not a Python UDF".to_owned(), + )); + }; + let args: Vec<&Expr> = f.args.iter().map(AsRef::as_ref).collect(); + let args = compile_all(&args, layout, input, &values)?; + let arg_types = f + .args + .iter() + .map(|a| { + a.resolved_type().ok_or_else(|| { + internal("a Python UDF argument has no resolved type".to_owned()) + }) + }) + .collect::>>()?; + let canonical = f.data_type.to_arrow(); + let client = client_arrow_type(&canonical); + let name = format!("pythonUDAF{i}"); + let arg_names: Vec = (0..args.len()).map(|j| format!("_{j}")).collect(); + let spec = RelationCallSpec { + name: udf.name.clone(), + version: udf.version, + eval_type: udf.eval_type, + thread_safe: udf.thread_safe, + arity: arg_types.len(), + input_types: vec![arg_types], + input_names: vec![arg_names.clone()], + key_count: 0, + output_names: vec![name.clone()], + output_types: vec![f.data_type.clone()], + output_arrow: Arc::new(Schema::new(vec![Field::new(&name, client.clone(), true)])), + conf: conf.clone(), + }; + let prepared = self.prepare_relation(&udf, &spec)?; + fields.push(Field::new(&name, canonical, true)); + lanes.push(RelationLane { + name: udf.name.clone(), + udf: prepared, + names: vec![arg_names], + sides: vec![ + args.into_iter() + .map(|e| Arc::new(Compiled(e)) as Arc) + .collect(), + ], + client: vec![client], + }); + } + Ok((lanes, Arc::new(Schema::new(fields)))) + } +} diff --git a/crates/krishiv-exec/src/physical/mod.rs b/crates/krishiv-exec/src/physical/mod.rs index 60f5f38f7..fe4eb1e2a 100644 --- a/crates/krishiv-exec/src/physical/mod.rs +++ b/crates/krishiv-exec/src/physical/mod.rs @@ -421,6 +421,54 @@ pub enum PhysNode { /// The calls, in output column order. calls: Vec, }, + /// `mapInPandas` / `mapInArrow` (design P3-T085 §3.1): each partition + /// of `input` sent to a Python UDF task; the rows are what it yields. + PythonMap { + /// The input. + input: Box, + /// The UDF and its declared result. + udf: PythonRelationPlan, + }, + /// `applyInPandas` / `applyInArrow` over groups: `input` is clustered by + /// the key columns and sorted by them (one partition without keys), so + /// a driver sees a key's rows together. + PythonGroupMap { + /// The input. + input: Box, + /// The key columns of `input`. + keys: Vec, + /// The UDF and its declared result. + udf: PythonRelationPlan, + }, + /// A cogrouped `applyInPandas` / `applyInArrow`: `input` is both + /// relations in one stream (each row's columns of the other relation + /// null), clustered by the key columns and sorted by them then `side`. + PythonCoGroupMap { + /// The input. + input: Box, + /// The key columns of `input`. + keys: Vec, + /// The column (`INT`, 0 left, 1 right) naming each row's relation. + side: usize, + /// The UDF and its declared result; two frame sides. + udf: PythonRelationPlan, + }, + /// An aggregate whose calls are all grouped-aggregate Python UDFs, over + /// `input` clustered by the key columns and sorted by them: per group, + /// the values of `key_values` in its first row, then one value per call. + PythonAggregate { + /// The input. + input: Box, + /// The attribute of each input column. + layout: Vec, + /// The columns of `input` groups are compared by: the keys, or a + /// collated key's comparison key. + keys: Vec, + /// The key columns of `input` the output shows, as many as `keys`. + key_values: Vec, + /// The calls, in output column order after the keys. + calls: Vec, + }, /// `TABLESAMPLE`: each row of `input` kept when its draw of input /// partition `i`'s `XORShiftRandom(seed + i)` is in `[lower, upper)` /// (Spark's `SampleExec` without replacement). An unseeded sample's seed @@ -1038,6 +1086,10 @@ plan_data!(enum PhysNode { ReplicateRows { input, layout, count, subqueries }, Sample { input, lower, upper, seed }, PythonEval { input, layout, calls }, + PythonMap { input, udf }, + PythonGroupMap { input, keys, udf }, + PythonCoGroupMap { input, keys, side, udf }, + PythonAggregate { input, layout, keys, key_values, calls }, Generate { input, layout, kind, outer, arg, subqueries }, Expand { input, layout, projections, nullable, subqueries }, Sort { input, keys }, @@ -1058,6 +1110,9 @@ plan_data!(enum PhysNode { LoopRef { id, layout }, }); plan_data!(struct PythonCallPlan { call, output }); +plan_data!(struct PythonRelationPlan { udf, sides, output }); +plan_data!(struct PythonFrameSide { columns, names, types }); +plan_data!(struct PythonOutputColumn { name, data_type, nullable }); plan_data!(struct LoopLimits { level, rows, pushed_limit }); plan_data!(struct WritePlan { format, @@ -1143,6 +1198,10 @@ impl PhysNode { | Self::WindowGroupLimit { input, .. } | Self::Aggregate { input, .. } | Self::PythonEval { input, .. } + | Self::PythonMap { input, .. } + | Self::PythonGroupMap { input, .. } + | Self::PythonCoGroupMap { input, .. } + | Self::PythonAggregate { input, .. } | Self::WriteFiles { input, .. } => vec![input.as_ref()], Self::Exchange(e) => vec![e.input.as_ref()], Self::HashJoin(j) => vec![j.probe.as_ref(), j.build.as_ref()], @@ -1190,6 +1249,10 @@ impl PhysNode { | Self::WindowGroupLimit { input, .. } | Self::Aggregate { input, .. } | Self::PythonEval { input, .. } + | Self::PythonMap { input, .. } + | Self::PythonGroupMap { input, .. } + | Self::PythonCoGroupMap { input, .. } + | Self::PythonAggregate { input, .. } | Self::WriteFiles { input, .. } => vec![input.as_mut()], Self::Exchange(e) => vec![e.input.as_mut()], Self::HashJoin(j) => vec![j.probe.as_mut(), j.build.as_mut()], @@ -1214,6 +1277,43 @@ pub struct PythonCallPlan { pub output: AttrId, } +/// The Python UDF of a relation node (design P3-T085 §3.1): a reference by +/// name and version (never the callable), the columns of its frames and its +/// declared result. +#[derive(Debug, Clone, PartialEq)] +pub struct PythonRelationPlan { + /// The UDF. + pub udf: krishiv_plan::PythonUdfRef, + /// The frame columns of each side: one, two for a cogroup. + pub sides: Vec, + /// The declared result columns. + pub output: Vec, +} + +/// The columns of one side's frames: a grouped kind's key columns first +/// (computed keys included, as many as the node has keys), then the side's +/// data columns, as Spark lays them out. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PythonFrameSide { + /// The input columns, in frame order. + pub columns: Vec, + /// Their names, as the UDF sees them. + pub names: Vec, + /// Their Spark types. + pub types: Vec, +} + +/// One declared result column of a Python relation node. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct PythonOutputColumn { + /// The name. + pub name: String, + /// The Spark type. + pub data_type: krishiv_types::DataType, + /// Whether it may hold nulls. + pub nullable: bool, +} + /// An executable plan as data: the root node, the schema of its result /// batches, the printable tree and the expression environment. Hand it to /// a [`Scheduler`](crate::Scheduler); [`instantiate`] builds the operators. diff --git a/crates/krishiv-exec/src/planner/build.rs b/crates/krishiv-exec/src/planner/build.rs index 3f7a2bf0d..b568b2e50 100644 --- a/crates/krishiv-exec/src/planner/build.rs +++ b/crates/krishiv-exec/src/planner/build.rs @@ -586,6 +586,11 @@ impl<'a> Planner<'a> { output, .. } => { + if *kind == krishiv_plan::GeneratorKind::PythonUdtf { + return Err(super::python_relation::later_kind( + "a lateral Python table function", + )); + } let input = self.plan(input)?; self.generate(input, *kind, *outer, arg, output) } @@ -603,6 +608,14 @@ impl<'a> Planner<'a> { partition_by, order_by, } => { + if exprs + .iter() + .any(|e| super::python_udf::has_python_aggregate_udf(e)) + { + return Err(super::python_relation::later_kind( + "a Python aggregate UDF over a window", + )); + } let input = self.plan(input)?; self.window(input, exprs, partition_by, order_by, None) } @@ -623,6 +636,29 @@ impl<'a> Planner<'a> { num_slices, output, } => Ok(range_node(*start, *end, *step, *num_slices, output)), + PlanNode::MapPartitions { + input, + udf, + output, + barrier, + } => self.map_partitions(input, udf, output, *barrier), + PlanNode::GroupMap { + input, + keys, + udf, + output, + } => self.group_map(input, keys, udf, output), + PlanNode::CoGroupMap { + left, + right, + left_keys, + right_keys, + udf, + output, + } => self.cogroup_map(left, right, left_keys, right_keys, udf, output), + PlanNode::PythonTableFunction { .. } => Err(super::python_relation::later_kind( + "a Python table function", + )), _ => { let text = p.to_string(); let name = text.split_whitespace().next().unwrap_or_default(); diff --git a/crates/krishiv-exec/src/planner/build_project.rs b/crates/krishiv-exec/src/planner/build_project.rs index bb6c703fa..644f1c127 100644 --- a/crates/krishiv-exec/src/planner/build_project.rs +++ b/crates/krishiv-exec/src/planner/build_project.rs @@ -123,6 +123,17 @@ impl Planner<'_> { grouping: &Grouping, aggregates: &[Arc], ) -> Result { + // Grouped-aggregate Python UDFs (design P3-T085 §3.1): alone in their + // aggregate, on a node of their own. + if aggregates.iter().any(has_python_udaf) { + krishiv_plan::check_python_aggregates(aggregates)?; + let Grouping::GroupBy(keys) = grouping else { + return Err(super::build::unsupported( + "grouped-aggregate Python UDFs with ROLLUP, CUBE or GROUPING SETS", + )); + }; + return self.python_aggregate(input, keys, aggregates); + } if let Grouping::GroupBy(keys) = grouping { self.request_moves(node, input, aggregates, true, keys.clone()); } @@ -539,6 +550,23 @@ fn fixed_width(t: &DataType) -> Option { }) } +/// Whether `e` calls a grouped-aggregate Python UDF anywhere (shared nodes +/// visited once). +fn has_python_udaf(e: &Arc) -> bool { + let mut seen = std::collections::HashSet::new(); + let mut stack = vec![Arc::clone(e)]; + while let Some(x) = stack.pop() { + if !seen.insert(Arc::as_ptr(&x).addr()) { + continue; + } + if krishiv_plan::is_python_udaf(&x) { + return true; + } + x.for_each_child(&mut |c| stack.push(Arc::clone(c))); + } + false +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/krishiv-exec/src/planner/mod.rs b/crates/krishiv-exec/src/planner/mod.rs index 8a9a25ed5..c8a92af52 100644 --- a/crates/krishiv-exec/src/planner/mod.rs +++ b/crates/krishiv-exec/src/planner/mod.rs @@ -39,6 +39,7 @@ mod orient; mod probe_select; mod properties; pub(crate) mod pushdown; +mod python_relation; mod python_udf; mod runtime_filter; pub(crate) mod subquery; diff --git a/crates/krishiv-exec/src/planner/python_relation.rs b/crates/krishiv-exec/src/planner/python_relation.rs new file mode 100644 index 000000000..05fd9962f --- /dev/null +++ b/crates/krishiv-exec/src/planner/python_relation.rs @@ -0,0 +1,680 @@ +// SPDX-License-Identifier: Apache-2.0 +//! The Python UDF relation nodes planned (design P3-T085 §3.1): map, +//! grouped map, cogrouped map and grouped aggregate, with the data movement +//! Spark gives them (`FlatMapGroupsInBatchExec.scala` L62-71, +//! `FlatMapCoGroupsInBatchExec.scala` L56-64, `ArrowAggregatePythonExec.scala` +//! L82-95 at v4.2.0): a map runs where its input is; a grouped kind gets a +//! hash exchange on its keys (one partition without keys) and a local sort +//! on them, the engine's spillable sort, so no group is held whole. +//! +//! A cogroup's two inputs go through one stream: each side projected to +//! `keys, side, left columns, right columns` (the other side's columns +//! null), a union, then the exchange on the keys and the sort on the keys +//! and the side. Each key's left rows then come before its right rows on +//! one driver, which is what Spark's two co-partitioned sorted inputs and +//! its `CoGroupedIterator` give the UDF. +//! +//! Barrier mode, window aggregates and UDTFs are refused here +//! (`KRISHIV_UNSUPPORTED_FEATURE`): the window and table kinds run with +//! PR-26c. + +use std::collections::HashSet; +use std::sync::Arc; + +use arrow_schema::{Field, Schema, SchemaRef}; +use krishiv_plan::{ + AttrId, BoundColumn, Expr, ExprNode, LogicalPlan, PythonUdfKind, PythonUdfRef, ResolvedExpr, + is_python_udaf, +}; +use krishiv_types::DataType; + +use super::build::{Planned, Planner, internal, unsupported}; +use super::explain::ExplainNode; +use super::properties::{Distribution, Placement}; +use crate::error::Result; +use crate::operator::{Operator as _, Source as _}; +use crate::operators::SortKey; +use crate::physical::{ + PhysNode, ProjectItemPlan, PythonCallPlan, PythonFrameSide, PythonOutputColumn, + PythonRelationPlan, compile_all, project_op, sort_op, +}; + +/// The error of a Python UDF kind that runs only from PR-26c. +pub(crate) fn later_kind(what: &str) -> krishiv_common::error::KrishivError { + unsupported(&format!( + "{what} (Python window aggregate and table UDFs run from P3-T014 PR-26c)" + )) +} + +/// Refuses `udf` unless its eval type is of `kind`. +fn check_kind(udf: &PythonUdfRef, kind: PythonUdfKind, node: &str) -> Result<()> { + if udf.eval_type.kind() == kind { + Ok(()) + } else { + Err(internal(&format!( + "{node} over a Python UDF of eval type {}", + udf.eval_type.code() + ))) + } +} + +/// The declared columns as plan data. +fn outputs(output: &[BoundColumn]) -> Vec { + output + .iter() + .map(|c| PythonOutputColumn { + name: c.name.clone(), + data_type: c.data_type.clone(), + nullable: c.nullable, + }) + .collect() +} + +/// The output schema of declared columns (canonical carriers). +fn output_schema(output: &[BoundColumn]) -> SchemaRef { + Arc::new(Schema::new( + output + .iter() + .map(|c| Field::new(&c.name, c.data_type.to_arrow(), c.nullable)) + .collect::>(), + )) +} + +/// A frame side: the key columns `key_at` (computed keys included), then +/// the data columns `columns` of a relation whose declared columns are +/// `cols`. Spark's layout (`PandasGroupUtils.resolveArgOffsets` at v4.2.0): +/// the worker builds `key` from the leading columns and the frame from the +/// rest. +fn side_of( + keys: &[Arc], + key_at: &[usize], + columns: &[usize], + cols: &[BoundColumn], +) -> Result { + let mut side = PythonFrameSide { + columns: key_at.iter().chain(columns).copied().collect(), + names: Vec::with_capacity(keys.len() + cols.len()), + types: Vec::with_capacity(keys.len() + cols.len()), + }; + for k in keys { + let name = match &k.node { + ExprNode::Resolved(ResolvedExpr::Column(c)) => c.name.clone(), + _ => k.to_string(), + }; + side.names.push(name); + side.types.push( + k.resolved_type() + .ok_or_else(|| internal("a Python UDF grouping key without a type"))?, + ); + } + side.names.extend(cols.iter().map(|c| c.name.clone())); + side.types.extend(cols.iter().map(|c| c.data_type.clone())); + Ok(side) +} + +/// The columns of `p`'s declared output. +fn columns_of(p: &LogicalPlan) -> Result> { + p.resolved_output() + .ok_or_else(|| internal("a Python UDF relation over an unresolved input")) +} + +impl Planner<'_> { + /// `mapInPandas` / `mapInArrow`. + pub(crate) fn map_partitions( + &mut self, + input: &Arc, + udf: &PythonUdfRef, + output: &[BoundColumn], + barrier: bool, + ) -> Result { + if barrier { + return Err(unsupported( + "barrier execution mode (isBarrier) of mapInPandas and mapInArrow", + )); + } + check_kind(udf, PythonUdfKind::Map, "MapPartitions")?; + let cols = columns_of(input)?; + let planned = self.plan(input)?; + let columns = cols + .iter() + .map(|c| planned.column(c.id)) + .collect::>>()?; + let plan = PythonRelationPlan { + udf: udf.clone(), + sides: vec![side_of(&[], &[], &columns, &cols)?], + output: outputs(output), + }; + let explain = format!("{}#{}", udf.name, udf.version); + Ok(relation( + planned, + |input| PhysNode::PythonMap { input, udf: plan }, + output, + ("PythonMap", explain), + )) + } + + /// `groupBy(keys).applyInPandas` / `applyInArrow`. Scalar Python calls + /// in the keys run in a `PythonEval` below (Spark computes the keys in + /// a `Project` under `FlatMapGroupsInPandas`, which `ExtractPythonUDFs` + /// at v4.2.0 turns into an `ArrowEvalPython`); the frames keep the + /// keys' own names and types. + pub(crate) fn group_map( + &mut self, + input: &Arc, + keys: &[Arc], + udf: &PythonUdfRef, + output: &[BoundColumn], + ) -> Result { + check_kind(udf, PythonUdfKind::GroupedMap, "GroupMap")?; + let cols = columns_of(input)?; + let planned = self.plan(input)?; + let (planned, computed) = self.extract_python(planned, keys)?; + let (planned, key_columns) = self.with_columns(planned, &computed)?; + let (planned, compared) = self.compared(planned, &computed, &key_columns)?; + let columns = cols + .iter() + .map(|c| planned.column(c.id)) + .collect::>>()?; + let sorted = self.clustered_and_sorted(planned, &compared, &[])?; + let plan = PythonRelationPlan { + udf: udf.clone(), + sides: vec![side_of(keys, &key_columns, &columns, &cols)?], + output: outputs(output), + }; + let explain = format!( + "[{}], {}#{}", + key_names(&sorted, &key_columns), + udf.name, + udf.version + ); + Ok(relation( + sorted, + |input| PhysNode::PythonGroupMap { + input, + keys: compared, + udf: plan, + }, + output, + ("PythonGroupMap", explain), + )) + } + + /// `cogroup(...).applyInPandas` / `applyInArrow`; each side's scalar + /// Python key calls run below its branch, as in [`Self::group_map`]. + pub(crate) fn cogroup_map( + &mut self, + left: &Arc, + right: &Arc, + left_keys: &[Arc], + right_keys: &[Arc], + udf: &PythonUdfRef, + output: &[BoundColumn], + ) -> Result { + check_kind(udf, PythonUdfKind::CogroupedMap, "CoGroupMap")?; + if left_keys.len() != right_keys.len() { + return Err(internal("a cogroup with different key counts")); + } + let (left_cols, right_cols) = (columns_of(left)?, columns_of(right)?); + let l = self.plan(left)?; + let (l, left_computed) = self.extract_python(l, left_keys)?; + let (l, lk) = self.with_columns(l, &left_computed)?; + let (l, lc) = self.compared(l, &left_computed, &lk)?; + let r = self.plan(right)?; + let (r, right_computed) = self.extract_python(r, right_keys)?; + let (r, rk) = self.with_columns(r, &right_computed)?; + let (r, rc) = self.compared(r, &right_computed, &rk)?; + let k = lk.len(); + // One layout for both branches: keys, side, left, right, then (for + // collated keys) the keys compared. + let width = k + 1 + left_cols.len() + right_cols.len(); + let keyed = lc != lk || rc != rk; + let total = if keyed { width + k } else { width }; + let layout: Vec = (0..total).map(|_| self.fresh_id()).collect(); + let nullable = vec![true; total]; + let mut branches = Vec::with_capacity(2); + for (side, (planned, keys, compared)) in + [(l, &lk, &lc), (r, &rk, &rc)].into_iter().enumerate() + { + let mut items = Vec::with_capacity(total); + for (i, &c) in keys.iter().enumerate() { + items.push(ProjectItemPlan::Column { + index: c, + name: Some(format!("_k{i}")), + }); + } + items.push(ProjectItemPlan::Expr { + name: "_side".to_owned(), + expr: literal(krishiv_plan::Literal::Integer(i32::from(side == 1))), + }); + for (own, (cols, prefix)) in [(&left_cols, "_l"), (&right_cols, "_r")] + .into_iter() + .enumerate() + { + for (i, c) in cols.iter().enumerate() { + let name = format!("{prefix}{i}"); + if own == side { + items.push(ProjectItemPlan::Column { + index: planned.column(c.id)?, + name: Some(name), + }); + } else { + items.push(ProjectItemPlan::Expr { + name, + expr: literal(krishiv_plan::Literal::Null(Some(c.data_type.clone()))), + }); + } + } + } + if keyed { + for (i, &c) in compared.iter().enumerate() { + items.push(ProjectItemPlan::Column { + index: c, + name: Some(format!("_c{i}")), + }); + } + } + let op = project_op( + &planned.schema, + &planned.layout, + &items, + Some(&nullable), + &self.no_values(), + )?; + let input_layout = planned.layout.clone(); + let nullable = nullable.clone(); + branches.push(planned.wrap( + |input| PhysNode::Project { + input, + layout: input_layout, + items, + nullable: Some(nullable), + subqueries: Vec::new(), + }, + op.schema(), + layout.clone(), + ("Project", format!("cogroup side {side}")), + )); + } + let (Some(rb), Some(lb)) = (branches.pop(), branches.pop()) else { + return Err(internal("a cogroup without its two sides")); + }; + if lb.schema.fields() != rb.schema.fields() { + return Err(unsupported( + "a cogroup whose keys or columns are stored in different types on its two sides", + )); + } + let union = Planned { + ordering: Vec::new(), + schema: Arc::clone(&lb.schema), + explain: ExplainNode::new("Union", String::new(), vec![lb.explain, rb.explain]), + node: PhysNode::Union { + inputs: vec![lb.node, rb.node], + }, + layout, + sites: Vec::new(), + placement: Placement::Unknown, + }; + let key_columns: Vec = (0..k).collect(); + let compared: Vec = if keyed { + (width..total).collect() + } else { + key_columns.clone() + }; + let sorted = self.clustered_and_sorted(union, &compared, &[k])?; + let plan = PythonRelationPlan { + udf: udf.clone(), + sides: vec![ + side_of( + left_keys, + &key_columns, + &(k + 1..k + 1 + left_cols.len()).collect::>(), + &left_cols, + )?, + side_of( + right_keys, + &key_columns, + &(k + 1 + left_cols.len()..width).collect::>(), + &right_cols, + )?, + ], + output: outputs(output), + }; + let explain = format!("{k} keys, {}#{}", udf.name, udf.version); + Ok(relation( + sorted, + |input| PhysNode::PythonCoGroupMap { + input, + keys: compared, + side: k, + udf: plan, + }, + output, + ("PythonCoGroupMap", explain), + )) + } + + /// An `Aggregate` whose aggregate calls are grouped-aggregate Python + /// UDFs (checked by the caller to be the only aggregates): the keys, + /// one column per distinct call, then the aggregate's own list over them. + pub(crate) fn python_aggregate( + &mut self, + input: &Arc, + keys: &[Arc], + aggregates: &[Arc], + ) -> Result { + let mut calls: Vec> = Vec::new(); + let mut seen = HashSet::new(); + for a in aggregates { + udaf_calls(a, &mut calls, &mut seen); + } + for call in &calls { + let f = function_of(call)?; + refuse_call_syntax(f)?; + if f.args.iter().any(|a| super::build::has_subquery(a)) { + return Err(unsupported("a subquery in the arguments of a Python UDF")); + } + } + // Scalar Python calls in the keys and the calls' arguments run below + // (Spark's `ExtractGroupingPythonUDFFromAggregate` and + // `ExtractPythonUDFs` at v4.2.0); the node reads their columns. + let mut wanted: Vec> = keys.to_vec(); + for call in &calls { + wanted.extend(function_of(call)?.args.iter().cloned()); + } + let planned = self.plan(input)?; + let (planned, wanted) = self.extract_python(planned, &wanted)?; + let (key_exprs, mut arg_exprs) = wanted + .split_at_checked(keys.len()) + .ok_or_else(|| internal("Python aggregate key arity"))?; + let (planned, key_columns) = self.with_columns(planned, key_exprs)?; + // A collated key groups by its comparison key; each group shows its + // first row's key (Spark's `GroupedIterator` key at v4.2.0). The + // collation rewrite leaves this aggregate alone for that. + let (planned, compared) = self.compared(planned, key_exprs, &key_columns)?; + let prepared = self.prepare(&[])?; + // Each call as the node runs it: over the arguments' columns. + let mut run = Vec::with_capacity(calls.len()); + for call in &calls { + let f = function_of(call)?; + let (args, rest) = arg_exprs + .split_at_checked(f.args.len()) + .ok_or_else(|| internal("Python aggregate argument arity"))?; + arg_exprs = rest; + let refs: Vec<&Expr> = args.iter().map(AsRef::as_ref).collect(); + compile_all( + &refs, + &planned.layout, + &planned.schema, + &self.values(&prepared), + )?; + run.push(Arc::new( + Expr::new(ExprNode::Resolved(ResolvedExpr::Function( + krishiv_plan::ResolvedFunction { + args: args.to_vec(), + ..f.clone() + }, + ))) + .with_origin(call.origin), + )); + } + let sorted = self.clustered_and_sorted(planned, &compared, &[])?; + let mut layout: Vec = key_columns + .iter() + .map(|&c| sorted.column_attr(c)) + .collect::>()?; + let mut fields: Vec = key_columns + .iter() + .map(|&c| sorted.schema.field(c).clone()) + .collect(); + let mut plans = Vec::with_capacity(calls.len()); + let mut columns: Vec<(Arc, BoundColumn)> = Vec::with_capacity(calls.len()); + for (i, (call, ran)) in calls.iter().zip(&run).enumerate() { + let f = function_of(call)?; + let column = BoundColumn { + id: self.fresh_id(), + name: format!("pythonUDAF{i}"), + data_type: f.data_type.clone(), + nullable: true, + }; + fields.push(Field::new(&column.name, f.data_type.to_arrow(), true)); + layout.push(column.id); + plans.push(PythonCallPlan { + call: Arc::clone(ran), + output: column.id, + }); + columns.push((Arc::clone(call), column)); + } + let detail = run + .iter() + .zip(&columns) + .map(|(c, (_, col))| format!("{c} AS {}", col.name)) + .collect::>() + .join(", "); + let input_layout = sorted.layout.clone(); + let key_values = key_columns.clone(); + let schema = Arc::new(Schema::new(fields)); + let node = sorted.wrap( + |input| PhysNode::PythonAggregate { + input, + layout: input_layout, + keys: compared, + key_values, + calls: plans, + }, + schema, + layout.clone(), + ("PythonAggregate", detail), + ); + // The aggregate's list over the keys and the calls' columns. + let key_bound: Vec<(Arc, BoundColumn)> = keys + .iter() + .zip(&layout) + .filter(|(k, _)| column_of(k).is_none()) + .map(|(k, id)| { + let t = k.resolved_type().unwrap_or(DataType::Null); + ( + Arc::clone(k), + BoundColumn { + id: *id, + name: k.to_string(), + data_type: t, + nullable: true, + }, + ) + }) + .collect(); + let exprs = aggregates + .iter() + .map(|e| replaced(e, &columns, &key_bound)) + .collect::>>()?; + // Scalar Python calls over the results run above (Spark's + // `ExtractPythonUDFFromAggregate` at v4.2.0). + let (node, exprs) = self.extract_python( + Planned { + placement: Placement::Unknown, + ordering: Vec::new(), + sites: Vec::new(), + ..node + }, + &exprs, + )?; + let mut out = self.project(node, &exprs)?; + out.sites.clear(); + Ok(out) + } + + /// The columns a relation's grouping `keys` (at `columns` of `input`) + /// are compared by: a collated key's comparison key, appended + /// (`krishiv_optimizer::collation_key`, the key the aggregates, sorts + /// and exchanges of collated values use), else the key itself. Spark + /// groups and cogroups by the collation (`GroupedIterator`'s and + /// `CoGroupedIterator`'s orderings at v4.2.0); the frames keep each + /// row's own key values. + fn compared( + &mut self, + input: Planned, + keys: &[Arc], + columns: &[usize], + ) -> Result<(Planned, Vec)> { + let mut exprs = Vec::with_capacity(keys.len()); + let mut changed = false; + for (key, &at) in keys.iter().zip(columns) { + let compared = krishiv_optimizer::collation_key(key)?; + if Arc::ptr_eq(&compared, key) { + // The key's column as it is (a computed key reads its own). + let column = input.schema.field(at).clone(); + exprs.push(Arc::new(Expr::new(ExprNode::Resolved( + ResolvedExpr::Column(BoundColumn { + id: input.column_attr(at)?, + name: column.name().clone(), + data_type: key + .resolved_type() + .ok_or_else(|| internal("a Python UDF grouping key without a type"))?, + nullable: column.is_nullable(), + }), + )))); + } else { + changed = true; + exprs.push(compared); + } + } + if !changed { + return Ok((input, columns.to_vec())); + } + self.with_columns(input, &exprs) + } + + /// `input` placed by `keys` (one partition without) and sorted by + /// `keys` then `then`, ascending. + fn clustered_and_sorted( + &mut self, + input: Planned, + keys: &[usize], + then: &[usize], + ) -> Result { + let input = if keys.is_empty() { + self.require(input, Distribution::Single)? + } else { + let attrs: Vec = keys + .iter() + .map(|&c| input.column_attr(c)) + .collect::>()?; + self.require(input, Distribution::Clustered(&attrs))? + }; + let sort: Vec = keys.iter().chain(then).map(|&c| SortKey::asc(c)).collect(); + let schema = sort_op(&input.schema, &sort, None)?.schema(); + let layout = input.layout.clone(); + let detail = sort + .iter() + .map(|k| input.schema.field(k.column).name().clone()) + .collect::>() + .join(", "); + Ok(input.wrap( + |input| PhysNode::Sort { input, keys: sort }, + schema, + layout, + ("Sort", format!("[{detail}]")), + )) + } +} + +/// A relation node over `input` declaring `output`. +fn relation( + input: Planned, + make: impl FnOnce(Box) -> PhysNode, + output: &[BoundColumn], + explain: (&str, String), +) -> Planned { + Planned { + node: make(Box::new(input.node)), + schema: output_schema(output), + layout: output.iter().map(|c| c.id).collect(), + explain: ExplainNode::new(explain.0, explain.1, vec![input.explain]), + sites: Vec::new(), + placement: Placement::Unknown, + ordering: Vec::new(), + } +} + +/// The names of `keys`, columns of `input`, joined. +fn key_names(input: &Planned, keys: &[usize]) -> String { + keys.iter() + .map(|&k| input.schema.field(k).name().clone()) + .collect::>() + .join(", ") +} + +fn literal(l: krishiv_plan::Literal) -> Arc { + Arc::new(Expr::new(ExprNode::Literal(l))) +} + +/// Refuses a Python UDAF call written with `DISTINCT`, `FILTER (WHERE ...)` +/// or `WITHIN GROUP`: the operator feeds every row of a group to the UDF, +/// so running such a call would silently ignore the clause. Spark refuses +/// them too (a `PythonUDF` is not an `AggregateFunction`); its class is not +/// verified against a Spark run yet, so the refusal is +/// `KRISHIV_UNSUPPORTED_FEATURE`. +fn refuse_call_syntax(f: &krishiv_plan::ResolvedFunction) -> Result<()> { + let clause = if f.distinct { + "DISTINCT" + } else if f.filter.is_some() { + "FILTER (WHERE ...)" + } else if !f.within_group.is_empty() { + "WITHIN GROUP (ORDER BY ...)" + } else { + return Ok(()); + }; + Err(unsupported(&format!( + "{clause} on a call of the Python aggregate UDF {}", + f.id.name() + ))) +} + +/// The function `call` is. +fn function_of(call: &Expr) -> Result<&krishiv_plan::ResolvedFunction> { + match &call.node { + ExprNode::Resolved(ResolvedExpr::Function(f)) => Ok(f), + _ => Err(internal("a Python UDAF call is not a function")), + } +} + +/// The attribute `e` reads if it is a plain column. +fn column_of(e: &Expr) -> Option { + match &e.node { + ExprNode::Resolved(ResolvedExpr::Column(c)) => Some(c.id), + _ => None, + } +} + +/// The distinct Python UDAF calls of `e`, in first-seen order. +fn udaf_calls(e: &Arc, out: &mut Vec>, seen: &mut HashSet) { + if !seen.insert(std::ptr::from_ref::(e).addr()) { + return; + } + if is_python_udaf(e) { + if !out.iter().any(|o| o == e) { + out.push(Arc::clone(e)); + } + return; + } + e.for_each_child(&mut |c| udaf_calls(c, out, seen)); +} + +/// `e` reading each call's column and each computed key's column. +fn replaced( + e: &Arc, + calls: &[(Arc, BoundColumn)], + keys: &[(Arc, BoundColumn)], +) -> Result> { + Arc::clone(e).transform_down(&mut |n: Arc| -> Result> { + Ok(calls + .iter() + .chain(keys) + .find(|(x, _)| *x == n) + .map_or(n, |(_, column)| { + Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Column( + column.clone(), + )))) + })) + }) +} diff --git a/crates/krishiv-exec/src/planner/python_udf.rs b/crates/krishiv-exec/src/planner/python_udf.rs index 547e40800..561d5c23e 100644 --- a/crates/krishiv-exec/src/planner/python_udf.rs +++ b/crates/krishiv-exec/src/planner/python_udf.rs @@ -54,18 +54,38 @@ fn node_key(e: &Expr) -> usize { /// Whether `e` holds a Python UDF call anywhere (not inside subqueries, /// which are planned on their own). Each shared node is visited once. pub(crate) fn has_python_udf(e: &Expr) -> bool { - has_python_udf_in(e, &mut HashSet::new()) + has_python_udf_in(e, &|_| true, &mut HashSet::new()) } -fn has_python_udf_in(e: &Expr, seen: &mut HashSet) -> bool { - if python_call(e).is_some() { +/// Whether `e` holds a call of a Python aggregate UDF (the grouped or +/// window aggregate kinds): in a window, Spark's window aggregate, which +/// runs from P3-T014 PR-26c. +pub(crate) fn has_python_aggregate_udf(e: &Expr) -> bool { + has_python_udf_in( + e, + &|t| { + matches!( + t.kind(), + krishiv_plan::PythonUdfKind::GroupedAgg | krishiv_plan::PythonUdfKind::WindowAgg + ) + }, + &mut HashSet::new(), + ) +} + +fn has_python_udf_in( + e: &Expr, + of: &dyn Fn(PythonEvalType) -> bool, + seen: &mut HashSet, +) -> bool { + if python_call(e).is_some_and(of) { return true; } if !seen.insert(node_key(e)) { return false; } let mut found = false; - e.for_each_child(&mut |c| found = found || has_python_udf_in(c, seen)); + e.for_each_child(&mut |c| found = found || has_python_udf_in(c, of, seen)); found } @@ -179,6 +199,16 @@ impl Planner<'_> { let ExprNode::Resolved(ResolvedExpr::Function(f)) = &call.node else { return Err(internal("a Python UDF call is not a function")); }; + // A relation kind (design P3-T085) as a column expression: only + // the scalar kinds run in a `PythonEval`. + if let FunctionId::PythonUdf { eval_type, .. } = &f.id + && !eval_type.is_scalar() + { + return Err(super::build::unsupported(&format!( + "a Python UDF of eval type {} as a column expression", + eval_type.code() + ))); + } // The arguments must compile over the input, as on the worker. let args: Vec<&Expr> = f.args.iter().map(AsRef::as_ref).collect(); compile_all(&args, &input.layout, &input.schema, &self.values(&prepared))?; diff --git a/crates/krishiv-exec/src/python_eval.rs b/crates/krishiv-exec/src/python_eval.rs index be094a03a..bc93606de 100644 --- a/crates/krishiv-exec/src/python_eval.rs +++ b/crates/krishiv-exec/src/python_eval.rs @@ -54,7 +54,7 @@ use arrow_array::{Array, ArrayRef, RecordBatch, RecordBatchOptions}; use arrow_schema::{DataType as ArrowType, Field, Schema, SchemaRef}; use futures::future::BoxFuture; use futures::task::noop_waker_ref; -use krishiv_common::error::KrishivError; +use krishiv_common::error::{KrishivError, KrishivUdfNotAvailable}; use krishiv_memory::{Reservation, SpillScope}; use krishiv_plan::PythonEvalType; use krishiv_types::DataType; @@ -95,6 +95,24 @@ pub trait UdfRuntime: Send + Sync + fmt::Debug { /// `KRISHIV_UDF_NOT_AVAILABLE` when this runtime has no such /// `(name, version)`; any other error fails the query. fn prepare(&self, spec: &UdfCallSpec) -> Result>; + + /// The relation call site of `spec` (design P3-T085 §3.2): map, + /// grouped, cogrouped and aggregate kinds. The default has none. + /// + /// # Errors + /// `KRISHIV_UDF_NOT_AVAILABLE` when this runtime has no such + /// `(name, version)` or runs no relation kind; any other error fails the + /// query. + fn prepare_relation( + &self, + spec: &crate::python_relation::RelationCallSpec, + ) -> Result> { + Err(KrishivUdfNotAvailable { + name: spec.name.clone(), + version: spec.version.to_string(), + } + .into()) + } } /// One prepared call site. diff --git a/crates/krishiv-exec/src/python_relation.rs b/crates/krishiv-exec/src/python_relation.rs new file mode 100644 index 000000000..3a6fa6d12 --- /dev/null +++ b/crates/krishiv-exec/src/python_relation.rs @@ -0,0 +1,178 @@ +// SPDX-License-Identifier: Apache-2.0 +//! The runtime seam of the Python UDF relation kinds (design P3-T085 §3.2): +//! map, grouped map, cogrouped map, grouped aggregate (and, later, window +//! aggregate and UDTF). +//! +//! The scalar kinds keep [`PreparedUdf::call`](crate::python_eval::PreparedUdf) +//! (a column per argument in, a column out). A relation kind sees a +//! relation, so its call is a **task fed by frames**: the worker prepares +//! the call site once ([`UdfRuntime::prepare_relation`]), opens one +//! [`UdfTask`] per driver (a partition read by several drivers opens one +//! task per driver, an empty one none), sends it the driver's rows as +//! [`Frame`]s and reads whatever batches it yields. Output rows are not +//! aligned to input rows. +//! +//! [`Frame`] mirrors Spark's grouped wire form (`PythonArrowInput.scala` +//! L251-300 at v4.2.0): a group's batches, cut by +//! `spark.sql.execution.arrow.maxRecordsPerBatch` (and +//! `maxBytesPerBatch`), then a marker. A map task gets batches then +//! [`Frame::End`]; a grouped task ends each group with +//! [`Frame::EndGroup`]; a cogrouped task gets a key's left rows (side 0), +//! then its right rows (side 1), then `EndGroup`. The embedded runtime +//! (PR-48) and the worker pool (PR-26c) carry the same frames: one +//! protocol, two transports. +//! +//! **Keys.** A grouped or cogrouped kind's frames lead with the +//! [`RelationCallSpec::key_count`] key columns (computed keys included), +//! then the side's data columns: Spark's layout +//! (`PandasGroupUtils.resolveArgOffsets`), from which the Python side builds +//! `key` for `f(key, pdf)` and the frame from the rest. Every row of a +//! group has the same key values. +//! +//! **Cogroup sides.** Every key gets a frame for both sides: a side without +//! rows for a key gets one zero-row [`Frame::Batch`] (its columns' names and +//! types in the batch's schema, as [`RelationCallSpec::input_names`] and +//! [`RelationCallSpec::input_types`] list them), so the Python side builds +//! Spark's empty DataFrame for it. The side order stays left then right. +//! +//! [`UdfRuntime::prepare_relation`] has a default that refuses with +//! `KRISHIV_UDF_NOT_AVAILABLE`, so a runtime that implements only the +//! scalar kinds keeps compiling and a plan calling a relation kind fails +//! on it as on a worker without a runtime. +//! +//! [`UdfRuntime::prepare_relation`]: crate::python_eval::UdfRuntime::prepare_relation + +use std::collections::BTreeMap; +use std::fmt; + +use arrow_array::RecordBatch; +use arrow_schema::SchemaRef; +use futures::future::BoxFuture; +use krishiv_plan::PythonEvalType; +use krishiv_types::DataType; + +use crate::error::Result; + +mod lane; +mod pair; +pub(crate) mod sink; + +pub use sink::{ + DEFAULT_WINDOW, InputColumn, RelationLane, RelationLimits, RelationShape, RelationSink, +}; + +/// What a runtime prepares one relation call site from. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RelationCallSpec { + /// The UDF's canonical (lower-case) name. + pub name: String, + /// The registration version the plan was analyzed with. + pub version: u32, + /// How it is converted and called (a relation kind). + pub eval_type: PythonEvalType, + /// Whether tasks may run concurrently. + pub thread_safe: bool, + /// The Spark types of the columns of each side's frames, in frame + /// order: one side, two for a cogroup. A grouped kind's frames lead + /// with its [`Self::key_count`] key columns, then the side's data + /// columns. + pub input_types: Vec>, + /// The names of the columns of each side's frames, as + /// [`Self::input_types`] orders them: what an empty side's DataFrame is + /// built from. Names need not be unique: a column key is also a data + /// column (`groupBy("a")` gives `a, a, b`), so a runtime splits the + /// columns by [`Self::key_count`], never by name. + pub input_names: Vec>, + /// The grouping keys (0 for a map kind and an aggregate's arguments): + /// the leading columns of every side's frames. With [`Self::arity`] + /// what the Python side needs to call `f(pdf)` or `f(key, pdf)`. + pub key_count: usize, + /// The data columns of a side's frames, keys excluded (the first + /// side's for a cogroup; the other's is its `input_types` length less + /// `key_count`). + pub arity: usize, + /// The declared result: names and Spark types, one per column. + pub output_names: Vec, + /// The declared result's Spark types. + pub output_types: Vec, + /// The Arrow schema the task's batches must have: the declared types' + /// client forms (decimals on `Decimal128`, strings and binaries as + /// views), as [`UdfCallSpec::return_arrow`](crate::python_eval::UdfCallSpec). + pub output_arrow: SchemaRef, + /// The runner configuration (Spark's `PythonRunner` conf): the session's + /// explicit settings the Python side reads. + pub conf: RunnerConf, +} + +/// The settings a Python runner reads (Spark's runner conf). +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct RunnerConf { + /// `spark.sql.execution.arrow.maxRecordsPerBatch`: most rows per frame; + /// 0 is no limit. + pub max_records_per_batch: usize, + /// `spark.sql.execution.arrow.maxBytesPerBatch`: a frame ends at the + /// row that brings its columns to this many bytes (at least one row); + /// 0 is no limit. + pub max_bytes_per_batch: usize, + /// Every explicit session setting, by key. + pub settings: BTreeMap, +} + +/// One message to a [`UdfTask`]. +#[derive(Debug, Clone)] +pub enum Frame { + /// Rows of the current group (of the partition, for a map kind); for + /// a cogroup, zero rows for a side the key has no rows on. + Batch { + /// The relation the rows come from: 0, or 1 for a cogroup's right. + side: u8, + /// The rows, every column in its client form. + batch: RecordBatch, + }, + /// The current group has no more rows. + EndGroup, + /// The partition has no more rows: the last frame. + End, +} + +/// One prepared relation call site. +pub trait PreparedRelationUdf: Send + Sync { + /// One task per driver: the frames in, the result batches out. + /// + /// # Errors + /// What the runtime raises starting it. + fn open(&self) -> Result>; +} + +/// A running relation call. The two methods may be awaited at once (the +/// operator sends while it reads, so a runtime may wait for its results to +/// be taken before it takes more frames); each is awaited by one caller at a +/// time. The work runs off the caller, on the runtime's own threads: +/// polling only moves frames and batches. Dropping the task cancels it. +/// +/// Both take `&self` (and the task is `Sync`) so the operator can await +/// them at once; the operator polls their futures while it holds a lock of +/// its own, so a runtime must not call back into the operator from them. +pub trait UdfTask: Send + Sync { + /// Queues `frame`; waits while the task's bounded input window is full. + /// + /// # Errors + /// The task's failure (a Python error, a crashed worker). + fn send(&self, frame: Frame) -> BoxFuture<'_, Result<()>>; + + /// The next result batch (the client forms of the declared types); + /// `None` once [`Frame::End`] is processed and every result was read. + /// + /// # Errors + /// The task's failure. + fn next(&self) -> BoxFuture<'_, Result>>; +} + +impl fmt::Debug for dyn UdfTask { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("UdfTask") + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/krishiv-exec/src/python_relation/lane.rs b/crates/krishiv-exec/src/python_relation/lane.rs new file mode 100644 index 000000000..6243814e6 --- /dev/null +++ b/crates/krishiv-exec/src/python_relation/lane.rs @@ -0,0 +1,378 @@ +// SPDX-License-Identifier: Apache-2.0 +//! One lane of a relation sink's driver (design P3-T085 §3.2): the frames +//! queued for a [`UdfTask`], the results it yielded, and the future (the +//! *pump*) that moves both while the driver, its `Blocked` future or +//! `finalize` polls it. +//! +//! **Memory (I11).** A queued frame holds the reservation of its rows (the +//! input chunk's, shared by every lane's frame of that chunk and released by +//! the last task that takes it) and of the copies made for its client form. +//! Every result batch is reserved in the driver's consumer as the pump +//! receives it; when the pool refuses, the batch waits in [`LaneIo::refused`] +//! and the pump reads no more results until the driver has taken it +//! (`sink` returns `NeedMemory`; `finalize` stores a map or grouped result +//! as the others, written out when the pool is full, and waits for an +//! aggregate value's memory while the other lanes move): back pressure into +//! the task, never unaccounted results. +//! +//! **Peers.** An aggregate's lanes yield one value per group and a group is +//! stored only once every lane has its value (`pair.rs`), so a lane's values +//! past its peers' cannot be released, nor spilled. A lane therefore reads +//! no result while it holds more rows than a peer whose task still runs +//! *and still has frames it has not taken*: it waits on its own task, +//! holding at most one result batch past that peer, so the peer's values +//! do not wait for memory this lane pins. A peer that took every frame it +//! was given may need later frames before it yields (a runtime batching +//! values across groups); holding back then would stop this lane's task +//! taking frames, fill its window and so stop the frames to every lane, a +//! hang. So the hold lasts only while the peer has frames to take, and is +//! released (the waiting pump woken) when it took its last one, yielded or +//! stopped; the pool then decides (a spill, or a classed error), never a +//! hang. The lane with the fewest rows is never held, and a lane is held +//! only by a peer with strictly fewer rows, so the holds form no cycle. +//! Only an aggregate's lanes are peers ([`Lane::group`]): a map or grouped +//! shape's results are stored one lane at a time. + +use std::collections::VecDeque; +use std::future::poll_fn; +use std::sync::Arc; +use std::task::{Context, Poll, Waker}; + +use arrow_array::RecordBatch; +use futures::future::BoxFuture; +use krishiv_common::error::KrishivError; +use krishiv_memory::Reservation; + +use super::{Frame, PreparedRelationUdf, UdfTask}; +use crate::error::Result; +use crate::sync::{Mutex, lock}; + +/// The memory a queued frame holds until its task takes it. +pub(super) struct FrameMem { + /// The input chunk's rows, shared by the lanes' frames of the chunk. + pub(super) _shared: Option>>, + /// The copies made for this lane's frame (client forms). + pub(super) _own: Reservation, +} + +/// A lane's queues, shared by the driver and the lane's pump. +pub(super) struct LaneIo { + /// Frames not taken by the task yet, each with its memory. + pub(super) inbound: VecDeque<(Frame, FrameMem)>, + /// The memory of the frame the task is taking. + pub(super) sending: Option, + /// Batches the task yielded, each with its reservation, not stored yet. + pub(super) outputs: VecDeque<(RecordBatch, Reservation)>, + /// A batch the pool had no memory for: the pump reads no more results + /// until the driver reserved it and moved it to [`Self::outputs`]. + pub(super) refused: Option, + /// An empty reservation of the driver's consumer, the results' are + /// split from. + pub(super) memory: Reservation, + /// The rows of [`Self::outputs`] and [`Self::refused`], kept as they + /// change (a peer's hold reads it on every result). + rows: usize, + /// The task's output ended, or the pump ended (its last result or a + /// failure). + stopped: bool, + /// The pump waiting while a peer holds the lane back, woken when a + /// peer yields, takes its last frame or stops. + waiting: Option, +} + +impl LaneIo { + pub(super) fn in_flight(&self) -> usize { + self.inbound.len() + usize::from(self.sending.is_some()) + } + + /// Whether results wait for the driver (stored or reserved). + pub(super) fn has_results(&self) -> bool { + !self.outputs.is_empty() || self.refused.is_some() + } + + /// The next result, taken by the driver. + pub(super) fn pop_output(&mut self) -> Option<(RecordBatch, Reservation)> { + let next = self.outputs.pop_front(); + if let Some((b, _)) = &next { + self.took(b.num_rows()); + } + next + } + + /// `rows` rows of the results taken by the driver (the front batch + /// sliced or popped by the caller). + pub(super) fn took(&mut self, rows: usize) { + debug_assert!(rows <= self.rows, "took {rows} of {} rows", self.rows); + self.rows = self.rows.saturating_sub(rows); + } +} + +/// The queues of a driver's lanes, one per lane. +type Peers = Arc<[Arc>]>; + +/// The future that moves a lane's frames to its task and its results back. +enum Pump { + /// No frame yet: the task is not open. + Idle, + Running(BoxFuture<'static, Result<()>>), + Done(Result<()>), +} + +/// One lane of one driver. +#[derive(Clone)] +pub(super) struct Lane { + pub(super) io: Arc>, + pump: Arc>, + /// Every lane's queues of the driver, this one's at `index`. + peers: Peers, + index: usize, +} + +impl Lane { + /// The lanes of one driver, one per reservation, whose results reserve + /// in that reservation's consumer. With `peered` (an aggregate's lanes) + /// each reads no result while it holds more rows than a peer still + /// running that has frames to take; else every lane is alone. + pub(super) fn group( + memories: impl IntoIterator, + peered: bool, + ) -> Vec { + if !peered { + return memories + .into_iter() + .flat_map(|memory| Self::group([memory], true)) + .collect(); + } + let peers: Peers = memories + .into_iter() + .map(|memory| { + Arc::new(Mutex::new(LaneIo { + inbound: VecDeque::new(), + sending: None, + outputs: VecDeque::new(), + refused: None, + memory, + rows: 0, + stopped: false, + waiting: None, + })) + }) + .collect(); + peers + .iter() + .enumerate() + .map(|(index, io)| Self { + io: Arc::clone(io), + pump: Arc::new(Mutex::new(Pump::Idle)), + peers: Arc::clone(&peers), + index, + }) + .collect() + } + + /// Polls the pump; its outcome once it ended (`None` while it runs or + /// before the task opened). + pub(super) fn poll(&self, cx: &mut Context<'_>) -> Option> { + let mut p = lock(&self.pump); + if let Pump::Running(f) = &mut *p + && let Poll::Ready(r) = f.as_mut().poll(cx) + { + *p = Pump::Done(r); + } + match &*p { + Pump::Done(r) => Some(r.clone()), + Pump::Idle | Pump::Running(_) => None, + } + } + + pub(super) fn opened(&self) -> bool { + !matches!(*lock(&self.pump), Pump::Idle) + } + + /// Whether the lane is never opened or its task ended well. + pub(super) fn finished(&self) -> bool { + matches!(*lock(&self.pump), Pump::Idle | Pump::Done(Ok(()))) + } + + /// Queues `frame`, opening the task first if it is not. + pub(super) fn queue( + &self, + udf: &dyn PreparedRelationUdf, + frame: Frame, + memory: FrameMem, + ) -> Result<()> { + { + let mut p = lock(&self.pump); + if matches!(*p, Pump::Idle) { + let task = udf.open()?; + *p = Pump::Running(pump(task, Arc::clone(&self.peers), self.index)); + } + } + lock(&self.io).inbound.push_back((frame, memory)); + Ok(()) + } + + /// Reserves the refused result, if any, without waiting, and moves it + /// to the outputs; the bytes still missing if the pool refuses again. + pub(super) fn try_take_refused(&self) -> Option { + let mut q = lock(&self.io); + let bytes = q.refused.as_ref()?.get_array_memory_size(); + let mut r = q.memory.split(0); + if r.try_grow(bytes).is_err() { + return Some(bytes); + } + if let Some(batch) = q.refused.take() { + q.outputs.push_back((batch, r)); + } + None + } + + /// The refused result, if any, taken without a reservation (its + /// caller stores it at once, writing it out when the pool is full). + pub(super) fn take_refused_unreserved(&self) -> Option { + let mut q = lock(&self.io); + let batch = q.refused.take()?; + q.took(batch.num_rows()); + Some(batch) + } + + /// The future that reserves the refused result, if any, waiting for + /// the memory, and moves it to the outputs. Cancel-safe. + pub(super) fn reserve_refused(&self) -> Option>> { + let (bytes, mut r) = { + let mut q = lock(&self.io); + let bytes = q.refused.as_ref()?.get_array_memory_size(); + (bytes, q.memory.split(0)) + }; + let io = Arc::clone(&self.io); + Some(Box::pin(async move { + r.grow(bytes).await.map_err(KrishivError::from)?; + let mut q = lock(&io); + if let Some(batch) = q.refused.take() { + q.outputs.push_back((batch, r)); + } + Ok(()) + })) + } +} + +/// Whether lane `index` holds more rows than a peer whose task still runs +/// and has frames it has not taken (a peer waiting for input never holds +/// it back, see the module doc). Each queue is locked alone, so two lanes +/// asking at once never deadlock. +fn held_back(peers: &[Arc>], index: usize) -> bool { + let Some(own) = peers.get(index).map(|io| lock(io).rows) else { + return false; + }; + peers.iter().enumerate().any(|(i, io)| { + let q = lock(io); + i != index && !q.stopped && q.in_flight() > 0 && q.rows < own + }) +} + +/// Wakes every peer of lane `index` waiting for it to yield, take its last +/// frame or stop; each waker is taken under the peer's lock and woken +/// after it is released. +fn wake_peers(peers: &[Arc>], index: usize) { + for (i, io) in peers.iter().enumerate() { + if i == index { + continue; + } + let waker = lock(io).waiting.take(); + if let Some(w) = waker { + w.wake(); + } + } +} + +/// Sends the queued frames of lane `index` to `task` while it reads the +/// task's results. +fn pump(task: Box, peers: Peers, index: usize) -> BoxFuture<'static, Result<()>> { + Box::pin(async move { + let task = task; + let Some(io) = peers.get(index).map(Arc::clone) else { + return Err(crate::error::internal(format!("no lane {index}"))); + }; + let send = async { + loop { + // Never woken: whoever queues a frame polls the pump next. + let frame = poll_fn(|_| { + let mut q = lock(&io); + match q.inbound.pop_front() { + Some((frame, memory)) => { + q.sending = Some(memory); + Poll::Ready(frame) + } + None => Poll::Pending, + } + }) + .await; + let end = matches!(frame, Frame::End); + task.send(frame).await?; + let idle = { + let mut q = lock(&io); + q.sending = None; + q.in_flight() == 0 + }; + if idle { + // A peer held back by this lane's frames moves on. + wake_peers(&peers, index); + } + if end { + return Ok::<(), KrishivError>(()); + } + } + }; + let receive = async { + loop { + // Not woken while refused: whoever reserves the refused + // batch polls the pump next. Held back by a peer: woken + // when a peer yields, takes its last frame or stops (the + // waker is set before the check, so a peer that moves + // after it wakes this pump). + poll_fn(|cx| { + if lock(&io).refused.is_some() { + return Poll::Pending; + } + lock(&io).waiting = Some(cx.waker().clone()); + if held_back(&peers, index) { + Poll::Pending + } else { + lock(&io).waiting = None; + Poll::Ready(()) + } + }) + .await; + let Some(batch) = task.next().await? else { + // The output ended: it yields no more rows, so it holds + // no peer back, even before its send side finished. + // `UdfTask::next` ends only once `End` was processed, + // so the send side is about to finish; a runtime that + // ended early while a send hangs (a write to a dead + // worker that blocks) would otherwise hold every peer + // with more rows, a hang instead of the mismatch. + lock(&io).stopped = true; + wake_peers(&peers, index); + return Ok(()); + }; + let bytes = batch.get_array_memory_size(); + { + let mut q = lock(&io); + q.rows += batch.num_rows(); + let mut r = q.memory.split(0); + if r.try_grow(bytes).is_ok() { + q.outputs.push_back((batch, r)); + } else { + q.refused = Some(batch); + } + } + wake_peers(&peers, index); + } + }; + let outcome = futures::future::try_join(send, receive).await.map(drop); + // Ended or failed: this lane holds no peer back any more. + lock(&io).stopped = true; + wake_peers(&peers, index); + outcome + }) +} diff --git a/crates/krishiv-exec/src/python_relation/pair.rs b/crates/krishiv-exec/src/python_relation/pair.rs new file mode 100644 index 000000000..e6e2b0990 --- /dev/null +++ b/crates/krishiv-exec/src/python_relation/pair.rs @@ -0,0 +1,236 @@ +// SPDX-License-Identifier: Apache-2.0 +//! The rows of a Python grouped aggregate as its tasks yield them (design +//! P3-T085 §3.1, review Md2): each lane's task yields one value per group, +//! in group order, so the oldest groups whose keys are still pending are +//! joined with the values every lane has yielded so far and stored at once. +//! A driver keeps one key batch per input batch (the group-start rows taken +//! together), reserved until its last group is joined. + +use std::collections::VecDeque; +use std::sync::Arc; + +use arrow_array::{Array, ArrayRef, RecordBatch, RecordBatchOptions, UInt32Array}; +use arrow_schema::SchemaRef; +use krishiv_memory::Reservation; +use krishiv_types::batch::Batch; + +use super::lane::Lane; +use super::sink::{RelationLane, mismatch}; +use crate::error::{Result, internal}; +use crate::sync::lock; + +/// The group keys of one input batch, not joined with their values yet. +pub(super) struct KeyPart { + rows: RecordBatch, + /// Rows already joined. + offset: usize, + _memory: Reservation, +} + +impl KeyPart { + fn left(&self) -> usize { + self.rows.num_rows() - self.offset + } +} + +/// The group-start rows `starts` of `batch` projected to `keys`, as one +/// batch of `schema`; `None` without starts. +pub(super) fn key_rows( + batch: &RecordBatch, + keys: &[usize], + schema: &SchemaRef, + starts: &[u32], +) -> Result> { + if starts.is_empty() { + return Ok(None); + } + let indices = UInt32Array::from(starts.to_vec()); + let mut columns = Vec::with_capacity(keys.len()); + for &k in keys { + let column = batch + .columns() + .get(k) + .ok_or_else(|| internal(format!("key column {k} out of range")))?; + columns.push( + arrow_select::take::take(column.as_ref(), &indices, None) + .map_err(|e| internal(format!("aggregate keys: {e}")))?, + ); + } + let options = RecordBatchOptions::new().with_row_count(Some(starts.len())); + RecordBatch::try_new_with_options(Arc::clone(schema), columns, &options) + .map(Some) + .map_err(|e| internal(format!("aggregate keys: {e}"))) +} + +/// A driver's pending groups. +#[derive(Default)] +pub(super) struct Pending { + keys: VecDeque, + /// Groups started, and joined, so far. + groups: usize, + joined: usize, +} + +impl Pending { + /// Adds the keys of an input batch's new groups, reserved in `memory`. + pub(super) fn push(&mut self, rows: RecordBatch, memory: Reservation) { + self.groups += rows.num_rows(); + self.keys.push_back(KeyPart { + rows, + offset: 0, + _memory: memory, + }); + } + + fn waiting(&self) -> usize { + self.keys.iter().map(KeyPart::left).sum() + } + + /// The first `n` pending keys, columns concatenated; parts used up are + /// dropped (their memory with them). + fn take_keys(&mut self, n: usize, schema: &SchemaRef) -> Result> { + let mut parts: Vec = Vec::new(); + let mut want = n; + while want > 0 { + let Some(front) = self.keys.front_mut() else { + return Err(internal("aggregate keys ran out".to_owned())); + }; + let len = want.min(front.left()); + parts.push(front.rows.slice(front.offset, len)); + front.offset += len; + want -= len; + if front.left() == 0 { + self.keys.pop_front(); + } + } + let keys = arrow_select::concat::concat_batches(schema, &parts) + .map_err(|e| internal(format!("aggregate keys: {e}")))?; + self.joined += n; + Ok(keys.columns().to_vec()) + } +} + +/// The rows a lane's results hold, each checked to be one column of the +/// lane's declared type. +fn yielded(lane: &RelationLane, state: &Lane) -> Result { + let q = lock(&state.io); + let mut rows = 0; + for (b, _) in &q.outputs { + check_value(lane, b)?; + rows += b.num_rows(); + } + Ok(rows) +} + +fn check_value(lane: &RelationLane, b: &RecordBatch) -> Result<()> { + if b.num_columns() != 1 { + return Err(mismatch( + lane, + format!("{} columns where 1 was declared", b.num_columns()), + )); + } + match lane.client.first() { + Some(client) if b.column(0).data_type() == client => Ok(()), + Some(client) => Err(mismatch( + lane, + format!("{} where {client} was declared", b.column(0).data_type()), + )), + None => Err(internal("an aggregate lane without a type".to_owned())), + } +} + +/// The first `n` values of a lane's results; batches used up are dropped +/// from the queue, their reservations into `memory`. +fn take_values(state: &Lane, n: usize, memory: &mut Vec) -> Result { + let mut q = lock(&state.io); + let mut parts: Vec = Vec::new(); + let mut want = n; + while want > 0 { + let Some((front, _)) = q.outputs.front_mut() else { + return Err(internal("aggregate values ran out".to_owned())); + }; + let len = want.min(front.num_rows()); + parts.push(Arc::clone(front.column(0)).slice(0, len)); + want -= len; + if len == front.num_rows() { + if let Some((_, r)) = q.pop_output() { + memory.push(r); + } + } else { + *front = front.slice(len, front.num_rows() - len); + q.took(len); + } + } + let refs: Vec<&dyn Array> = parts.iter().map(AsRef::as_ref).collect(); + arrow_select::concat::concat(&refs).map_err(|e| internal(format!("aggregate results: {e}"))) +} + +/// The output rows of every group whose values each lane has yielded: the +/// keys, then a value per lane (canonical carriers), with the memory of the +/// results they came from; `None` when no group is complete. +pub(super) fn joined( + pending: &mut Pending, + key_schema: &SchemaRef, + schema: &SchemaRef, + lanes: &[RelationLane], + states: &[Lane], +) -> Result)>> { + let mut n = pending.waiting(); + for (lane, state) in lanes.iter().zip(states) { + let rows = yielded(lane, state)?; + if rows + pending.joined > pending.groups { + return Err(too_many(lane, rows + pending.joined, pending.groups)); + } + n = n.min(rows); + } + if n == 0 { + return Ok(None); + } + let mut columns = pending.take_keys(n, key_schema)?; + let width = columns.len(); + let mut memory = Vec::new(); + for (i, (lane, state)) in lanes.iter().zip(states).enumerate() { + let values = take_values(state, n, &mut memory)?; + let (Some(client), Some(field)) = (lane.client.first(), schema.fields().get(width + i)) + else { + return Err(internal("an aggregate result out of its schema".to_owned())); + }; + columns.push(super::sink::recarried( + lane, + &values, + client, + field.data_type(), + )?); + } + let options = RecordBatchOptions::new().with_row_count(Some(n)); + let out = RecordBatch::try_new_with_options(Arc::clone(schema), columns, &options) + .map_err(|e| internal(format!("a Python aggregate result: {e}")))?; + Ok(Some((Batch::new(out), memory))) +} + +/// Refuses an aggregate driver whose tasks all ended with groups or values +/// left over: a task yielded another number of values than groups. +pub(super) fn check_complete( + pending: &Pending, + lanes: &[RelationLane], + states: &[Lane], +) -> Result<()> { + for (lane, state) in lanes.iter().zip(states) { + let left = yielded(lane, state)?; + if left > 0 || pending.joined < pending.groups { + return Err(too_many(lane, left + pending.joined, pending.groups)); + } + } + Ok(()) +} + +fn too_many( + lane: &RelationLane, + rows: usize, + groups: usize, +) -> krishiv_common::error::KrishivError { + mismatch( + lane, + format!("{rows} values for {groups} groups (one value per group)"), + ) +} diff --git a/crates/krishiv-exec/src/python_relation/sink.rs b/crates/krishiv-exec/src/python_relation/sink.rs new file mode 100644 index 000000000..9a4edea0d --- /dev/null +++ b/crates/krishiv-exec/src/python_relation/sink.rs @@ -0,0 +1,1421 @@ +// SPDX-License-Identifier: Apache-2.0 +//! The sink of the Python relation nodes (`PythonMap`, `PythonGroupMap`, +//! `PythonCoGroupMap`, `PythonAggregate`; design P3-T085 §3.1-§3.2). +//! +//! **Shape.** As `PythonEval`'s (P2-T059): a pipeline break whose sink sends +//! the input to the runtime and returns at once, and whose results are kept +//! in a shared plan's list (`crate::shared`: memory under the driver's +//! consumer, spilled under pressure) that a `SharedSource` reads after +//! `finalize`. Compute workers never run Python and never wait inside +//! `sink`: each driver opens one [`UdfTask`](super::UdfTask) per lane (one +//! UDF; an aggregate has a lane per call) at its first frame, queues frames, +//! and while a lane's bounded window of frames in flight is full returns +//! [`SinkStatus::Blocked`], whose future moves frames and results until the +//! window has room. Unlike `PythonEval`, output rows are not aligned to +//! input rows: a task's batches are stored as it yields them, by `sink` and +//! then by `finalize` while the tasks finish. +//! +//! **Groups.** The planner sorts a grouped input by its keys after a hash +//! exchange on them (one partition without keys), so each driver sees a +//! key's rows together; the sink cuts groups where the key changes +//! (compared as the sort compares them, floats canonical) and ends each +//! with [`Frame::EndGroup`]. A cogroup's input is both relations in one +//! stream sorted by key then side: a key's left rows, then its right rows; +//! a side without rows for a key gets one zero-row frame. +//! Rows are cut into frames of `maxRecordsPerBatch` rows (and, when set, +//! closed by the row that brings the frame's columns to `maxBytesPerBatch` +//! bytes as Spark's Arrow writer sizes them) or at the end of their group +//! or side, as Spark cuts a group's Arrow stream; a group is never held +//! whole. +//! +//! **Window.** The window bounds the frames queued per lane and is checked +//! frame by frame: input rows cut but not framed yet wait in the driver's +//! backlog (raw rows, reserved), are framed as the window frees, and the +//! driver is `Blocked` while a lane's window is full or cuts wait: its +//! future frames the backlog as the windows free, so no input is taken +//! before the backlog drained and the cuts held stay within about one +//! input batch. +//! +//! **Memory (I11).** Every byte the sink holds is reserved in the driver's +//! consumer: input rows held or in the backlog, queued frames (the chunk's +//! rows shared by the lanes' frames, released by the last lane's task to +//! take it, and what a frame's columns hold beyond those rows: a computed +//! argument's result, a client form's copy), the results a task +//! yielded until they are stored (`lane.rs`: a refused result stops the +//! task's reads until reserved), and an aggregate's pending group keys +//! (`pair.rs`). A refusal in `sink` returns `NeedMemory` before the input +//! is taken, so the retry is exact. In `finalize` a refused map or grouped +//! result is stored as the others, written to a spill file when the pool +//! is full; an aggregate's refused value waits for its memory while the +//! other lanes move, which hold at most one result batch past it while it +//! has frames to take (`lane.rs`, Peers). A stored result's bytes pass to the store's +//! reservation through the consumer's credit, never counted twice. +//! +//! **Errors and cancellation.** The first failure of a task fails the +//! query. Dropping the state (query end, cancel, error) drops every task, +//! which cancels it, and every reservation. + +use std::collections::VecDeque; +use std::fmt; +use std::future::poll_fn; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow_array::{Array, ArrayRef, RecordBatch, RecordBatchOptions}; +use arrow_row::OwnedRow; +use arrow_schema::{DataType as ArrowType, Field, Schema, SchemaRef}; +use futures::future::BoxFuture; +use futures::task::noop_waker_ref; +use krishiv_common::error::{KrishivError, KrishivPythonUdfResultMismatch}; +use krishiv_memory::{ConsumerHandle, Reservation, SpillScope}; +use krishiv_types::batch::Batch; + +use super::lane::{FrameMem, Lane}; +use super::pair::{self, Pending}; +use super::{Frame, PreparedRelationUdf}; +use crate::context::TaskContext; +use crate::error::{Result, internal}; +use crate::operator::{GlobalState, LocalState, Sink, SinkStatus}; +use crate::operators::{KeyEncoder, ScalarExpr, SortKey, global_ref, local_mut}; +use crate::shared::SharedSink; +use crate::sync::{Mutex, lock}; + +/// How a relation node's input reaches its tasks. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum RelationShape { + /// Every row of the partition, in batches, then [`Frame::End`] + /// (`mapInPandas`, `mapInArrow`). + Map, + /// Groups of equal `keys` (input columns), each ended by + /// [`Frame::EndGroup`]; no keys is one group. + Grouped { + /// The key columns. + keys: Vec, + }, + /// Two relations in one stream sorted by `keys` then the `side` column + /// (0 left, 1 right): per key, the left rows then the right rows. + CoGrouped { + /// The key columns. + keys: Vec, + /// The column (`INT`) naming each row's relation. + side: usize, + }, + /// As [`Self::Grouped`]; each lane's task yields one value per group, + /// and the output is the values of `key_values` in the group's first + /// row, then one column per lane. + Aggregate { + /// The columns groups are compared by (a collated key's comparison + /// key). + keys: Vec, + /// The key columns the output shows, as many as `keys`. + key_values: Vec, + }, +} + +impl RelationShape { + fn keys(&self) -> Option<&[usize]> { + match self { + Self::Map => None, + Self::Grouped { keys } + | Self::CoGrouped { keys, .. } + | Self::Aggregate { keys, .. } => Some(keys), + } + } +} + +/// One UDF of a relation node: the prepared call site and the columns of its +/// frames, per side (expressions over the input; a column read for a map or +/// grouped kind, an argument for an aggregate). +pub struct RelationLane { + /// The UDF's name, for errors. + pub name: String, + /// The prepared call site. + pub udf: Arc, + /// Each side's frame columns, evaluated over the input. + pub sides: Vec>>, + /// Each side's frame column names. + pub names: Vec>, + /// The Arrow types the task's batches must have (client forms), one per + /// result column. + pub client: Vec, +} + +impl fmt::Debug for RelationLane { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("RelationLane") + .field("name", &self.name) + .field("sides", &self.sides) + .field("names", &self.names) + .field("client", &self.client) + .finish_non_exhaustive() + } +} + +/// Frames in flight per lane when nothing else is configured. +pub const DEFAULT_WINDOW: usize = 4; + +/// The sink side of a Python relation node. +pub struct RelationSink { + shape: RelationShape, + /// Shared with the drivers' `Blocked` futures, which frame with them. + lanes: Arc<[RelationLane]>, + encoder: Option>, + /// The output (canonical carriers). + schema: SchemaRef, + /// An aggregate's key columns, as their rows are kept. + key_schema: Option, + input: SchemaRef, + max_rows: usize, + max_bytes: usize, + window: usize, + store: Arc, +} + +impl fmt::Debug for RelationSink { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("RelationSink") + .field("shape", &self.shape) + .field("lanes", &self.lanes) + .field("max_rows", &self.max_rows) + .field("window", &self.window) + .finish_non_exhaustive() + } +} + +/// Batching and buffering limits of a [`RelationSink`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct RelationLimits { + /// Most rows per frame; 0 is no limit. + pub max_rows: usize, + /// The bytes that end a frame: its rows up to the one that brings its + /// columns to them (Spark's `maxBytesPerBatch`); 0 is no limit. + pub max_bytes: usize, + /// Frames in flight per lane before the driver blocks. + pub window: usize, +} + +impl RelationSink { + /// A sink over input of `input` running `lanes` as `shape` says; its + /// rows have `schema` (canonical carriers). Finished rows spill to + /// `scope`. + /// + /// # Errors + /// `INTERNAL_ERROR` for lanes that do not fit the shape; + /// `KRISHIV_UNSUPPORTED_FEATURE` for a key type that cannot be compared. + pub fn try_new( + input: SchemaRef, + shape: RelationShape, + lanes: Vec, + schema: SchemaRef, + limits: RelationLimits, + scope: Option, + ) -> Result> { + let sides = match shape { + RelationShape::CoGrouped { .. } => 2, + _ => 1, + }; + if lanes.is_empty() + || lanes + .iter() + .any(|l| l.sides.len() != sides || l.names.len() != sides) + { + return Err(internal( + "Python relation lanes do not fit their shape".to_owned(), + )); + } + let encoder = match shape.keys() { + Some(keys) => { + let keys: Vec = keys.iter().map(|&k| SortKey::asc(k)).collect(); + Some(KeyEncoder::new(&input, &keys)?) + } + None => None, + }; + let key_schema = match &shape { + RelationShape::Aggregate { keys, key_values } => { + if key_values.len() != keys.len() { + return Err(internal( + "a Python aggregate's key values do not match its keys".to_owned(), + )); + } + let fields: Vec = key_values + .iter() + .map(|&k| { + input + .fields() + .get(k) + .map(|f| (**f).clone()) + .ok_or_else(|| internal(format!("key column {k} out of range"))) + }) + .collect::>()?; + Some(Arc::new(Schema::new(fields))) + } + _ => None, + }; + Ok(Arc::new(Self { + shape, + lanes: lanes.into(), + encoder, + store: SharedSink::new(Arc::clone(&schema), scope), + schema, + key_schema, + input, + max_rows: limits.max_rows, + max_bytes: limits.max_bytes, + window: limits.window.max(1), + })) + } + + /// The output schema. + #[must_use] + pub fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +/// Input rows cut into what a lane's task gets, not framed yet. +enum Cut { + /// Rows of one side (zero for a cogroup side without rows), with the + /// reservation of their bytes. + Rows { + side: u8, + rows: RecordBatch, + memory: Reservation, + }, + EndGroup, + End, +} + +/// One driver's state (and, after `combine`, what `finalize` finishes). +struct Driver { + lanes: Vec, + /// A finished batch the store could not take yet (it asked for + /// memory), with the memory of the results it came from. + ready: Option<(Batch, Vec)>, + partition: u32, + /// The key of the open group's last row seen. + current: Option, + /// A group is open (grouped shapes). + group_open: bool, + /// Rows of the open group (and side) not cut yet. + held: Vec, + held_rows: usize, + held_side: u8, + /// The cogroup sides the open group has rows on. + sides_seen: [bool; 2], + /// The reservation of the bytes of the input rows taken and not cut + /// yet: the held rows and the rest of the batch being fed. + held_bytes: Reservation, + /// The rows `held_bytes` covers: each cut takes its share by rows of + /// all of them, not of the held rows only (a batch's later groups keep + /// theirs while its first group is cut). + held_bytes_rows: usize, + /// Cut, waiting for room in the lanes' windows (shared with the + /// driver's `Blocked` future, which frames it as the windows free). + feed: Arc>, + /// Memory reserved for the next frame's copies (`finalize` waits for + /// it; `sink` asks the driver for it). + spare: Reservation, + /// In `finalize`, the bytes the next frame's copies wait for (0: none). + wants: usize, + /// An aggregate's groups waiting for their values. + pending: Pending, +} + +/// A driver's backlog, shared with its `Blocked` future. +struct Feed { + /// Cut, not framed yet. + backlog: VecDeque, + /// A failure of the `Blocked` future's framing: the next call's error. + failed: Option, +} + +/// What [`flush_feed`] stopped at. +#[derive(Clone, Copy)] +enum Flushed { + /// The backlog is empty. + All, + /// A lane's window is full. + Window, + /// The next frame's copies need this many more bytes. + Memory(usize), +} + +/// One driver's sink state. +struct RelLocal { + store: LocalState, + driver: Driver, + /// The driver's consumer, whose credit hands results to the store. + consumer: Arc, +} + +/// The global state: the store's, and what drivers left. +struct RelGlobal { + store: GlobalState, + leftover: Mutex>, +} + +/// A run of input rows within one group and side. +struct Run { + start: usize, + len: usize, + /// It starts a group (its first row's key differs from the open one). + new_group: bool, + side: u8, +} + +impl RelationSink { + /// The runs of `batch` and the key of its last row: rows of one group + /// and side, in order. + fn runs( + &self, + batch: &RecordBatch, + current: Option<&OwnedRow>, + ) -> Result<(Vec, Option)> { + let n = batch.num_rows(); + let sides: Option<&arrow_array::Int32Array> = match &self.shape { + RelationShape::CoGrouped { side, .. } => Some( + batch + .column(*side) + .as_any() + .downcast_ref() + .ok_or_else(|| internal("a cogroup side column is not INT".to_owned()))?, + ), + _ => None, + }; + let side_of = |i: usize| -> u8 { sides.map_or(0, |s| u8::from(s.value(i) != 0)) }; + let Some(encoder) = &self.encoder else { + let run = Run { + start: 0, + len: n, + new_group: false, + side: 0, + }; + return Ok((vec![run], None)); + }; + let rows = encoder.encode(batch)?; + let mut runs: Vec = Vec::new(); + for i in 0..n { + let new_group = if i == 0 { + current.is_none_or(|c| c.row() != rows.row(0)) + } else { + rows.row(i) != rows.row(i - 1) + }; + let side = side_of(i); + match runs.last_mut() { + Some(r) if !new_group && r.side == side => r.len += 1, + _ => runs.push(Run { + start: i, + len: 1, + new_group, + side, + }), + } + } + let last = n.checked_sub(1).map(|i| rows.row(i).owned()); + Ok((runs, last.or_else(|| current.cloned()))) + } + + /// The frame of `lane` over `rows` (one side): its columns evaluated, + /// in their client forms; and the bytes they hold beyond the rows' + /// buffers (a computed argument's result, a client form's copy). + fn frame( + lane: &RelationLane, + side: u8, + rows: &RecordBatch, + partition: u32, + ) -> Result<(Frame, usize)> { + let exprs = &lane.sides[usize::from(side)]; + let names = &lane.names[usize::from(side)]; + let mut fields = Vec::with_capacity(exprs.len()); + let mut columns = Vec::with_capacity(exprs.len()); + let mut copied = 0; + let mut input = Vec::new(); + for c in rows.columns() { + buffer_ranges(&c.to_data(), &mut input); + } + for (e, name) in exprs.iter().zip(names) { + let a = client_column(e.as_ref(), rows, partition)?; + copied += fresh_bytes(&a.to_data(), &input); + fields.push(Field::new(name, a.data_type().clone(), e.nullable())); + columns.push(a); + } + let options = RecordBatchOptions::new().with_row_count(Some(rows.num_rows())); + let batch = + RecordBatch::try_new_with_options(Arc::new(Schema::new(fields)), columns, &options) + .map_err(|e| internal(format!("Python UDF frame: {e}")))?; + Ok((Frame::Batch { side, batch }, copied)) + } + + /// Cuts the held rows into frames of at most `maxRecordsPerBatch` rows + /// (and `maxBytesPerBatch` bytes); with `all` the rows left over too + /// (the end of a group or side). + fn cut_held(&self, d: &mut Driver, all: bool) -> Result<()> { + if d.held_rows == 0 { + return Ok(()); + } + let held = std::mem::take(&mut d.held); + let rows = std::mem::take(&mut d.held_rows); + let data = if held.len() == 1 { + held.into_iter() + .next() + .ok_or_else(|| internal("no held rows".to_owned()))? + } else { + arrow_select::concat::concat_batches(&self.input, &held) + .map_err(|e| internal(format!("held Python UDF rows: {e}")))? + }; + let step = if self.max_rows == 0 { + rows + } else { + self.max_rows + }; + // The frames' columns as the tasks get them (computed arguments, + // client forms), sized as Spark's Arrow writer sizes them. + let sized = if self.max_bytes > 0 { + let mut columns = Vec::new(); + for lane in self.lanes.iter() { + for e in &lane.sides[usize::from(d.held_side)] { + columns.push(client_column(e.as_ref(), &data, d.partition)?); + } + } + Some(columns) + } else { + None + }; + // A map's rows travel as one struct column (its validity counts). + let wrapped = if matches!(self.shape, RelationShape::Map) { + self.lanes.len() + } else { + 0 + }; + let mut feed = lock(&d.feed); + let mut offset = 0; + while offset < rows { + let most = step.min(rows - offset); + let (len, capped) = match &sized { + Some(columns) => sized_rows(columns, wrapped, offset, most, self.max_bytes), + None => (most, false), + }; + let full = capped || len == step; + if !full && !all { + // Too few for a full frame: held for the next input. + d.held.push(data.slice(offset, len)); + d.held_rows = len; + break; + } + let share = d.held_bytes.size() * len / d.held_bytes_rows.max(len); + d.held_bytes_rows = d.held_bytes_rows.saturating_sub(len); + feed.backlog.push_back(Cut::Rows { + side: d.held_side, + rows: data.slice(offset, len), + memory: d.held_bytes.split(share), + }); + offset += len; + } + Ok(()) + } + + /// A zero-row frame of `side`: a cogroup key without rows on that side + /// (Spark gives the UDF an empty DataFrame for it). The side counts as + /// seen, so a key whose rows span input batches gets it once. + fn empty_side(&self, d: &mut Driver, side: u8) { + d.sides_seen[usize::from(side)] = true; + lock(&d.feed).backlog.push_back(Cut::Rows { + side, + rows: RecordBatch::new_empty(Arc::clone(&self.input)), + memory: d.held_bytes.split(0), + }); + } + + fn cogrouped(&self) -> bool { + matches!(self.shape, RelationShape::CoGrouped { .. }) + } + + /// Ends the open group: its held rows (a cogroup's empty right side), + /// then `EndGroup`. + fn end_group(&self, d: &mut Driver) -> Result<()> { + self.cut_held(d, true)?; + if self.cogrouped() && !d.sides_seen[1] { + self.empty_side(d, 1); + } + lock(&d.feed).backlog.push_back(Cut::EndGroup); + Ok(()) + } + + /// Cuts `batch` (reserved in `held_bytes`) into the backlog. + fn feed(&self, d: &mut Driver, batch: &RecordBatch, runs: Vec) -> Result<()> { + let grouped = self.encoder.is_some(); + for run in runs { + if run.new_group { + if d.group_open { + self.end_group(d)?; + } + d.group_open = true; + d.sides_seen = [false; 2]; + } else if grouped && !d.group_open { + return Err(internal("Python UDF rows before their group".to_owned())); + } + if run.side != d.held_side { + self.cut_held(d, true)?; + d.held_side = run.side; + } + if self.cogrouped() && run.side == 1 && !d.sides_seen[0] { + // A key on the right only: its empty left side first. + self.empty_side(d, 0); + } + d.sides_seen[usize::from(run.side)] = true; + d.held.push(batch.slice(run.start, run.len)); + d.held_rows += run.len; + self.cut_held(d, false)?; + } + Ok(()) + } + + /// Frames the driver's backlog as [`flush_feed`] does; the failure of + /// a `Blocked` future's framing first. + fn flush(&self, d: &mut Driver) -> Result { + let mut feed = lock(&d.feed); + if let Some(e) = feed.failed.take() { + return Err(e); + } + flush_feed( + &self.lanes, + self.window, + d.partition, + &d.lanes, + &mut feed.backlog, + &mut d.spare, + ) + } + + /// A task's batch as an output batch: its columns checked against the + /// declared client types, then moved to the canonical carriers. The + /// runtime checks the result against the declared schema first; this is + /// the engine's own guard. + fn output(&self, lane: &RelationLane, batch: &RecordBatch) -> Result { + if batch.num_columns() != lane.client.len() { + return Err(mismatch( + lane, + format!( + "{} columns where {} were declared", + batch.num_columns(), + lane.client.len() + ), + )); + } + let mut columns = Vec::with_capacity(batch.num_columns()); + for ((c, client), field) in batch + .columns() + .iter() + .zip(&lane.client) + .zip(self.schema.fields()) + { + columns.push(recarried(lane, c, client, field.data_type())?); + } + let options = RecordBatchOptions::new().with_row_count(Some(batch.num_rows())); + let out = RecordBatch::try_new_with_options(Arc::clone(&self.schema), columns, &options) + .map_err(|e| mismatch(lane, format!("{e}")))?; + Ok(Batch::new(out)) + } + + /// The next finished output batch of a driver, with the memory of the + /// results it came from: a lane's next result (map and grouped shapes) + /// or the groups every lane has a value for (aggregate). + fn next_output(&self, d: &mut Driver) -> Result)>> { + if let (RelationShape::Aggregate { .. }, Some(key_schema)) = (&self.shape, &self.key_schema) + { + return pair::joined( + &mut d.pending, + key_schema, + &self.schema, + &self.lanes, + &d.lanes, + ); + } + for (lane, state) in self.lanes.iter().zip(&d.lanes) { + let next = lock(&state.io).pop_output(); + if let Some((batch, memory)) = next { + return Ok(Some((self.output(lane, &batch)?, vec![memory]))); + } + } + Ok(None) + } + + /// Stores `batch` in the list, its results' bytes (`memory`) handed to + /// the store's reservation; or keeps it, reserved again, for the retry + /// after the memory the store asked for. + fn store( + &self, + l: &mut RelLocal, + batch: Batch, + memory: Vec, + g: &RelGlobal, + ) -> Result> { + // As the store keeps it, so the bytes lent are the bytes it takes. + let batch = crate::operators::owned_rows(batch)?; + let lent = lend(memory, &l.consumer, batch.data().get_array_memory_size()); + match self.store.sink(batch.clone(), &mut l.store, &g.store)? { + SinkStatus::NeedMemory { bytes } => { + let back = l.consumer.reservation_from_credit(lent); + l.driver.ready = Some((batch, vec![back])); + Ok(Some(SinkStatus::NeedMemory { + bytes: bytes.saturating_sub(lent).max(1), + })) + } + _ => Ok(None), + } + } + + /// Moves each lane's frames and results without waiting; the first + /// failure of a lane. + fn progress(d: &Driver) -> Result<()> { + let mut cx = Context::from_waker(noop_waker_ref()); + for lane in &d.lanes { + if let Some(Err(e)) = lane.poll(&mut cx) { + return Err(e); + } + } + Ok(()) + } + + /// Stores what the driver's tasks yielded, reserving refused results + /// first; a status to return if the pool or the store asked for memory. + fn store_results(&self, l: &mut RelLocal, g: &RelGlobal) -> Result> { + if let Some((ready, memory)) = l.driver.ready.take() + && let Some(status) = self.store(l, ready, memory, g)? + { + return Ok(Some(status)); + } + Self::progress(&l.driver)?; + while let Some((batch, memory)) = self.next_output(&mut l.driver)? { + if let Some(status) = self.store(l, batch, memory, g)? { + return Ok(Some(status)); + } + } + for lane in &l.driver.lanes { + if let Some(bytes) = lane.try_take_refused() { + return Ok(Some(SinkStatus::NeedMemory { bytes })); + } + } + // The pumps read on past a result just reserved. + Self::progress(&l.driver)?; + Ok(None) + } + + /// Takes `input` into the backlog: `NeedMemory` (nothing changed) if + /// the pool refuses its bytes. + fn take_input(&self, d: &mut Driver, input: Batch) -> Result> { + if !input.is_unweighted() { + // A retracting row sent to Python would come back as an insert. + return Err(krishiv_common::error::KrishivUnsupportedFeature { + feature: "a weighted (change-stream) input to a Python UDF relation".to_owned(), + } + .into()); + } + let input = crate::operators::owned_rows(input)?; + let data = input.data(); + let (runs, last) = self.runs(data, d.current.as_ref())?; + let keys = match (&self.shape, &self.key_schema) { + (RelationShape::Aggregate { key_values, .. }, Some(schema)) => { + let starts: Vec = runs + .iter() + .filter(|r| r.new_group) + .map(|r| u32::try_from(r.start)) + .collect::>() + .map_err(|_| internal("a batch over u32 rows".to_owned()))?; + pair::key_rows(data, key_values, schema, &starts)? + } + _ => None, + }; + let key_bytes = keys.as_ref().map_or(0, RecordBatch::get_array_memory_size); + let bytes = data.get_array_memory_size() + key_bytes; + if d.held_bytes.try_grow(bytes).is_err() { + return Ok(Some(SinkStatus::NeedMemory { bytes })); + } + if let Some(keys) = keys { + let memory = d.held_bytes.split(key_bytes); + d.pending.push(keys, memory); + } + d.held_bytes_rows += data.num_rows(); + self.feed(d, data, runs)?; + d.current = last; + Ok(None) + } +} + +/// A result that does not match what the UDF declared. +pub(super) fn mismatch(lane: &RelationLane, reason: String) -> KrishivError { + KrishivPythonUdfResultMismatch { + name: lane.name.clone(), + reason, + } + .into() +} + +/// `c` checked to be `client`, moved to the canonical `canonical`. +pub(super) fn recarried( + lane: &RelationLane, + c: &ArrayRef, + client: &ArrowType, + canonical: &ArrowType, +) -> Result { + if c.data_type() != client { + return Err(mismatch( + lane, + format!("{} where {client} was declared", c.data_type()), + )); + } + crate::carrier::recarried_array(c, canonical) +} + +/// Lends the bytes `parts` hold (results being stored) to `handle`'s grant +/// credit, up to `bytes`, so the reservation that stores the same rows takes +/// them without asking the pool again; the rest is freed. The bytes lent: +/// what to take back with `reservation_from_credit` if that reservation is +/// refused. +fn lend(parts: Vec, handle: &Arc, bytes: usize) -> usize { + let mut lent = 0; + for mut part in parts { + if part.consumer() != handle.id() && part.transfer_to(handle).is_err() { + // Freed as it drops. + continue; + } + let take = part.size().min(bytes - lent); + part.shrink_to_credit(take); + lent += take; + } + lent +} + +/// Frames `backlog` into `lanes` while every lane's window has room, the +/// copies reserved in `spare`. Nothing changes for a frame whose copies the +/// pool refuses. +fn flush_feed( + specs: &[RelationLane], + window: usize, + partition: u32, + lanes: &[Lane], + backlog: &mut VecDeque, + spare: &mut Reservation, +) -> Result { + loop { + if lanes.iter().any(|l| lock(&l.io).in_flight() >= window) { + return Ok(Flushed::Window); + } + let Some(cut) = backlog.front() else { + return Ok(Flushed::All); + }; + match cut { + Cut::Rows { side, rows, .. } => { + let mut frames = Vec::with_capacity(specs.len()); + for lane in specs { + frames.push(RelationSink::frame(lane, *side, rows, partition)?); + } + let copied: usize = frames.iter().map(|(_, c)| c).sum(); + let more = copied.saturating_sub(spare.size()); + if more > 0 && spare.try_grow(more).is_err() { + return Ok(Flushed::Memory(more)); + } + let Some(Cut::Rows { memory, .. }) = backlog.pop_front() else { + return Err(internal("the Python UDF backlog changed".to_owned())); + }; + let shared = Arc::new(Mutex::new(memory)); + for ((lane, state), (frame, copied)) in specs.iter().zip(lanes).zip(frames) { + let memory = FrameMem { + _shared: Some(Arc::clone(&shared)), + _own: spare.split(copied), + }; + state.queue(lane.udf.as_ref(), frame, memory)?; + } + } + Cut::EndGroup | Cut::End => { + let end = matches!(cut, Cut::End); + backlog.pop_front(); + for (lane, state) in specs.iter().zip(lanes) { + // A task never opened (no rows) gets no `End`. + if end && !state.opened() { + continue; + } + let frame = if end { Frame::End } else { Frame::EndGroup }; + let memory = FrameMem { + _shared: None, + _own: spare.split(0), + }; + state.queue(lane.udf.as_ref(), frame, memory)?; + } + } + } + } +} + +/// `e` over `rows`, in its client form. +fn client_column(e: &dyn ScalarExpr, rows: &RecordBatch, partition: u32) -> Result { + let a = e.evaluate_in(rows, partition)?; + let client = krishiv_types::arrow::client_arrow_type(a.data_type()); + if a.data_type() == &client { + Ok(a) + } else { + crate::carrier::recarried_array(&a, &client) + } +} + +/// The rows of the frame starting at row `start` of `columns`, at most +/// `most`, under a `maxBytesPerBatch` of `cap`; and whether the cap closed +/// it. +/// +/// Spark writes rows into the frame while its Arrow size is under the cap +/// and closes it once a row brings it to the cap or past +/// (`BatchedPythonArrowInput.writeSizedBatch`, PythonArrowInput.scala:226-233 +/// at v4.2.0): the frame is the fewest rows whose size reaches the cap, at +/// least one, so the row that crosses the cap is the frame's last. The size +/// grows with the rows, so the cut is a binary search over the columns' +/// offsets, never a copy. +fn sized_rows( + columns: &[ArrayRef], + wrapped: usize, + start: usize, + most: usize, + cap: usize, +) -> (usize, bool) { + let size = |rows: usize| -> usize { + let wrap = wrapped * validity_bytes(rows); + columns.iter().fold(wrap, |n, c| { + n.saturating_add(arrow_size(c.as_ref(), start, rows)) + }) + }; + if size(most) < cap { + return (most, false); + } + // The fewest rows in `1..=most` whose size reaches the cap. + let (mut lo, mut hi) = (1, most); + while lo < hi { + let mid = lo + (hi - lo) / 2; + if size(mid) >= cap { + hi = mid; + } else { + lo = mid + 1; + } + } + (lo, true) +} + +/// The bytes of a validity bitmap of `rows` rows. +fn validity_bytes(rows: usize) -> usize { + rows.div_ceil(8) +} + +/// The size Spark's Arrow writer gives rows `start..start + rows` of `a` +/// (`ArrowFieldWriter.getSizeInBytes`, ArrowWriter.scala:183-187 at +/// v4.2.0: Arrow Java's `getBufferSizeFor`): validity, offsets and values, +/// a validity bitmap for every vector, `rows + 1` offsets, the values the +/// offsets span; nothing for zero rows. +fn arrow_size(a: &dyn Array, start: usize, rows: usize) -> usize { + use arrow_array::cast::AsArray; + if rows == 0 { + return 0; + } + let validity = validity_bytes(rows); + // Offsets, then the child rows they span. + let nested = |width: usize, (from, len): (usize, usize), child: &dyn Array| { + validity + (rows + 1) * width + arrow_size(child, from, len) + }; + match a.data_type() { + ArrowType::Null => 0, + ArrowType::Boolean => 2 * validity, + ArrowType::Utf8 => { + validity + (rows + 1) * 4 + span(a.as_string::().value_offsets(), start, rows).1 + } + ArrowType::Binary => { + validity + (rows + 1) * 4 + span(a.as_binary::().value_offsets(), start, rows).1 + } + ArrowType::LargeUtf8 => { + validity + (rows + 1) * 8 + span(a.as_string::().value_offsets(), start, rows).1 + } + ArrowType::LargeBinary => { + validity + (rows + 1) * 8 + span(a.as_binary::().value_offsets(), start, rows).1 + } + ArrowType::FixedSizeBinary(width) => validity + rows * usize::try_from(*width).unwrap_or(0), + ArrowType::List(_) => { + let l = a.as_list::(); + nested(4, span(l.value_offsets(), start, rows), l.values().as_ref()) + } + ArrowType::LargeList(_) => { + let l = a.as_list::(); + nested(8, span(l.value_offsets(), start, rows), l.values().as_ref()) + } + ArrowType::Map(..) => { + let m = a.as_map(); + nested(4, span(m.value_offsets(), start, rows), m.entries()) + } + ArrowType::FixedSizeList(_, width) => { + let width = usize::try_from(*width).unwrap_or(0); + let l = a.as_fixed_size_list(); + validity + arrow_size(l.values().as_ref(), start * width, rows * width) + } + ArrowType::Struct(_) => a.as_struct().columns().iter().fold(validity, |n, c| { + n.saturating_add(arrow_size(c.as_ref(), start, rows)) + }), + t => match t.primitive_width() { + Some(width) => validity + rows * width, + // A form Spark's writer has no vector for: the rows' share of + // the array's buffers. + None => a + .get_buffer_memory_size() + .saturating_mul(rows) + .div_ceil(a.len().max(1)), + }, + } +} + +/// The first child row and the count of child rows (or bytes) that rows +/// `start..start + rows` span in `offsets`. +fn span( + offsets: &[O], + start: usize, + rows: usize, +) -> (usize, usize) { + let at = |i: usize| offsets.get(i).map_or(0, |o| o.as_usize()); + let from = at(start); + (from, at(start + rows).saturating_sub(from)) +} + +/// The address ranges of `data`'s buffers (validity and children too). +fn buffer_ranges(data: &arrow_data::ArrayData, out: &mut Vec<(usize, usize)>) { + let mut add = |b: &arrow_buffer::Buffer| { + let start = b.as_ptr() as usize; + out.push((start, start + b.len())); + }; + if let Some(nulls) = data.nulls() { + add(nulls.buffer()); + } + for b in data.buffers() { + add(b); + } + for c in data.child_data() { + buffer_ranges(c, out); + } +} + +/// The bytes of `data`'s buffers that lie outside every range of `input`: +/// what evaluating a frame column allocated (a column read shares the +/// input's buffers, a computed argument does not). +fn fresh_bytes(data: &arrow_data::ArrayData, input: &[(usize, usize)]) -> usize { + let size = |b: &arrow_buffer::Buffer| { + let at = b.as_ptr() as usize; + if input.iter().any(|&(start, end)| start <= at && at < end) { + 0 + } else { + b.capacity() + } + }; + let mut bytes = data.nulls().map_or(0, |n| size(n.buffer())); + bytes += data.buffers().iter().map(size).sum::(); + bytes + + data + .child_data() + .iter() + .map(|c| fresh_bytes(c, input)) + .sum::() +} + +/// The future of a `Blocked` driver. It frames the driver's backlog as the +/// windows free (the driver takes no input before the backlog drained, so +/// the cuts held stay within about one input batch), and is ready once the +/// backlog is empty and every lane's window has room; or early when a +/// lane ended (its failure is the next call's), a result waits for memory, +/// a frame's copies need memory or framing failed (the next call asks for +/// the memory or returns the failure). +fn blocked( + specs: Arc<[RelationLane]>, + lanes: Vec, + feed: Arc>, + mut spare: Reservation, + window: usize, + partition: u32, +) -> BoxFuture<'static, ()> { + Box::pin(poll_fn(move |cx| { + loop { + for lane in &lanes { + if lane.poll(cx).is_some() || lock(&lane.io).refused.is_some() { + return Poll::Ready(()); + } + } + let mut f = lock(&feed); + if f.failed.is_some() { + return Poll::Ready(()); + } + let before = f.backlog.len(); + match flush_feed( + &specs, + window, + partition, + &lanes, + &mut f.backlog, + &mut spare, + ) { + Err(e) => { + f.failed = Some(e); + return Poll::Ready(()); + } + Ok(Flushed::Memory(_)) => return Poll::Ready(()), + Ok(Flushed::All | Flushed::Window) => { + let (moved, empty) = (f.backlog.len() < before, f.backlog.is_empty()); + drop(f); + if moved { + // The pumps take the new frames, then frame again. + continue; + } + let room = lanes.iter().all(|l| lock(&l.io).in_flight() < window); + return if empty && room { + Poll::Ready(()) + } else { + Poll::Pending + }; + } + } + } + })) +} + +impl Sink for RelationSink { + fn name(&self) -> &'static str { + "PythonRelation" + } + + fn global_state(&self, ctx: &TaskContext) -> Result { + Ok(Arc::new(RelGlobal { + store: self.store.global_state(ctx)?, + leftover: Mutex::new(Vec::new()), + })) + } + + fn local_state(&self, ctx: &TaskContext, global: &GlobalState) -> Result { + let g = global_ref::(global, "PythonRelation")?; + let consumer = ctx.consumer().consumer(); + Ok(Box::new(RelLocal { + consumer: Arc::clone(consumer), + store: self.store.local_state(ctx, &g.store)?, + driver: Driver { + lanes: Lane::group( + self.lanes.iter().map(|_| consumer.reservation()), + matches!(self.shape, RelationShape::Aggregate { .. }), + ), + ready: None, + partition: ctx.partition(), + current: None, + group_open: false, + held: Vec::new(), + held_rows: 0, + held_side: 0, + sides_seen: [false; 2], + held_bytes: consumer.reservation(), + held_bytes_rows: 0, + feed: Arc::new(Mutex::new(Feed { + backlog: VecDeque::new(), + failed: None, + })), + spare: consumer.reservation(), + wants: 0, + pending: Pending::default(), + }, + })) + } + + fn sink( + &self, + input: Batch, + local: &mut LocalState, + global: &GlobalState, + ) -> Result { + let g = global_ref::(global, "PythonRelation")?; + let l = local_mut::(local, "PythonRelation")?; + // Results and the backlog first: a `NeedMemory` here leaves the + // input untaken, so the retry with the same batch is exact. + if let Some(status) = self.store_results(l, g)? { + return Ok(status); + } + if let Flushed::Memory(bytes) = self.flush(&mut l.driver)? { + return Ok(SinkStatus::NeedMemory { bytes }); + } + if input.num_rows() > 0 { + if let Some(status) = self.take_input(&mut l.driver, input)? { + return Ok(status); + } + // A refusal here waits for the next call, which asks for it. + self.flush(&mut l.driver)?; + Self::progress(&l.driver)?; + } + // Blocked while cuts wait or a window is full: the future frames + // the backlog, so no more input is taken before it drained. + let d = &mut l.driver; + if !lock(&d.feed).backlog.is_empty() + || d.lanes + .iter() + .any(|lane| lock(&lane.io).in_flight() >= self.window) + { + return Ok(SinkStatus::Blocked(blocked( + Arc::clone(&self.lanes), + d.lanes.clone(), + Arc::clone(&d.feed), + d.spare.split(0), + self.window, + d.partition, + ))); + } + Ok(SinkStatus::NeedMoreInput) + } + + fn combine(&self, local: LocalState, global: &GlobalState) -> Result<()> { + let g = global_ref::(global, "PythonRelation")?; + let l = local + .downcast::() + .map_err(|_| internal("wrong local state for PythonRelation".to_owned()))?; + let RelLocal { + store, mut driver, .. + } = *l; + // The input ended: the open group ends, then the partition. What the + // windows do not take now, `finalize` frames. + if driver.group_open { + self.end_group(&mut driver)?; + } else { + self.cut_held(&mut driver, true)?; + } + lock(&driver.feed).backlog.push_back(Cut::End); + self.flush(&mut driver)?; + Self::progress(&driver)?; + lock(&g.leftover).push(driver); + self.store.combine(store, &g.store) + } + + fn finalize<'a>( + &'a self, + global: &'a GlobalState, + ctx: &'a TaskContext, + ) -> BoxFuture<'a, Result<()>> { + Box::pin(async move { + let g = global_ref::(global, "PythonRelation")?; + let mut drivers = std::mem::take(&mut *lock(&g.leftover)); + // Every task at once; results are stored as they arrive. The + // first failure (or a cancel) ends the phase, and the error + // drops the other tasks, which cancels them. + loop { + ctx.query().check_cancelled()?; + let mut done = true; + for d in &mut drivers { + done &= self.finish_step(d, g, ctx).await?; + } + if done { + break; + } + // A frame waiting for its copies' memory: every driver's + // results are stored by now (written out where a frame + // waits), so only the pool can give it. + if let Some(d) = drivers.iter_mut().find(|d| d.wants > 0) { + let bytes = std::mem::take(&mut d.wants); + d.spare.grow(bytes).await.map_err(KrishivError::from)?; + continue; + } + self.wait(&drivers, ctx).await?; + } + if matches!(self.shape, RelationShape::Aggregate { .. }) { + for d in &drivers { + pair::check_complete(&d.pending, &self.lanes, &d.lanes)?; + } + } + drop(drivers); + self.store.finalize(&g.store, ctx).await + }) + } + + fn spill<'a>( + &'a self, + local: &'a mut LocalState, + target: usize, + ctx: &'a TaskContext, + ) -> BoxFuture<'a, Result> { + Box::pin(async move { + let l = local_mut::(local, "PythonRelation")?; + self.store.spill(&mut l.store, target, ctx).await + }) + } +} + +impl RelationSink { + /// One step of a finished driver in `finalize`: the backlog framed as + /// far as the windows and the pool allow, the results stored, a refused + /// map or grouped result with them (written to a spill file when the + /// pool has no room, or when a frame waits for its copies' memory: the + /// results' memory is what it can get), the backlog framed again, then + /// an aggregate's refused values reserved if the pool has room now. A + /// frame still refused leaves its bytes in `wants`, which `finalize` + /// waits for once every driver stored its results. Whether the driver + /// is done: no backlog, every task ended, every result stored (an + /// aggregate's values left over are `check_complete`'s error). + async fn finish_step(&self, d: &mut Driver, g: &RelGlobal, ctx: &TaskContext) -> Result { + d.wants = wanted(self.flush(d)?); + let mut batches = Vec::new(); + let mut held = Vec::new(); + if let Some((batch, memory)) = d.ready.take() { + batches.push(batch); + held.extend(memory); + } + while let Some((batch, memory)) = self.next_output(d)? { + batches.push(batch); + held.extend(memory); + } + let aggregate = matches!(self.shape, RelationShape::Aggregate { .. }); + if !aggregate { + // After the lane's reserved results, in its order. + for (lane, state) in self.lanes.iter().zip(&d.lanes) { + if let Some(batch) = state.take_refused_unreserved() { + batches.push(self.output(lane, &batch)?); + } + } + } + if !batches.is_empty() { + let bytes = batches + .iter() + .map(|b| b.data().get_array_memory_size()) + .sum(); + // Kept in memory if the pool has room, else written out: the + // list can spill only once the store finalized. + let handle = ctx.consumer().consumer(); + let lent = lend(std::mem::take(&mut held), handle, bytes); + let mut reservation = ctx.consumer().reservation(); + let free = d.wants > 0 && self.store.spills(); + if !free && reservation.try_grow(bytes).is_ok() { + self.store.add_part(&g.store, batches, reservation)?; + } else { + drop((reservation, handle.reservation_from_credit(lent))); + self.store.add_spilled_part(&g.store, batches, ctx).await?; + } + } + if d.wants > 0 { + d.wants = wanted(self.flush(d)?); + } + if aggregate { + // The stored results' memory is free again; a value still + // refused waits in `wait` while the other lanes move. + for lane in &d.lanes { + let _missing = lane.try_take_refused(); + } + } + Ok(lock(&d.feed).backlog.is_empty() + && d.ready.is_none() + && d.lanes + .iter() + .all(|l| l.finished() && (aggregate || !lock(&l.io).has_results()))) + } + + /// Waits until a driver can move: a lane yielded (for an aggregate: + /// every lane, as `pair::joined` needs, or a refused value got its + /// memory), failed, ended or freed room in its window for a backlog; or + /// the query was cancelled. + async fn wait(&self, drivers: &[Driver], ctx: &TaskContext) -> Result<()> { + let aggregate = matches!(self.shape, RelationShape::Aggregate { .. }); + let mut reserving: Vec>> = if aggregate { + drivers + .iter() + .flat_map(|d| d.lanes.iter().filter_map(Lane::reserve_refused)) + .collect() + } else { + Vec::new() + }; + let mut cancelled = std::pin::pin!(ctx.query().cancelled()); + poll_fn(|cx| { + if cancelled.as_mut().poll(cx).is_ready() { + return Poll::Ready(Err(ctx.query().cancellation())); + } + for f in &mut reserving { + if let Poll::Ready(r) = f.as_mut().poll(cx) { + return Poll::Ready(r); + } + } + let mut ready = false; + let mut all_done = true; + for d in drivers { + let mut room = true; + let mut joinable = aggregate; + for lane in &d.lanes { + if let Some(Err(e)) = lane.poll(cx) { + return Poll::Ready(Err(e)); + } + all_done &= lane.finished(); + let q = lock(&lane.io); + if aggregate { + // Only rows can be joined: a lane holding nothing but + // empty batches would make `finalize` spin. + joinable &= q.outputs.iter().any(|(b, _)| b.num_rows() > 0); + } else { + ready |= q.has_results(); + } + room &= q.in_flight() < self.window; + } + let backlog = !lock(&d.feed).backlog.is_empty(); + ready |= joinable || (room && backlog); + all_done &= !backlog && d.ready.is_none(); + } + if ready || all_done { + Poll::Ready(Ok(())) + } else { + Poll::Pending + } + }) + .await + } +} + +/// The bytes a [`Flushed`] frame waits for (0: none). +fn wanted(flushed: Flushed) -> usize { + match flushed { + Flushed::Memory(bytes) => bytes, + Flushed::All | Flushed::Window => 0, + } +} + +/// Per lane the frames in flight, and the cuts in the backlog, of a +/// driver's state. +#[cfg(test)] +pub(super) fn queued(local: &mut LocalState) -> Result<(Vec, usize)> { + let l = local_mut::(local, "PythonRelation")?; + let d = &l.driver; + let lanes = d.lanes.iter().map(|l| lock(&l.io).in_flight()).collect(); + let backlog = lock(&d.feed).backlog.len(); + Ok((lanes, backlog)) +} + +/// The global state a relation node's reader takes its list from: the +/// store's own; `None` for a state of another sink. +pub(crate) fn store_state(global: &GlobalState) -> Option { + global + .downcast_ref::() + .map(|g| Arc::clone(&g.store)) +} + +/// A column of the input, as a frame reads it. +#[derive(Debug)] +pub struct InputColumn { + /// The column. + pub index: usize, + /// Its type. + pub data_type: ArrowType, + /// Whether it may hold nulls. + pub nullable: bool, +} + +impl ScalarExpr for InputColumn { + fn data_type(&self) -> ArrowType { + self.data_type.clone() + } + + fn nullable(&self) -> bool { + self.nullable + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + batch + .columns() + .get(self.index) + .cloned() + .ok_or_else(|| internal(format!("input column {} out of range", self.index))) + } +} diff --git a/crates/krishiv-exec/src/python_relation/tests.rs b/crates/krishiv-exec/src/python_relation/tests.rs new file mode 100644 index 000000000..dcc097fb5 --- /dev/null +++ b/crates/krishiv-exec/src/python_relation/tests.rs @@ -0,0 +1,1711 @@ +// SPDX-License-Identifier: Apache-2.0 +//! The relation sink driven by hand over a fake runtime whose tasks log the +//! frames they see (design P3-T085 §4): batching, group cuts, cogroup sides, +//! the bounded window, failures, cancellation and memory. +#![expect( + clippy::disallowed_methods, + reason = "test-only runtime plumbing; never runs on a compute worker" +)] + +use std::collections::VecDeque; +use std::future::poll_fn; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::task::Poll; +use std::time::Duration; + +use arrow_array::{ArrayRef, Int32Array, Int64Array, RecordBatch}; +use arrow_schema::{DataType as ArrowType, Field, Schema, SchemaRef}; +use futures::future::BoxFuture; +use futures::task::AtomicWaker; +use krishiv_common::error::{ErrorClass, InternalError}; +use krishiv_memory::{MemoryConfig, MemoryPool, SpillOwner, TreePool}; +use krishiv_types::batch::Batch; + +use super::sink::{DEFAULT_WINDOW, InputColumn, RelationLimits}; +use super::{Frame, PreparedRelationUdf, RelationLane, RelationShape, RelationSink, UdfTask}; +use crate::operator::{Sink, SinkStatus}; +use crate::operators::ScalarExpr; +use crate::{ + ExecConfig, Executor, OperatorId, QueryContext, QueryId, TaskContext, TokioClock, + TokioExecutor, WorkloadClass, +}; + +/// What the fake's tasks share with the test. +#[derive(Default)] +struct Shared { + /// Every frame every task saw: `B:`, `G`, `E`. + log: std::sync::Mutex>, + /// Tasks opened. + opened: std::sync::atomic::AtomicUsize, + /// While closed, `send` waits. + closed: AtomicBool, + /// While set, `send` waits once the log has this many frames. + admit: std::sync::Mutex>, + gate: AtomicWaker, + /// Fail the n-th frame (counted from 1). + fail_at: std::sync::Mutex>, + /// Tasks dropped (cancelled or finished). + dropped: std::sync::atomic::AtomicUsize, + /// Result batches the operator read from the tasks. + taken: std::sync::atomic::AtomicUsize, + /// Polls of `next`: past [`SPUN`] the task fails, so a busy loop ends. + polls: std::sync::atomic::AtomicUsize, + /// While set, `send` waits while this many results are not read (a + /// runtime whose output pipe is full takes no more frames). + out_cap: std::sync::Mutex>, + /// While set, `send` wakes `next` only when the frame gave a result or + /// ended the task (a runtime that signals its reader only on output). + quiet: AtomicBool, +} + +/// Polls of a task's `next` that only a busy loop reaches. +const SPUN: usize = 100_000; + +impl Shared { + fn open_gate(&self) { + self.closed.store(false, Ordering::SeqCst); + self.gate.wake(); + } + + fn log(&self) -> Vec { + self.log.lock().unwrap().clone() + } +} + +/// What a fake task yields. +#[derive(Clone, Copy)] +enum Yield { + /// Each batch back, twice (output not aligned to input). + Twice, + /// Nothing at all. + Nothing, + /// One row per group: its row count (`BIGINT`). + Count, + /// Two rows per group (an aggregate UDF yielding too many values). + CountTwice, + /// Each batch back this many times. + Burst(usize), + /// As `Count`, the values of this many groups yielded together (a + /// runtime batching its results across groups); the rest at the end. + Coalesce(usize), +} + +struct Fake { + shared: Arc, + yields: Yield, +} + +impl PreparedRelationUdf for Fake { + fn open(&self) -> crate::error::Result> { + self.shared.opened.fetch_add(1, Ordering::SeqCst); + Ok(Box::new(Task { + shared: Arc::clone(&self.shared), + yields: self.yields, + state: std::sync::Mutex::new(TaskState::default()), + ready: AtomicWaker::new(), + drained: AtomicWaker::new(), + })) + } +} + +#[derive(Default)] +struct TaskState { + out: VecDeque, + rows: i64, + ended: bool, + /// `Coalesce`: the values not yielded yet. + held: Vec, +} + +struct Task { + shared: Arc, + yields: Yield, + state: std::sync::Mutex, + ready: AtomicWaker, + /// Woken as `next` takes a result (`Shared::out_cap`). + drained: AtomicWaker, +} + +impl Drop for Task { + fn drop(&mut self) { + self.shared.dropped.fetch_add(1, Ordering::SeqCst); + } +} + +fn count(n: i64) -> RecordBatch { + counts(vec![n]) +} + +fn counts(n: Vec) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("n", ArrowType::Int64, true)])); + RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(n)) as ArrayRef]).unwrap() +} + +impl UdfTask for Task { + fn send(&self, frame: Frame) -> BoxFuture<'_, crate::error::Result<()>> { + Box::pin(async move { + poll_fn(|cx| { + self.drained.register(cx.waker()); + let cap = *self.shared.out_cap.lock().unwrap(); + if cap.is_some_and(|c| self.state.lock().unwrap().out.len() >= c) { + return Poll::Pending; + } + self.shared.gate.register(cx.waker()); + let full = (*self.shared.admit.lock().unwrap()) + .is_some_and(|n| self.shared.log.lock().unwrap().len() >= n); + if full || self.shared.closed.load(Ordering::SeqCst) { + Poll::Pending + } else { + Poll::Ready(()) + } + }) + .await; + let n = { + let mut log = self.shared.log.lock().unwrap(); + log.push(match &frame { + Frame::Batch { side, batch } => format!("B{side}:{}", batch.num_rows()), + Frame::EndGroup => "G".to_owned(), + Frame::End => "E".to_owned(), + }); + log.len() + }; + if *self.shared.fail_at.lock().unwrap() == Some(n) { + return Err(InternalError { + message: "the UDF raised".to_owned(), + } + .into()); + } + let mut s = self.state.lock().unwrap(); + let before = s.out.len(); + match frame { + Frame::Batch { batch, .. } => { + s.rows += i64::try_from(batch.num_rows()).unwrap(); + let copies = match self.yields { + Yield::Twice => 2, + Yield::Burst(n) => n, + Yield::Nothing | Yield::Count | Yield::CountTwice | Yield::Coalesce(_) => 0, + }; + for _ in 0..copies { + s.out.push_back(batch.clone()); + } + } + Frame::EndGroup => { + let n = std::mem::take(&mut s.rows); + match self.yields { + Yield::Count => s.out.push_back(count(n)), + Yield::CountTwice => { + s.out.push_back(count(n)); + s.out.push_back(count(n)); + } + Yield::Coalesce(groups) => { + s.held.push(n); + if s.held.len() >= groups { + let values = std::mem::take(&mut s.held); + s.out.push_back(counts(values)); + } + } + _ => {} + } + } + Frame::End => { + if !s.held.is_empty() { + let values = std::mem::take(&mut s.held); + s.out.push_back(counts(values)); + } + s.ended = true; + } + } + let output = s.out.len() != before || s.ended; + drop(s); + if output || !self.shared.quiet.load(Ordering::SeqCst) { + self.ready.wake(); + } + Ok(()) + }) + } + + fn next(&self) -> BoxFuture<'_, crate::error::Result>> { + Box::pin(poll_fn(|cx| { + if self.shared.polls.fetch_add(1, Ordering::SeqCst) > SPUN { + return Poll::Ready(Err(InternalError { + message: "the operator spun".to_owned(), + } + .into())); + } + self.ready.register(cx.waker()); + let mut s = self.state.lock().unwrap(); + match s.out.pop_front() { + Some(b) => { + self.shared.taken.fetch_add(1, Ordering::SeqCst); + self.drained.wake(); + Poll::Ready(Ok(Some(b))) + } + None if s.ended => Poll::Ready(Ok(None)), + None => Poll::Pending, + } + })) + } +} + +fn pool(limit: usize) -> Arc { + TreePool::with_clock( + &MemoryConfig { + limit: Some(limit), + quantum: 1 << 10, + spill_wait_timeout: Duration::from_secs(5), + abort_release_timeout: Duration::from_secs(30), + spill_scratch_reserve: krishiv_memory::DEFAULT_SPILL_SCRATCH_RESERVE, + resident_guard: true, + }, + Arc::new(TokioClock), + ) +} + +/// `k INT, v BIGINT`; for a cogroup also `side INT`. +fn input(cogroup: bool) -> SchemaRef { + let mut fields = vec![ + Field::new("k", ArrowType::Int32, true), + Field::new("v", ArrowType::Int64, true), + ]; + if cogroup { + fields.push(Field::new("side", ArrowType::Int32, false)); + } + Arc::new(Schema::new(fields)) +} + +fn column(index: usize, data_type: ArrowType) -> Arc { + Arc::new(InputColumn { + index, + data_type, + nullable: true, + }) +} + +struct Rig { + sink: Arc, + root: Arc, + query: Arc, + ctx: TaskContext, + shared: Arc, +} + +fn rig(shape: RelationShape, yields: Yield, limits: RelationLimits, limit: usize) -> Rig { + rig_in(shape, yields, limits, limit, None) +} + +/// As [`rig`], every frame of the map lane being `columns` (the task yields +/// nothing). +fn rig_columns( + limits: RelationLimits, + limit: usize, + columns: &[(Arc, &str)], +) -> Rig { + let r = rig(RelationShape::Map, Yield::Nothing, limits, limit); + let lane = RelationLane { + name: "f".to_owned(), + udf: Arc::new(Fake { + shared: Arc::clone(&r.shared), + yields: Yield::Nothing, + }), + sides: vec![columns.iter().map(|(e, _)| Arc::clone(e)).collect()], + names: vec![columns.iter().map(|(_, n)| (*n).to_owned()).collect()], + client: vec![ArrowType::Int32, ArrowType::Int64], + }; + let sink = RelationSink::try_new( + input(false), + RelationShape::Map, + vec![lane], + r.sink.schema(), + limits, + None, + ) + .unwrap(); + Rig { sink, ..r } +} + +/// A computed argument (as `repeat(s, n)`): a new string of `width` bytes +/// per row, already in its client form. +#[derive(Debug)] +struct Wide { + width: usize, +} + +impl ScalarExpr for Wide { + fn data_type(&self) -> ArrowType { + ArrowType::Utf8 + } + + fn nullable(&self) -> bool { + false + } + + fn evaluate(&self, batch: &RecordBatch) -> crate::error::Result { + let s = "x".repeat(self.width); + let a: arrow_array::StringArray = (0..batch.num_rows()).map(|_| Some(s.as_str())).collect(); + Ok(Arc::new(a)) + } +} + +/// As [`rig`], its finished rows spilling to `scope`. +fn rig_in( + shape: RelationShape, + yields: Yield, + limits: RelationLimits, + limit: usize, + scope: Option, +) -> Rig { + let shared = Arc::new(Shared::default()); + let cogroup = matches!(shape, RelationShape::CoGrouped { .. }); + let schema = input(cogroup); + let both = || vec![column(0, ArrowType::Int32), column(1, ArrowType::Int64)]; + let names = || vec!["k".to_owned(), "v".to_owned()]; + let sides = if cogroup { 2 } else { 1 }; + let (client, output) = match yields { + Yield::Count | Yield::CountTwice | Yield::Coalesce(_) => ( + vec![ArrowType::Int64], + match shape { + RelationShape::Aggregate { .. } => vec![ + Field::new("k", ArrowType::Int32, true), + Field::new("n", ArrowType::Int64, true), + ], + _ => vec![Field::new("n", ArrowType::Int64, true)], + }, + ), + Yield::Twice | Yield::Nothing | Yield::Burst(_) => ( + vec![ArrowType::Int32, ArrowType::Int64], + vec![ + Field::new("k", ArrowType::Int32, true), + Field::new("v", ArrowType::Int64, true), + ], + ), + }; + let lane = RelationLane { + name: "f".to_owned(), + udf: Arc::new(Fake { + shared: Arc::clone(&shared), + yields, + }), + sides: (0..sides).map(|_| both()).collect(), + names: (0..sides).map(|_| names()).collect(), + client, + }; + let sink = RelationSink::try_new( + schema, + shape, + vec![lane], + Arc::new(Schema::new(output)), + limits, + scope, + ) + .unwrap(); + let root = pool(limit); + let query = QueryContext::new( + QueryId(85), + WorkloadClass::Batch, + ExecConfig::default(), + root.child("q", Some(limit)), + ); + let executor: Arc = TokioExecutor::deterministic(1).unwrap(); + let ctx = TaskContext::new( + Arc::clone(&query), + executor, + SpillOwner::new(), + OperatorId(0), + "python_relation", + ); + Rig { + sink, + root, + query, + ctx, + shared, + } +} + +fn limits(max_rows: usize, window: usize) -> RelationLimits { + RelationLimits { + max_rows, + max_bytes: 0, + window, + } +} + +/// Rows `(k, v)`, with a side column for a cogroup. +fn rows(keys: &[i32], side: Option) -> Batch { + let k: Int32Array = keys.iter().map(|&k| Some(k)).collect(); + let v: Int64Array = (0..keys.len()).map(|i| Some(i as i64)).collect(); + let mut columns: Vec = vec![Arc::new(k), Arc::new(v)]; + if let Some(s) = side { + columns.push(Arc::new(Int32Array::from(vec![s; keys.len()]))); + } + Batch::new(RecordBatch::try_new(input(side.is_some()), columns).unwrap()) +} + +/// Runs `batches` through the rig's sink as one driver. +async fn drive(r: &Rig, batches: Vec) -> crate::error::Result<()> { + let g = r.sink.global_state(&r.ctx)?; + let mut l = r.sink.local_state(&r.ctx, &g)?; + for b in batches { + // `Blocked` took the batch: wait, then the next one. `NeedMemory` + // took nothing: the same batch again. + match r.sink.sink(b, &mut l, &g)? { + SinkStatus::Blocked(wait) => wait.await, + SinkStatus::NeedMemory { bytes } => panic!("no memory for {bytes}"), + SinkStatus::NeedMoreInput | SinkStatus::Finished => {} + } + } + r.sink.combine(l, &g)?; + r.sink.finalize(&g, &r.ctx).await +} + +fn done(r: Rig) { + drop((r.ctx, r.query)); + assert_eq!(r.root.reserved(), 0, "memory left after the phase"); +} + +#[tokio::test] +async fn a_map_task_gets_full_frames_then_the_rest_then_end() { + let r = rig(RelationShape::Map, Yield::Twice, limits(4, 8), 64 << 20); + drive( + &r, + vec![rows(&[1, 2, 3], None), rows(&[4, 5, 6, 7, 8, 9], None)], + ) + .await + .unwrap(); + assert_eq!(r.shared.log(), ["B0:4", "B0:4", "B0:1", "E"]); + assert_eq!( + r.shared.opened.load(Ordering::SeqCst), + 1, + "one task per driver" + ); + done(r); +} + +#[tokio::test] +async fn an_empty_partition_opens_no_task() { + let r = rig(RelationShape::Map, Yield::Twice, limits(4, 8), 64 << 20); + drive(&r, Vec::new()).await.unwrap(); + assert!(r.shared.log().is_empty()); + assert_eq!(r.shared.opened.load(Ordering::SeqCst), 0); + done(r); +} + +#[tokio::test] +async fn a_grouped_task_sees_each_key_s_rows_together_and_one_end_group_per_key() { + let r = rig( + RelationShape::Grouped { keys: vec![0] }, + Yield::Count, + limits(2, 16), + 64 << 20, + ); + // Key 2 spans the two batches; key 3 has 3 rows (cut by maxRecords 2). + drive(&r, vec![rows(&[1, 2, 2], None), rows(&[2, 3, 3, 3], None)]) + .await + .unwrap(); + assert_eq!( + r.shared.log(), + ["B0:1", "G", "B0:2", "B0:1", "G", "B0:2", "B0:1", "G", "E"] + ); + done(r); +} + +#[tokio::test] +async fn a_grouped_task_without_keys_sees_one_group() { + let r = rig( + RelationShape::Grouped { keys: Vec::new() }, + Yield::Count, + limits(0, 16), + 64 << 20, + ); + drive(&r, vec![rows(&[1, 2], None), rows(&[3], None)]) + .await + .unwrap(); + assert_eq!(r.shared.log(), ["B0:2", "B0:1", "G", "E"]); + done(r); +} + +#[tokio::test] +async fn a_cogroup_key_on_one_side_gets_a_zero_row_frame_for_the_other() { + let r = rig( + RelationShape::CoGrouped { + keys: vec![0], + side: 2, + }, + Yield::Count, + limits(0, 16), + 64 << 20, + ); + // Sorted by key then side: key 1 left only, key 2 both, key 3 right only. + drive( + &r, + vec![ + rows(&[1, 2], Some(0)), + rows(&[2, 2], Some(1)), + rows(&[3], Some(1)), + ], + ) + .await + .unwrap(); + assert_eq!( + r.shared.log(), + [ + "B0:1", "B1:0", "G", "B0:1", "B1:2", "G", "B0:0", "B1:1", "G", "E" + ] + ); + done(r); +} + +#[tokio::test] +async fn a_one_sided_cogroup_key_over_two_batches_gets_one_empty_frame_first() { + // Key 1 left only and key 3 right only, each split over two batches. + for (side, expected) in [ + (0, ["B0:1", "B0:1", "B1:0", "G", "E"]), + (1, ["B0:0", "B1:1", "B1:1", "G", "E"]), + ] { + let r = rig( + RelationShape::CoGrouped { + keys: vec![0], + side: 2, + }, + Yield::Count, + limits(0, 16), + 64 << 20, + ); + let key = 1 + 2 * side; + drive(&r, vec![rows(&[key], Some(side)), rows(&[key], Some(side))]) + .await + .unwrap(); + assert_eq!(r.shared.log(), expected, "side {side}"); + done(r); + } +} + +#[tokio::test] +async fn an_aggregate_lane_gets_one_group_per_key() { + let r = rig( + RelationShape::Aggregate { + keys: vec![0], + key_values: vec![0], + }, + Yield::Count, + limits(0, 16), + 64 << 20, + ); + drive(&r, vec![rows(&[5, 5, 7], None)]).await.unwrap(); + assert_eq!(r.shared.log(), ["B0:2", "G", "B0:1", "G", "E"]); + done(r); +} + +/// An aggregate shows as many key values as it compares keys: a plan +/// that says otherwise is refused, not run with misaligned keys. +#[test] +fn an_aggregate_with_key_values_unlike_its_keys_is_refused() { + let lane = RelationLane { + name: "f".to_owned(), + udf: Arc::new(Fake { + shared: Arc::new(Shared::default()), + yields: Yield::Count, + }), + sides: vec![vec![ + column(0, ArrowType::Int32), + column(1, ArrowType::Int64), + ]], + names: vec![vec!["k".to_owned(), "v".to_owned()]], + client: vec![ArrowType::Int64], + }; + let err = RelationSink::try_new( + input(false), + RelationShape::Aggregate { + keys: vec![0], + key_values: Vec::new(), + }, + vec![lane], + Arc::new(Schema::new(vec![Field::new("n", ArrowType::Int64, true)])), + limits(0, 16), + None, + ) + .err() + .unwrap(); + assert!(err.to_string().contains("key values"), "{err}"); +} + +#[tokio::test] +async fn a_full_window_blocks_the_driver_until_the_task_takes_frames() { + let r = rig(RelationShape::Map, Yield::Nothing, limits(1, 2), 64 << 20); + r.shared.closed.store(true, Ordering::SeqCst); + let g = r.sink.global_state(&r.ctx).unwrap(); + let mut l = r.sink.local_state(&r.ctx, &g).unwrap(); + let SinkStatus::Blocked(wait) = r.sink.sink(rows(&[1, 2, 3], None), &mut l, &g).unwrap() else { + panic!("three frames over a window of two must block"); + }; + let waiting = tokio::spawn(wait); + tokio::task::yield_now().await; + assert!( + !waiting.is_finished(), + "the window freed while the task took nothing" + ); + assert!(r.shared.log().is_empty()); + r.shared.open_gate(); + tokio::time::timeout(Duration::from_secs(5), waiting) + .await + .unwrap() + .unwrap(); + r.sink.combine(l, &g).unwrap(); + r.sink.finalize(&g, &r.ctx).await.unwrap(); + assert_eq!(r.shared.log(), ["B0:1", "B0:1", "B0:1", "E"]); + drop(g); + done(r); +} + +#[tokio::test] +async fn a_failed_frame_fails_the_phase_and_releases_its_memory() { + let r = rig( + RelationShape::Grouped { keys: vec![0] }, + Yield::Count, + limits(0, 16), + 64 << 20, + ); + *r.shared.fail_at.lock().unwrap() = Some(2); + let err = drive(&r, vec![rows(&[1, 2, 3], None)]).await.unwrap_err(); + assert!(err.to_string().contains("the UDF raised"), "{err}"); + done(r); +} + +#[tokio::test] +async fn a_cancel_mid_group_ends_cleanly_and_drops_the_task() { + let r = rig( + RelationShape::Grouped { keys: vec![0] }, + Yield::Count, + limits(1, 64), + 64 << 20, + ); + r.shared.closed.store(true, Ordering::SeqCst); + let g = r.sink.global_state(&r.ctx).unwrap(); + let mut l = r.sink.local_state(&r.ctx, &g).unwrap(); + assert!(matches!( + r.sink.sink(rows(&[1, 1, 1], None), &mut l, &g).unwrap(), + SinkStatus::NeedMoreInput + )); + assert!(r.root.reserved() > 0, "the queued frames hold memory"); + r.sink.combine(l, &g).unwrap(); + r.query.cancel(); + let err = tokio::time::timeout(Duration::from_secs(5), r.sink.finalize(&g, &r.ctx)) + .await + .unwrap() + .unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::OperationCanceled); + drop(g); + assert_eq!( + r.shared.dropped.load(Ordering::SeqCst), + 1, + "the task was dropped" + ); + done(r); +} + +#[tokio::test] +async fn a_refused_reservation_sends_nothing_so_the_retry_is_exact() { + let r = rig(RelationShape::Map, Yield::Twice, limits(0, 8), 512); + let g = r.sink.global_state(&r.ctx).unwrap(); + let mut l = r.sink.local_state(&r.ctx, &g).unwrap(); + let keys: Vec = (0..1000).collect(); + let status = r.sink.sink(rows(&keys, None), &mut l, &g).unwrap(); + assert!( + matches!(status, SinkStatus::NeedMemory { .. }), + "{status:?}" + ); + assert_eq!(r.shared.opened.load(Ordering::SeqCst), 0, "nothing sent"); + drop((l, g)); + done(r); +} + +/// A runtime of the scalar kinds only. +#[derive(Debug)] +struct ScalarOnly; + +impl crate::python_eval::UdfRuntime for ScalarOnly { + fn prepare( + &self, + spec: &crate::python_eval::UdfCallSpec, + ) -> crate::error::Result> { + Err(InternalError { + message: format!("no {}", spec.name), + } + .into()) + } +} + +#[test] +fn a_runtime_without_relation_kinds_refuses_them_as_not_available() { + let spec = super::RelationCallSpec { + name: "f".to_owned(), + version: 3, + eval_type: krishiv_plan::PythonEvalType::GroupedMapArrow, + thread_safe: true, + input_types: vec![vec![krishiv_types::DataType::Long]], + input_names: vec![vec!["k".to_owned()]], + key_count: 1, + arity: 1, + output_names: vec!["n".to_owned()], + output_types: vec![krishiv_types::DataType::Long], + output_arrow: Arc::new(Schema::new(vec![Field::new("n", ArrowType::Int64, true)])), + conf: super::RunnerConf::default(), + }; + let err = crate::python_eval::UdfRuntime::prepare_relation(&ScalarOnly, &spec) + .err() + .unwrap(); + assert_eq!(err.error_class(), ErrorClass::KrishivUdfNotAvailable); + assert!(err.to_string().contains("f@3"), "{err}"); +} + +#[tokio::test] +async fn max_bytes_per_batch_cuts_frames_between_rows() { + let limits = RelationLimits { + max_rows: 0, + max_bytes: 1, + window: 64, + }; + let r = rig(RelationShape::Map, Yield::Twice, limits, 64 << 20); + assert!(format!("{:?}", r.sink).contains("RelationSink")); + drive(&r, vec![rows(&[1, 2, 3], None)]).await.unwrap(); + // At least one row per frame, however small the byte cap. + assert_eq!(r.shared.log(), ["B0:1", "B0:1", "B0:1", "E"]); + done(r); +} + +/// A computed argument of uneven width: row `i` (its `v`) is a string of +/// `widths[i]` bytes. +#[derive(Debug)] +struct Uneven { + widths: Vec, +} + +impl ScalarExpr for Uneven { + fn data_type(&self) -> ArrowType { + ArrowType::Utf8 + } + + fn nullable(&self) -> bool { + false + } + + fn evaluate(&self, batch: &RecordBatch) -> crate::error::Result { + let v = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + let a: arrow_array::StringArray = v + .values() + .iter() + .map(|&i| Some("x".repeat(self.widths[usize::try_from(i).unwrap()]))) + .collect(); + Ok(Arc::new(a)) + } +} + +/// The frames a map task gets for one batch of rows of `widths` bytes under +/// a `maxBytesPerBatch` of 32 KiB. +async fn frames_of_widths(widths: Vec) -> Vec { + let limits = RelationLimits { + max_rows: 0, + max_bytes: 32 << 10, + window: 64, + }; + let keys: Vec = (0..i32::try_from(widths.len()).unwrap()).collect(); + let wide: Arc = Arc::new(Uneven { widths }); + let r = rig_columns(limits, 64 << 20, &[(wide, "s")]); + drive(&r, vec![rows(&keys, None)]).await.unwrap(); + let log = r.shared.log(); + done(r); + log +} + +#[tokio::test] +async fn max_bytes_per_batch_cuts_by_the_frame_columns_row_sizes() { + // Two 30 KiB strings, then 100 empty ones. The input rows are a few + // bytes each; the argument is not. As Spark's writer: the first frame + // closes at the row that brings it to the cap (row 2), not at the + // batch's average bytes per row. + let mut widths = vec![30 << 10, 30 << 10]; + widths.extend([0; 100]); + assert_eq!(frames_of_widths(widths).await, ["B0:2", "B0:100", "E"]); + // 20 KiB rows: two per frame (the second crosses the cap). + assert_eq!( + frames_of_widths(vec![20 << 10; 5]).await, + ["B0:2", "B0:2", "B0:1", "E"] + ); +} + +#[tokio::test] +async fn a_row_over_max_bytes_per_batch_closes_its_frame() { + // As Spark's writer: a row over the cap by itself ends the frame it + // joins (rows before it under the cap share it), and alone it is a + // frame of its own. + assert_eq!( + frames_of_widths(vec![10, 40 << 10, 10, 10]).await, + ["B0:2", "B0:2", "E"] + ); + assert_eq!( + frames_of_widths(vec![40 << 10, 10, 10]).await, + ["B0:1", "B0:2", "E"] + ); +} + +#[tokio::test] +async fn a_result_of_another_width_fails_the_phase() { + // The task yields `n BIGINT` where `k, v` were declared. + let shared = Arc::new(Shared::default()); + let lane = RelationLane { + name: "f".to_owned(), + udf: Arc::new(Fake { + shared: Arc::clone(&shared), + yields: Yield::Count, + }), + sides: vec![vec![ + column(0, ArrowType::Int32), + column(1, ArrowType::Int64), + ]], + names: vec![vec!["k".to_owned(), "v".to_owned()]], + client: vec![ArrowType::Int32, ArrowType::Int64], + }; + let mut r = rig( + RelationShape::Grouped { keys: vec![0] }, + Yield::Count, + limits(0, 8), + 64 << 20, + ); + r.sink = RelationSink::try_new( + input(false), + RelationShape::Grouped { keys: vec![0] }, + vec![lane], + Arc::new(Schema::new(vec![ + Field::new("k", ArrowType::Int32, true), + Field::new("v", ArrowType::Int64, true), + ])), + limits(0, 8), + None, + ) + .unwrap(); + let err = drive(&r, vec![rows(&[1, 2], None)]).await.unwrap_err(); + assert_eq!( + err.error_class(), + ErrorClass::KrishivPythonUdfResultMismatch, + "{err}" + ); + assert!( + err.to_string().contains("columns where 2 were declared"), + "{err}" + ); + done(r); +} + +#[tokio::test] +async fn the_window_bounds_the_frames_and_the_rest_waits_in_the_backlog() { + let r = rig(RelationShape::Map, Yield::Nothing, limits(1, 2), 64 << 20); + r.shared.closed.store(true, Ordering::SeqCst); + let g = r.sink.global_state(&r.ctx).unwrap(); + let mut l = r.sink.local_state(&r.ctx, &g).unwrap(); + let SinkStatus::Blocked(wait) = r + .sink + .sink(rows(&[1, 2, 3, 4, 5], None), &mut l, &g) + .unwrap() + else { + panic!("five frames over a window of two must block"); + }; + // Two frames in flight (one being sent), three rows cut and waiting. + assert_eq!(super::sink::queued(&mut l).unwrap(), (vec![2], 3)); + r.shared.open_gate(); + tokio::time::timeout(Duration::from_secs(5), wait) + .await + .unwrap(); + r.sink.combine(l, &g).unwrap(); + r.sink.finalize(&g, &r.ctx).await.unwrap(); + assert_eq!( + r.shared.log(), + ["B0:1", "B0:1", "B0:1", "B0:1", "B0:1", "E"] + ); + drop(g); + done(r); +} + +#[tokio::test] +async fn a_result_the_pool_refuses_stops_the_task_s_reads_until_reserved() { + let limit = 256 << 10; + let r = rig(RelationShape::Map, Yield::Burst(64), limits(0, 8), limit); + let g = r.sink.global_state(&r.ctx).unwrap(); + let mut l = r.sink.local_state(&r.ctx, &g).unwrap(); + let keys: Vec = (0..1000).collect(); + let status = r.sink.sink(rows(&keys, None), &mut l, &g).unwrap(); + assert!(matches!(status, SinkStatus::NeedMoreInput), "{status:?}"); + // Each result is reserved as it is read: 64 copies of the input do + // not fit, so the task's reads stop at the refused one. + let taken = r.shared.taken.load(Ordering::SeqCst); + assert!(taken > 0 && taken < 64, "{taken} results read"); + assert!(r.root.reserved() <= limit); + // The next call asks the driver for memory before it takes its input. + let status = r.sink.sink(rows(&[1], None), &mut l, &g).unwrap(); + assert!( + matches!(status, SinkStatus::NeedMemory { .. }), + "{status:?}" + ); + assert_eq!(r.shared.taken.load(Ordering::SeqCst), taken); + drop((l, g)); + assert_eq!(r.shared.dropped.load(Ordering::SeqCst), 1); + done(r); +} + +#[tokio::test] +async fn an_aggregate_task_yielding_another_number_of_values_than_groups_fails() { + // Too many values fail as soon as they outnumber the groups. + for (yields, values) in [(Yield::CountTwice, " values"), (Yield::Nothing, "0 values")] { + let r = rig( + RelationShape::Aggregate { + keys: vec![0], + key_values: vec![0], + }, + yields, + limits(0, 16), + 64 << 20, + ); + let err = drive(&r, vec![rows(&[5, 5, 7], None)]).await.unwrap_err(); + assert_eq!( + err.error_class(), + ErrorClass::KrishivPythonUdfResultMismatch, + "{err}" + ); + assert!(err.to_string().contains(values), "{err}"); + assert!(err.to_string().contains("values for 2 groups"), "{err}"); + done(r); + } +} + +#[tokio::test] +async fn a_blocked_driver_sends_each_row_once() { + // Two frames per batch over a window of one: every batch blocks. + let r = rig( + RelationShape::Grouped { keys: vec![0] }, + Yield::Count, + limits(0, 1), + 64 << 20, + ); + drive(&r, vec![rows(&[1, 2], None), rows(&[2, 3], None)]) + .await + .unwrap(); + assert_eq!( + r.shared.log(), + ["B0:1", "G", "B0:1", "B0:1", "G", "B0:1", "G", "E"] + ); + done(r); +} + +#[tokio::test] +async fn a_weighted_batch_is_refused() { + let r = rig(RelationShape::Map, Yield::Twice, limits(0, 8), 64 << 20); + let g = r.sink.global_state(&r.ctx).unwrap(); + let mut l = r.sink.local_state(&r.ctx, &g).unwrap(); + let (data, _) = rows(&[1, 2], None).into_parts(); + let retracting = Batch::with_weights(data, krishiv_types::batch::Weights::Const(-1)).unwrap(); + let err = r.sink.sink(retracting, &mut l, &g).unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::KrishivUnsupportedFeature); + assert!(err.to_string().contains("weighted"), "{err}"); + drop((l, g)); + done(r); +} + +#[tokio::test] +async fn an_aggregate_waits_for_its_slowest_lane_without_spinning() { + let mut r = rig( + RelationShape::Aggregate { + keys: vec![0], + key_values: vec![0], + }, + Yield::Count, + limits(0, 16), + 64 << 20, + ); + let slow = Arc::new(Shared::default()); + slow.closed.store(true, Ordering::SeqCst); + let lane = |shared: &Arc, name: &str| RelationLane { + name: name.to_owned(), + udf: Arc::new(Fake { + shared: Arc::clone(shared), + yields: Yield::Count, + }), + sides: vec![vec![ + column(0, ArrowType::Int32), + column(1, ArrowType::Int64), + ]], + names: vec![vec!["k".to_owned(), "v".to_owned()]], + client: vec![ArrowType::Int64], + }; + r.sink = RelationSink::try_new( + input(false), + RelationShape::Aggregate { + keys: vec![0], + key_values: vec![0], + }, + vec![lane(&r.shared, "f"), lane(&slow, "g")], + Arc::new(Schema::new(vec![ + Field::new("k", ArrowType::Int32, true), + Field::new("n", ArrowType::Int64, true), + Field::new("m", ArrowType::Int64, true), + ])), + limits(0, 16), + None, + ) + .unwrap(); + let g = r.sink.global_state(&r.ctx).unwrap(); + let mut l = r.sink.local_state(&r.ctx, &g).unwrap(); + let status = r.sink.sink(rows(&[5, 5, 7], None), &mut l, &g).unwrap(); + assert!(matches!(status, SinkStatus::NeedMoreInput), "{status:?}"); + r.sink.combine(l, &g).unwrap(); + let mut finalize = r.sink.finalize(&g, &r.ctx); + // The fast lane's values wait for the slow lane's: pending, no busy loop. + let first = poll_fn(|cx| Poll::Ready(finalize.as_mut().poll(cx))).await; + assert!(first.is_pending(), "{first:?}"); + assert_eq!(r.shared.log(), ["B0:2", "G", "B0:1", "G", "E"]); + assert!(slow.log().is_empty()); + slow.open_gate(); + tokio::time::timeout(Duration::from_secs(5), finalize) + .await + .unwrap() + .unwrap(); + assert!( + r.shared.polls.load(Ordering::SeqCst) < 100, + "polled in a loop" + ); + drop(g); + done(r); +} + +/// An aggregate whose fast lane, with enough groups, would fill the pool +/// with values before the slow lane's first is reserved: no group can be +/// joined, and the lane queues are not spillable, so the slow lane's value +/// would wait for memory only its peer holds. A lane reads no result while it holds more values +/// than a running peer that has frames to take, so while the gated lane +/// has frames the fast one holds one result batch past it (checked before +/// the gate opens). (A deep window: the stall as first seen.) +/// +/// Once the gate opens, the gated lane takes its whole window at once, so +/// it has no frame left to take and the hold ends: past that point the +/// bound is the results of the at most `window` frames each lane already +/// took, and the two lanes share the pool, which writes stored groups out +/// when it is full. The test checks the hold, not the pool: an unheld fast +/// lane reads every value before the gate opens (`taken` would be +/// `groups`, not 1), so the groups are kept few enough that `finalize` +/// takes a fraction of a second and the short timeout catches a stall, not +/// a slow runner. +#[tokio::test] +async fn a_fast_aggregate_lane_waits_for_its_peers_instead_of_filling_the_pool() { + let dir = tempfile::tempdir().unwrap(); + let scope = krishiv_memory::SpillManager::open(&krishiv_memory::SpillConfig { + dir: dir.path().join("spill"), + ..krishiv_memory::SpillConfig::default() + }) + .unwrap() + .scope() + .unwrap(); + let shape = || RelationShape::Aggregate { + keys: vec![0], + key_values: vec![0], + }; + let mut r = rig(shape(), Yield::Count, limits(0, 32768), 1 << 20); + let slow = Arc::new(Shared::default()); + slow.closed.store(true, Ordering::SeqCst); + let lane = |shared: &Arc, name: &str| RelationLane { + name: name.to_owned(), + udf: Arc::new(Fake { + shared: Arc::clone(shared), + yields: Yield::Count, + }), + sides: vec![vec![ + column(0, ArrowType::Int32), + column(1, ArrowType::Int64), + ]], + names: vec![vec!["k".to_owned(), "v".to_owned()]], + client: vec![ArrowType::Int64], + }; + r.sink = RelationSink::try_new( + input(false), + shape(), + vec![lane(&r.shared, "f"), lane(&slow, "g")], + Arc::new(Schema::new(vec![ + Field::new("k", ArrowType::Int32, true), + Field::new("n", ArrowType::Int64, true), + Field::new("m", ArrowType::Int64, true), + ])), + limits(0, 32768), + Some(scope), + ) + .unwrap(); + let g = r.sink.global_state(&r.ctx).unwrap(); + let mut l = r.sink.local_state(&r.ctx, &g).unwrap(); + let groups = 3000; + let keys: Vec = (0..groups).collect(); + let status = r.sink.sink(rows(&keys, None), &mut l, &g).unwrap(); + assert!(matches!(status, SinkStatus::NeedMoreInput), "{status:?}"); + // The slow lane has no value yet: the fast one read one, then waits. + assert_eq!(r.shared.taken.load(Ordering::SeqCst), 1); + r.sink.combine(l, &g).unwrap(); + slow.open_gate(); + // The open slow lane takes its whole deep window at once, so the hold + // ends and the two lanes share the pool: never stuck. + tokio::time::timeout(Duration::from_secs(10), r.sink.finalize(&g, &r.ctx)) + .await + .unwrap() + .unwrap(); + let groups = usize::try_from(groups).unwrap(); + assert_eq!(r.shared.taken.load(Ordering::SeqCst), groups); + assert_eq!(slow.taken.load(Ordering::SeqCst), groups); + drop(g); + done(r); +} + +/// A two-lane aggregate on key `k`: lane `f` of `fast`, lane `g` of +/// `slow`, at `window` frames in flight per lane. +fn two_lane_aggregate( + lanes: [(&Arc, Yield); 2], + window: usize, + scope: Option, +) -> Arc { + let lane = |(shared, yields): (&Arc, Yield), name: &str| RelationLane { + name: name.to_owned(), + udf: Arc::new(Fake { + shared: Arc::clone(shared), + yields, + }), + sides: vec![vec![ + column(0, ArrowType::Int32), + column(1, ArrowType::Int64), + ]], + names: vec![vec!["k".to_owned(), "v".to_owned()]], + client: vec![ArrowType::Int64], + }; + let [fast, slow] = lanes; + RelationSink::try_new( + input(false), + RelationShape::Aggregate { + keys: vec![0], + key_values: vec![0], + }, + vec![lane(fast, "f"), lane(slow, "g")], + Arc::new(Schema::new(vec![ + Field::new("k", ArrowType::Int32, true), + Field::new("n", ArrowType::Int64, true), + Field::new("m", ArrowType::Int64, true), + ])), + limits(0, window), + scope, + ) + .unwrap() +} + +/// The groups `0..n`, one row each. +fn groups(n: i32) -> Batch { + rows(&(0..n).collect::>(), None) +} + +/// The production window: the slow lane's window, not its memory, stops +/// the feed. While it has frames it has not taken, the fast lane reads no +/// value past its first; once it takes them both lanes finish. +#[tokio::test] +async fn at_the_production_window_a_fast_lane_waits_while_its_peer_has_frames() { + let shape = RelationShape::Aggregate { + keys: vec![0], + key_values: vec![0], + }; + let mut r = rig(shape, Yield::Count, limits(0, DEFAULT_WINDOW), 1 << 20); + let slow = Arc::new(Shared::default()); + slow.closed.store(true, Ordering::SeqCst); + r.sink = two_lane_aggregate( + [(&r.shared, Yield::Count), (&slow, Yield::Count)], + DEFAULT_WINDOW, + None, + ); + let g = r.sink.global_state(&r.ctx).unwrap(); + let mut l = r.sink.local_state(&r.ctx, &g).unwrap(); + let SinkStatus::Blocked(wait) = r.sink.sink(groups(2_000), &mut l, &g).unwrap() else { + panic!("the slow lane's window is full"); + }; + assert_eq!(r.shared.taken.load(Ordering::SeqCst), 1); + slow.open_gate(); + tokio::time::timeout(Duration::from_secs(10), async { + wait.await; + r.sink.combine(l, &g).unwrap(); + r.sink.finalize(&g, &r.ctx).await.unwrap(); + }) + .await + .unwrap(); + assert_eq!(r.shared.taken.load(Ordering::SeqCst), 2_000); + assert_eq!(slow.taken.load(Ordering::SeqCst), 2_000); + drop(g); + done(r); +} + +/// A peer runtime batching its values across groups needs later frames +/// before it yields; the fast lane, whose runtime takes no frame while a +/// result is unread, must not hold back then, or no lane gets a frame. +#[tokio::test] +async fn a_peer_batching_values_across_groups_never_hangs_the_fast_lane() { + let shape = RelationShape::Aggregate { + keys: vec![0], + key_values: vec![0], + }; + let mut r = rig(shape, Yield::Count, limits(0, DEFAULT_WINDOW), 64 << 20); + *r.shared.out_cap.lock().unwrap() = Some(1); + let slow = Arc::new(Shared::default()); + r.sink = two_lane_aggregate( + [(&r.shared, Yield::Count), (&slow, Yield::Coalesce(40))], + DEFAULT_WINDOW, + None, + ); + tokio::time::timeout(Duration::from_secs(10), drive(&r, vec![groups(100)])) + .await + .expect("the query hung") + .unwrap(); + assert_eq!(r.shared.taken.load(Ordering::SeqCst), 100); + assert_eq!(slow.taken.load(Ordering::SeqCst), 3); + done(r); +} + +/// The "took its last frame" wake. The peer, gated, has frames it has not +/// taken, so the fast lane is held after its first value, its runtime +/// stops taking frames and its window stops the feed. The gate opens: in +/// one poll the held lane parks first, then the peer takes its last frames +/// without yielding (it batches 40 groups) and its runtime signals its +/// reader only on output. Only the peer's wake as it took its last frame +/// polls the held lane again; else nothing does and the query hangs. +#[tokio::test] +async fn a_quiet_peer_taking_its_last_frame_releases_the_fast_lane() { + let shape = RelationShape::Aggregate { + keys: vec![0], + key_values: vec![0], + }; + let mut r = rig(shape, Yield::Count, limits(0, DEFAULT_WINDOW), 64 << 20); + *r.shared.out_cap.lock().unwrap() = Some(1); + let slow = Arc::new(Shared::default()); + slow.closed.store(true, Ordering::SeqCst); + slow.quiet.store(true, Ordering::SeqCst); + r.sink = two_lane_aggregate( + [(&r.shared, Yield::Count), (&slow, Yield::Coalesce(40))], + DEFAULT_WINDOW, + None, + ); + let g = r.sink.global_state(&r.ctx).unwrap(); + let mut l = r.sink.local_state(&r.ctx, &g).unwrap(); + let SinkStatus::Blocked(wait) = r.sink.sink(groups(100), &mut l, &g).unwrap() else { + panic!("the windows are full"); + }; + assert_eq!(r.shared.taken.load(Ordering::SeqCst), 1); + slow.open_gate(); + tokio::time::timeout(Duration::from_secs(10), async { + wait.await; + r.sink.combine(l, &g).unwrap(); + r.sink.finalize(&g, &r.ctx).await.unwrap(); + }) + .await + .expect("the query hung"); + assert_eq!(r.shared.taken.load(Ordering::SeqCst), 100); + assert_eq!(slow.taken.load(Ordering::SeqCst), 3); + drop(g); + done(r); +} + +/// A peer that ends with fewer values no longer holds the fast lane back: +/// the phase fails with the mismatch, it does not hang. +#[tokio::test] +async fn a_peer_ending_with_fewer_values_fails_the_phase_instead_of_hanging() { + let shape = RelationShape::Aggregate { + keys: vec![0], + key_values: vec![0], + }; + let mut r = rig(shape, Yield::Count, limits(0, DEFAULT_WINDOW), 64 << 20); + *r.shared.out_cap.lock().unwrap() = Some(1); + let slow = Arc::new(Shared::default()); + r.sink = two_lane_aggregate( + [(&r.shared, Yield::Count), (&slow, Yield::Nothing)], + DEFAULT_WINDOW, + None, + ); + let err = tokio::time::timeout(Duration::from_secs(10), drive(&r, vec![groups(50)])) + .await + .expect("the query hung") + .unwrap_err(); + assert_eq!( + err.error_class(), + ErrorClass::KrishivPythonUdfResultMismatch, + "{err}" + ); + done(r); +} + +/// A peer failing while the fast lane waits for it: its error, no hang. +#[tokio::test] +async fn a_peer_failing_while_the_fast_lane_waits_gives_the_peer_s_error() { + let shape = RelationShape::Aggregate { + keys: vec![0], + key_values: vec![0], + }; + let mut r = rig(shape, Yield::Count, limits(0, DEFAULT_WINDOW), 64 << 20); + let slow = Arc::new(Shared::default()); + slow.closed.store(true, Ordering::SeqCst); + *slow.fail_at.lock().unwrap() = Some(1); + r.sink = two_lane_aggregate( + [(&r.shared, Yield::Count), (&slow, Yield::Count)], + DEFAULT_WINDOW, + None, + ); + let g = r.sink.global_state(&r.ctx).unwrap(); + let mut l = r.sink.local_state(&r.ctx, &g).unwrap(); + let SinkStatus::Blocked(wait) = r.sink.sink(groups(100), &mut l, &g).unwrap() else { + panic!("the slow lane's window is full"); + }; + // Held back: the slow lane has frames it has not taken. + assert_eq!(r.shared.taken.load(Ordering::SeqCst), 1); + slow.open_gate(); + let err = tokio::time::timeout(Duration::from_secs(10), async { + wait.await; + r.sink.combine(l, &g)?; + r.sink.finalize(&g, &r.ctx).await + }) + .await + .expect("the query hung") + .unwrap_err(); + assert!(err.to_string().contains("the UDF raised"), "{err}"); + drop(g); + done(r); +} + +#[tokio::test] +async fn a_result_refused_in_finalize_spills_instead_of_waiting() { + let dir = tempfile::tempdir().unwrap(); + let scope = krishiv_memory::SpillManager::open(&krishiv_memory::SpillConfig { + dir: dir.path().join("spill"), + ..krishiv_memory::SpillConfig::default() + }) + .unwrap() + .scope() + .unwrap(); + let r = rig_in( + RelationShape::Map, + Yield::Burst(64), + limits(0, 8), + 256 << 10, + Some(scope), + ); + let g = r.sink.global_state(&r.ctx).unwrap(); + let mut l = r.sink.local_state(&r.ctx, &g).unwrap(); + let keys: Vec = (0..1000).collect(); + let status = r.sink.sink(rows(&keys, None), &mut l, &g).unwrap(); + assert!(matches!(status, SinkStatus::NeedMoreInput), "{status:?}"); + assert!( + r.shared.taken.load(Ordering::SeqCst) < 64, + "a result refused" + ); + r.sink.combine(l, &g).unwrap(); + // Faster than the pool's spill wait: the refused results are written + // out, not waited for. + tokio::time::timeout(Duration::from_secs(4), r.sink.finalize(&g, &r.ctx)) + .await + .unwrap() + .unwrap(); + assert_eq!(r.shared.taken.load(Ordering::SeqCst), 64); + drop(g); + done(r); +} + +/// A frame whose copies the pool refuses in `finalize` while a lane's +/// results hold the rest of the pool: the results are stored first, written +/// out, which frees their memory, so the frame gets its copies instead of +/// waiting for memory nothing can release (the lane queues are not +/// spillable, nor the store's parts before it finalizes). +#[tokio::test] +async fn finalize_drains_results_before_waiting_for_frame_copies() { + let dir = tempfile::tempdir().unwrap(); + let scope = krishiv_memory::SpillManager::open(&krishiv_memory::SpillConfig { + dir: dir.path().join("spill"), + ..krishiv_memory::SpillConfig::default() + }) + .unwrap() + .scope() + .unwrap(); + let mut r = rig(RelationShape::Map, Yield::Nothing, limits(100, 8), 1 << 20); + // `k INT, v BIGINT, d DECIMAL(18, 2)`: lane `f` echoes `k, v` many + // times; lane `g` takes `d`, whose frames are copied to the client's + // `Decimal128`. + let schema = Arc::new(Schema::new(vec![ + Field::new("k", ArrowType::Int32, true), + Field::new("v", ArrowType::Int64, true), + Field::new("d", ArrowType::Decimal64(18, 2), true), + ])); + let quiet = Arc::new(Shared::default()); + let lane = + |shared: &Arc, yields, sides: Vec>, names| RelationLane { + name: "f".to_owned(), + udf: Arc::new(Fake { + shared: Arc::clone(shared), + yields, + }), + sides: vec![sides], + names: vec![names], + client: vec![ArrowType::Int32, ArrowType::Int64], + }; + r.sink = RelationSink::try_new( + Arc::clone(&schema), + RelationShape::Map, + vec![ + lane( + &r.shared, + Yield::Burst(2000), + vec![column(0, ArrowType::Int32), column(1, ArrowType::Int64)], + vec!["k".to_owned(), "v".to_owned()], + ), + lane( + &quiet, + Yield::Nothing, + vec![column(2, ArrowType::Decimal64(18, 2))], + vec!["d".to_owned()], + ), + ], + Arc::new(Schema::new(vec![ + Field::new("k", ArrowType::Int32, true), + Field::new("v", ArrowType::Int64, true), + ])), + limits(100, 8), + Some(scope), + ) + .unwrap(); + let g = r.sink.global_state(&r.ctx).unwrap(); + let mut l = r.sink.local_state(&r.ctx, &g).unwrap(); + // 150 rows: a frame of 100 now, 50 held for `combine`. + let n = 150; + let k: Int32Array = (0..n).map(Some).collect(); + let v: Int64Array = (0..i64::from(n)).map(Some).collect(); + let s = arrow_array::Decimal64Array::from_iter_values(0..i64::from(n)) + .with_precision_and_scale(18, 2) + .unwrap(); + let batch = Batch::new( + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(k), Arc::new(v), Arc::new(s)], + ) + .unwrap(), + ); + let status = r.sink.sink(batch, &mut l, &g).unwrap(); + assert!(matches!(status, SinkStatus::NeedMoreInput), "{status:?}"); + assert!( + r.shared.taken.load(Ordering::SeqCst) < 2000, + "a result refused" + ); + // Whatever the pool has left goes to another holder, so the held frame + // can get its copies only from memory the results release. + let mut hog = r.ctx.consumer().reservation(); + let mut step = 1 << 20; + while step >= 64 { + if hog.try_grow(step).is_err() { + step /= 2; + } + } + r.sink.combine(l, &g).unwrap(); + // The held frame is refused its copies. + assert_eq!(quiet.log(), ["B0:100"]); + // The other holder lets go once the held frame is sent: before that, + // only the results' memory can be had. + let release = async { + while !quiet.log().iter().any(|e| e == "B0:50") { + tokio::time::sleep(Duration::from_millis(1)).await; + } + drop(hog); + }; + let (finalized, ()) = tokio::time::timeout( + Duration::from_secs(4), + futures::future::join(r.sink.finalize(&g, &r.ctx), release), + ) + .await + .unwrap(); + finalized.unwrap(); + assert_eq!(quiet.log(), ["B0:100", "B0:50", "E"]); + assert_eq!(r.shared.taken.load(Ordering::SeqCst), 4000); + drop(g); + done(r); +} + +/// A batch of two groups: the first group's frame is taken while the +/// second waits in the backlog, which still holds its share of the batch's +/// reservation (the slices keep the batch's buffers). +#[tokio::test] +async fn a_batch_s_later_group_keeps_its_share_of_the_input_reservation() { + let r = rig( + RelationShape::Grouped { keys: vec![0] }, + Yield::Nothing, + limits(0, 1), + 64 << 20, + ); + r.shared.closed.store(true, Ordering::SeqCst); + let g = r.sink.global_state(&r.ctx).unwrap(); + let mut l = r.sink.local_state(&r.ctx, &g).unwrap(); + let keys: Vec = (0..20_000).map(|i| 1 + i / 10_000).collect(); + let batch = rows(&keys, None); + let bytes = batch.data().get_array_memory_size(); + let SinkStatus::Blocked(wait) = r.sink.sink(batch, &mut l, &g).unwrap() else { + panic!("a window of one must block"); + }; + // The first group's frame in flight; its end and the second group's + // rows wait. + assert_eq!(super::sink::queued(&mut l).unwrap(), (vec![1], 2)); + // The task takes the first frame only. + *r.shared.admit.lock().unwrap() = Some(1); + r.shared.open_gate(); + let mut wait = wait; + let first = poll_fn(|cx| Poll::Ready(wait.as_mut().poll(cx))).await; + assert!(first.is_pending(), "ready with the task stuck"); + // The first frame is taken and its memory freed; the second group's + // rows are still reserved. + assert_eq!(r.shared.log(), ["B0:10000"]); + assert_eq!(super::sink::queued(&mut l).unwrap(), (vec![1], 1)); + assert!( + r.root.reserved() >= bytes / 2, + "{} of {bytes} bytes reserved", + r.root.reserved() + ); + *r.shared.admit.lock().unwrap() = None; + r.shared.open_gate(); + tokio::time::timeout(Duration::from_secs(5), wait) + .await + .unwrap(); + r.sink.combine(l, &g).unwrap(); + r.sink.finalize(&g, &r.ctx).await.unwrap(); + assert_eq!(r.shared.log(), ["B0:10000", "G", "B0:10000", "G", "E"]); + drop(g); + done(r); +} + +/// A frame column computed from the rows (`python_udaf(repeat(s, 1000))`) +/// holds memory of its own while the frame is queued: it is reserved, as +/// the copies of a client form are, not only the input rows. A column read +/// shares the input's buffers and adds nothing. +#[tokio::test] +async fn a_computed_frame_column_is_reserved_while_queued() { + let width = 1000; + let r = rig_columns( + limits(0, 4), + 64 << 20, + &[ + (column(0, ArrowType::Int32), "k"), + (Arc::new(Wide { width }), "s"), + ], + ); + r.shared.closed.store(true, Ordering::SeqCst); + let g = r.sink.global_state(&r.ctx).unwrap(); + let mut l = r.sink.local_state(&r.ctx, &g).unwrap(); + let keys: Vec = (0..100).collect(); + let batch = rows(&keys, None); + let input = batch.data().get_array_memory_size(); + let status = r.sink.sink(batch, &mut l, &g).unwrap(); + assert!(matches!(status, SinkStatus::NeedMoreInput), "{status:?}"); + // The frame waits for the closed task: its rows and its strings. + assert_eq!(super::sink::queued(&mut l).unwrap(), (vec![1], 0)); + let reserved = r.root.reserved(); + assert!( + reserved >= input + 100 * width, + "{reserved} bytes reserved for {input} input bytes and {} computed", + 100 * width + ); + r.shared.open_gate(); + r.sink.combine(l, &g).unwrap(); + r.sink.finalize(&g, &r.ctx).await.unwrap(); + assert_eq!(r.shared.log(), ["B0:100", "E"]); + drop(g); + done(r); +} + +/// A column read adds nothing to the frame's reservation: it shares the +/// input's buffers, which the input's reservation covers. +#[tokio::test] +async fn a_frame_of_column_reads_reserves_only_the_input() { + let r = rig_columns( + limits(0, 4), + 64 << 20, + &[ + (column(0, ArrowType::Int32), "k"), + (column(1, ArrowType::Int64), "v"), + ], + ); + r.shared.closed.store(true, Ordering::SeqCst); + let g = r.sink.global_state(&r.ctx).unwrap(); + let mut l = r.sink.local_state(&r.ctx, &g).unwrap(); + let keys: Vec = (0..10_000).collect(); + let batch = rows(&keys, None); + let input = batch.data().get_array_memory_size(); + let status = r.sink.sink(batch, &mut l, &g).unwrap(); + assert!(matches!(status, SinkStatus::NeedMoreInput), "{status:?}"); + let reserved = r.root.reserved(); + // The quantum rounds the input's bytes up. + assert!( + reserved < input + (4 << 10), + "{reserved} bytes reserved for {input} input bytes" + ); + r.shared.open_gate(); + r.sink.combine(l, &g).unwrap(); + r.sink.finalize(&g, &r.ctx).await.unwrap(); + drop(g); + done(r); +} + +/// More frames per batch than the window takes (`maxRecordsPerBatch=1`): +/// the driver is blocked until its backlog drained, so the cuts held never +/// outgrow one input batch, however many batches come. +#[tokio::test] +async fn the_backlog_drains_before_the_next_input_is_taken() { + let r = rig(RelationShape::Map, Yield::Nothing, limits(1, 4), 64 << 20); + let g = r.sink.global_state(&r.ctx).unwrap(); + let mut l = r.sink.local_state(&r.ctx, &g).unwrap(); + let keys: Vec = (0..100).collect(); + for _ in 0..5 { + // `Blocked` took the batch: wait, then the next one. + if let SinkStatus::Blocked(wait) = r.sink.sink(rows(&keys, None), &mut l, &g).unwrap() { + tokio::time::timeout(Duration::from_secs(5), wait) + .await + .unwrap(); + } + let (_, backlog) = super::sink::queued(&mut l).unwrap(); + assert_eq!(backlog, 0, "cuts left when the next input comes"); + } + r.sink.combine(l, &g).unwrap(); + r.sink.finalize(&g, &r.ctx).await.unwrap(); + assert_eq!(r.shared.log().len(), 501); + drop(g); + done(r); +} + +/// A blocked driver's future frames the backlog only as the task takes +/// frames: while the task is stuck it stays pending, and the cuts keep +/// their reservation. +#[tokio::test] +async fn a_blocked_future_waits_while_the_task_takes_nothing() { + let r = rig(RelationShape::Map, Yield::Nothing, limits(1, 2), 64 << 20); + r.shared.closed.store(true, Ordering::SeqCst); + let g = r.sink.global_state(&r.ctx).unwrap(); + let mut l = r.sink.local_state(&r.ctx, &g).unwrap(); + let keys: Vec = (0..10).collect(); + let SinkStatus::Blocked(wait) = r.sink.sink(rows(&keys, None), &mut l, &g).unwrap() else { + panic!("ten frames over a window of two must block"); + }; + let waiting = tokio::spawn(wait); + tokio::task::yield_now().await; + assert!(!waiting.is_finished(), "ready while the task took nothing"); + assert_eq!(super::sink::queued(&mut l).unwrap(), (vec![2], 8)); + r.shared.open_gate(); + tokio::time::timeout(Duration::from_secs(5), waiting) + .await + .unwrap() + .unwrap(); + assert_eq!(super::sink::queued(&mut l).unwrap().1, 0); + r.sink.combine(l, &g).unwrap(); + r.sink.finalize(&g, &r.ctx).await.unwrap(); + assert_eq!(r.shared.log().len(), 11); + drop(g); + done(r); +} diff --git a/crates/krishiv-exec/src/shared.rs b/crates/krishiv-exec/src/shared.rs index 3f3f42481..fb1b5fb4f 100644 --- a/crates/krishiv-exec/src/shared.rs +++ b/crates/krishiv-exec/src/shared.rs @@ -359,6 +359,12 @@ impl SharedSink { Arc::new(Self { schema, scope }) } + /// Whether a part added by [`Self::add_spilled_part`] is written out + /// (the sink has a spill scope) rather than waited for. + pub(crate) fn spills(&self) -> bool { + self.scope.is_some() + } + /// Adds `batches`, charged to `reservation`, after the parts combined so /// far: what a `PythonEval` finishes once its drivers have combined. /// Call before `finalize`. @@ -381,6 +387,51 @@ impl SharedSink { }); Ok(()) } + + /// Adds `batches` as [`Self::add_part`] does when the pool has no + /// memory for them: written to a spill file (the writer's working + /// memory from the spill scratch reserve); without a spill scope, + /// after waiting for the memory. Call before `finalize`. + /// + /// # Errors + /// A failed write, the memory refused, or a cancel. + pub(crate) async fn add_spilled_part( + &self, + global: &GlobalState, + batches: Vec, + ctx: &TaskContext, + ) -> Result<()> { + let Some(scope) = &self.scope else { + let bytes = batches + .iter() + .map(|b| b.data().get_array_memory_size()) + .sum(); + let mut reservation = ctx.consumer().reservation(); + reservation.grow(bytes).await.map_err(KrishivError::from)?; + return self.add_part(global, batches, reservation); + }; + if batches.is_empty() { + return Ok(()); + } + let g = global_ref::(global, "SharedSink")?; + let scratch = ctx.consumer().reservation().for_spill_scratch(); + let writer = scope + .writer(Arc::clone(&self.schema), scratch) + .map_err(KrishivError::from)?; + let query = Arc::clone(ctx.query()); + let file = ctx + .run_blocking(move || { + let mut writer = writer; + for b in &batches { + query.check_cancelled()?; + writer.write(b).map_err(KrishivError::from)?; + } + writer.finish().map_err(KrishivError::from) + }) + .await?; + lock(&g.parts).push(Segment::File(Arc::new(file))); + Ok(()) + } } impl Sink for SharedSink { @@ -623,8 +674,12 @@ impl Source for SharedSource { internal("a shared plan was read before it was built".to_owned()) })?; if !state.is::() { - // A `PythonEval` keeps its rows in a shared plan's list. - state = crate::python_eval::store_state(&state)?; + // A `PythonEval` or a Python relation node keeps its + // rows in a shared plan's list. + state = match crate::python_relation::sink::store_state(&state) { + Some(store) => store, + None => crate::python_eval::store_state(&state)?, + }; } let sink = global_ref::(&state, "SharedSource")?; sink.done.get().cloned().ok_or_else(|| { diff --git a/crates/krishiv-exec/tests/planner/main.rs b/crates/krishiv-exec/tests/planner/main.rs index 203ab5ece..1215b2d5b 100644 --- a/crates/krishiv-exec/tests/planner/main.rs +++ b/crates/krishiv-exec/tests/planner/main.rs @@ -31,6 +31,7 @@ mod partial_aggregate; mod physical; mod plan_stats; mod prefix_cache; +mod python_relations; mod runtime_filters; mod scan_row_drop; mod semantics; diff --git a/crates/krishiv-exec/tests/planner/python_relations.rs b/crates/krishiv-exec/tests/planner/python_relations.rs new file mode 100644 index 000000000..accf3d7aa --- /dev/null +++ b/crates/krishiv-exec/tests/planner/python_relations.rs @@ -0,0 +1,1199 @@ +// SPDX-License-Identifier: Apache-2.0 +//! The Python UDF relation nodes planned and run over a fake runtime whose +//! tasks echo or count the frames they get (design P3-T085 §3.1, §4): rows, +//! groups, cogroup sides, one row per aggregate group, output not aligned to +//! input, the refusals, and a large skewed group through the spilling sort. + +use std::collections::VecDeque; +use std::future::poll_fn; +use std::sync::Arc; +use std::task::Poll; + +use arrow_array::{ArrayRef, Int32Array, Int64Array, RecordBatch}; +use arrow_schema::{DataType as ArrowType, Field, Schema}; +use futures::future::BoxFuture; +use futures::task::AtomicWaker; +use krishiv_common::error::{ErrorClass, KrishivError}; +use krishiv_exec::python_eval::{PreparedUdf, UdfCallSpec, UdfRuntime}; +use krishiv_exec::python_relation::{Frame, PreparedRelationUdf, RelationCallSpec, UdfTask}; +use krishiv_plan::{ + AttrId, BoundColumn, Expr, ExprNode, FunctionKind, GeneratorKind, Grouping, LogicalPlan, Plan, + PlanNode, PythonEvalType, PythonUdfRef, ResolvedExpr, ResolvedFunction, +}; +use krishiv_types::DataType; + +use super::{Env, Runnable, cell, env, run}; + +/// What a fake task does, by the UDF's name. +#[derive(Clone, Copy, PartialEq, Eq, Debug)] +enum Mode { + /// Each batch back once. + Echo, + /// Each batch back twice. + Twice, + /// Each batch back eight times. + Burst, + /// One `BIGINT` per group (per partition without groups): its rows. + Count, +} + +#[derive(Debug, Default)] +struct Runtime { + /// Every frame: `B:`, `G`, `E`. + log: Arc>>, + /// Each task's own frames. + per_task: Arc>>, + /// Per `Batch` frame of a grouped kind: `::` + /// (`-` for the key of a zero-row frame). + keys: Log, + /// Sends a slow task made wait (`slow_*` UDFs). + waits: Arc, +} + +type Log = Arc>>; + +impl UdfRuntime for Runtime { + fn prepare(&self, spec: &UdfCallSpec) -> Result, KrishivError> { + if spec.name == "plus1" && spec.arg_types == [DataType::Long] { + return Ok(Arc::new(PlusOne)); + } + Err(krishiv_common::error::InternalError { + message: format!("no scalar {}", spec.name), + } + .into()) + } + + fn prepare_relation( + &self, + spec: &RelationCallSpec, + ) -> Result, KrishivError> { + let slow = spec.name.starts_with("slow_"); + let mode = match spec.name.trim_start_matches("slow_") { + "echo" => Mode::Echo, + "twice" => Mode::Twice, + "burst" => Mode::Burst, + _ => Mode::Count, + }; + assert_eq!(spec.input_names.len(), spec.input_types.len(), "{spec:?}"); + for (names, types) in spec.input_names.iter().zip(&spec.input_types) { + assert_eq!(names.len(), types.len(), "{spec:?}"); + assert!(names.len() >= spec.key_count, "{spec:?}"); + } + assert_eq!( + spec.arity + spec.key_count, + spec.input_types[0].len(), + "{spec:?}" + ); + Ok(Arc::new(Prepared { + mode, + output: Arc::clone(&spec.output_arrow), + log: Arc::clone(&self.log), + per_task: Arc::clone(&self.per_task), + keys: Arc::clone(&self.keys), + waits: slow.then(|| Arc::clone(&self.waits)), + names: spec.input_names.clone(), + key_count: spec.key_count, + })) + } +} + +/// The scalar fake `plus1(BIGINT) -> BIGINT`. +struct PlusOne; + +impl PreparedUdf for PlusOne { + fn call(&self, args: RecordBatch) -> BoxFuture<'static, Result> { + let column = args.column(0); + let values = column.as_any().downcast_ref::().unwrap(); + let out: Int64Array = values.iter().map(|v| v.map(|v| v + 1)).collect(); + Box::pin(async move { Ok(Arc::new(out) as ArrayRef) }) + } +} + +/// The fake's prepared site. +struct Prepared { + mode: Mode, + output: arrow_schema::SchemaRef, + log: Arc>>, + per_task: Arc>>, + keys: Log, + waits: Option>, + names: Vec>, + key_count: usize, +} + +impl PreparedRelationUdf for Prepared { + fn open(&self) -> Result, KrishivError> { + let own: Log = Arc::default(); + self.per_task.lock().unwrap().push(Arc::clone(&own)); + Ok(Box::new(Task { + own, + mode: self.mode, + output: Arc::clone(&self.output), + log: Arc::clone(&self.log), + keys: Arc::clone(&self.keys), + waits: self.waits.clone(), + names: self.names.clone(), + key_count: self.key_count, + state: std::sync::Mutex::new(State::default()), + ready: AtomicWaker::new(), + })) + } +} + +#[derive(Default)] +struct State { + out: VecDeque, + rows: i64, + groups: usize, + ended: bool, +} + +struct Task { + own: Log, + mode: Mode, + output: arrow_schema::SchemaRef, + log: Arc>>, + keys: Log, + waits: Option>, + names: Vec>, + key_count: usize, + state: std::sync::Mutex, + ready: AtomicWaker, +} + +impl Task { + /// Checks a frame against the spec: its columns are the side's, led by + /// the keys, whose values are the same on every row; logs them. + fn check(&self, side: u8, batch: &RecordBatch) { + let names: Vec<&str> = batch + .schema_ref() + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect(); + assert_eq!(names, self.names[usize::from(side)], "frame columns"); + if self.key_count == 0 { + return; + } + let key = |r: usize| { + (0..self.key_count) + .map(|c| cell(batch.column(c).as_ref(), r)) + .collect::>() + .join(",") + }; + let first = if batch.num_rows() == 0 { + "-".to_owned() + } else { + key(0) + }; + // Equal under the key's collation: the tests' collated keys are + // `UTF8_LCASE`, and a group's rows keep their own values. + for r in 1..batch.num_rows() { + assert_eq!( + key(r).to_lowercase(), + first.to_lowercase(), + "one key per group" + ); + } + self.keys + .lock() + .unwrap() + .push(format!("{side}:{first}:{}", batch.num_rows())); + } + + fn push(&self, entry: String) { + self.own.lock().unwrap().push(entry.clone()); + self.log.lock().unwrap().push(entry); + } + + fn count(&self, n: i64) -> RecordBatch { + RecordBatch::try_new( + Arc::clone(&self.output), + vec![Arc::new(Int64Array::from(vec![n])) as ArrayRef], + ) + .unwrap() + } +} + +impl UdfTask for Task { + fn send(&self, frame: Frame) -> BoxFuture<'_, Result<(), KrishivError>> { + Box::pin(async move { + if let Some(waits) = &self.waits { + // A slow task: each frame waits one poll, so the operator's + // window fills and the driver blocks until it is woken. + let mut first = true; + poll_fn(|cx| { + if std::mem::take(&mut first) { + waits.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + cx.waker().wake_by_ref(); + Poll::Pending + } else { + Poll::Ready(()) + } + }) + .await; + } + let mut s = self.state.lock().unwrap(); + match frame { + Frame::Batch { side, batch } => { + self.check(side, &batch); + self.push(format!("B{side}:{}", batch.num_rows())); + s.rows += i64::try_from(batch.num_rows()).unwrap(); + let back = + RecordBatch::try_new(Arc::clone(&self.output), batch.columns().to_vec()); + match self.mode { + Mode::Echo => s.out.push_back(back.unwrap()), + Mode::Twice | Mode::Burst => { + let back = back.unwrap(); + let n = if self.mode == Mode::Twice { 2 } else { 8 }; + for _ in 0..n { + s.out.push_back(back.clone()); + } + } + Mode::Count => {} + } + } + Frame::EndGroup => { + self.push("G".to_owned()); + s.groups += 1; + if self.mode == Mode::Count { + let n = std::mem::take(&mut s.rows); + let b = self.count(n); + s.out.push_back(b); + } + } + Frame::End => { + self.push("E".to_owned()); + if self.mode == Mode::Count && s.groups == 0 { + let n = s.rows; + let b = self.count(n); + s.out.push_back(b); + } + s.ended = true; + } + } + drop(s); + self.ready.wake(); + Ok(()) + }) + } + + fn next(&self) -> BoxFuture<'_, Result, KrishivError>> { + Box::pin(poll_fn(|cx| { + self.ready.register(cx.waker()); + let mut s = self.state.lock().unwrap(); + match s.out.pop_front() { + Some(b) => Poll::Ready(Ok(Some(b))), + None if s.ended => Poll::Ready(Ok(None)), + None => Poll::Pending, + } + })) + } +} + +fn udf(name: &str, eval_type: PythonEvalType) -> PythonUdfRef { + PythonUdfRef { + name: name.to_owned(), + version: 1, + eval_type, + thread_safe: true, + deterministic: true, + } +} + +fn col(c: &BoundColumn) -> Arc { + Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Column( + c.clone(), + )))) +} + +fn bound(id: u64, name: &str, data_type: DataType) -> BoundColumn { + BoundColumn { + id: AttrId(id), + name: name.to_owned(), + data_type, + nullable: true, + } +} + +/// The analyzed `sql` and its columns. +fn input(env: &Env, sql: &str) -> (Arc, Vec) { + let plan = env.analyze(sql); + let root = Arc::clone(plan.root()); + let cols = root.resolved_output().unwrap(); + (root, cols) +} + +/// `n BIGINT` (id 900). +fn counted() -> Vec { + vec![bound(900, "n", DataType::Long)] +} + +/// `node` planned and run with `runtime` (none: no runtime). +fn run_with( + env: &Env, + node: LogicalPlan, + runtime: Option>, +) -> Result, KrishivError> { + let plan = krishiv_optimizer::optimize(&Plan::try_new(node)?)?; + let physical = krishiv_exec::planner::PhysicalPlanner::plan( + &krishiv_exec::planner::DefaultPlanner, + &plan, + &env.cx(), + )?; + let mut runnable: Runnable = env.runnable(physical); + if let Some(rt) = runtime { + runnable = runnable.with_worker(|w| w.with_python_udfs(rt)); + } + run(&runnable) +} + +fn log_of(rt: &Arc) -> Vec { + rt.log.lock().unwrap().clone() +} + +/// Every cell of `batches` as `a|b` lines, sorted. +fn lines(batches: &[RecordBatch]) -> Vec { + let mut rows = Vec::new(); + for b in batches { + for r in 0..b.num_rows() { + let cells: Vec = b.columns().iter().map(|c| cell(c.as_ref(), r)).collect(); + rows.push(cells.join("|")); + } + } + rows.sort(); + rows +} + +#[test] +fn a_map_keeps_every_row() { + let env = env(); + let (input, cols) = input(&env, "SELECT a, b FROM t"); + let rt = Arc::new(Runtime::default()); + let node = LogicalPlan::new(PlanNode::MapPartitions { + input, + udf: udf("echo", PythonEvalType::MapArrowIter), + output: vec![ + bound(900, "x", cols[0].data_type.clone()), + bound(901, "y", cols[1].data_type.clone()), + ], + barrier: false, + }); + let out = run_with(&env, node, Some(Arc::clone(&rt))).unwrap(); + let rows = lines(&out); + assert_eq!(rows.len(), 10, "{rows:?}"); + assert!(rows.contains(&"NULL|30".to_owned()), "{rows:?}"); + assert!(log_of(&rt).iter().any(|e| e == "E")); +} + +#[test] +fn output_rows_are_not_aligned_to_input_rows() { + let env = env(); + let (input, cols) = input(&env, "SELECT a, b FROM t"); + let node = LogicalPlan::new(PlanNode::MapPartitions { + input, + udf: udf("twice", PythonEvalType::MapPandasIter), + output: vec![ + bound(900, "x", cols[0].data_type.clone()), + bound(901, "y", cols[1].data_type.clone()), + ], + barrier: false, + }); + let out = run_with(&env, node, Some(Arc::new(Runtime::default()))).unwrap(); + assert_eq!(lines(&out).len(), 20); +} + +#[test] +fn a_grouped_map_sees_each_key_once_with_all_its_rows() { + let env = env(); + let (input, cols) = input(&env, "SELECT a, b FROM t"); + let rt = Arc::new(Runtime::default()); + let node = LogicalPlan::new(PlanNode::GroupMap { + input, + keys: vec![col(&cols[0])], + udf: udf("count", PythonEvalType::GroupedMapPandas), + output: counted(), + }); + let out = run_with(&env, node, Some(Arc::clone(&rt))).unwrap(); + // `a` is 0..9 with 3 and 7 NULL: eight keys of one row, NULL of two. + let rows = lines(&out); + assert_eq!(rows.len(), 9, "{rows:?}"); + assert_eq!(rows.iter().filter(|r| *r == "1").count(), 8, "{rows:?}"); + assert!(rows.contains(&"2".to_owned()), "{rows:?}"); + let log = log_of(&rt); + assert_eq!(log.iter().filter(|e| *e == "G").count(), 9, "{log:?}"); +} + +/// A slow task through the real driver: the window fills, the driver +/// pauses on `Blocked` and is woken, and every row arrives once. +#[test] +fn a_slow_task_blocks_the_driver_and_loses_nothing() { + let env = env(); + let (input, cols) = input(&env, "SELECT a, b FROM t"); + let rt = Arc::new(Runtime::default()); + let node = LogicalPlan::new(PlanNode::GroupMap { + input, + keys: vec![col(&cols[0])], + udf: udf("slow_count", PythonEvalType::GroupedMapPandas), + output: counted(), + }); + let out = run_with(&env, node, Some(Arc::clone(&rt))).unwrap(); + let rows = lines(&out); + assert_eq!(rows.len(), 9, "{rows:?}"); + assert_eq!(rows.iter().filter(|r| *r == "1").count(), 8, "{rows:?}"); + let frames = log_of(&rt).len(); + assert_eq!( + rt.waits.load(std::sync::atomic::Ordering::SeqCst), + frames, + "every frame waited" + ); +} + +#[test] +fn a_grouped_map_without_keys_is_one_group() { + let env = env(); + let (input, _) = input(&env, "SELECT a, b FROM t"); + let node = LogicalPlan::new(PlanNode::GroupMap { + input, + keys: Vec::new(), + udf: udf("count", PythonEvalType::GroupedMapArrow), + output: counted(), + }); + let out = run_with(&env, node, Some(Arc::new(Runtime::default()))).unwrap(); + assert_eq!(lines(&out), ["10"]); +} + +#[test] +fn a_cogroup_gives_each_key_both_sides_and_a_one_sided_key_its_side() { + let env = env(); + let (left, lc) = input(&env, "SELECT a, b FROM t"); + let (right, rc) = input(&env, "SELECT a AS c, s FROM t WHERE a < 2"); + let rt = Arc::new(Runtime::default()); + let node = LogicalPlan::new(PlanNode::CoGroupMap { + left, + right, + left_keys: vec![col(&lc[0])], + right_keys: vec![col(&rc[0])], + udf: udf("count", PythonEvalType::CogroupedMapArrow), + output: counted(), + }); + let out = run_with(&env, node, Some(Arc::clone(&rt))).unwrap(); + // Keys 0 and 1 have a row on each side; the others the left one only. + let rows = lines(&out); + assert_eq!(rows.len(), 9, "{rows:?}"); + assert_eq!(rows.iter().filter(|r| *r == "2").count(), 3, "{rows:?}"); + let log = log_of(&rt); + let count = |e: &str| log.iter().filter(|x| *x == e).count(); + assert_eq!(count("B1:1"), 2, "{log:?}"); + // The seven keys on the left only get a zero-row right side. + assert_eq!(count("B1:0"), 7, "{log:?}"); + assert_eq!(count("B0:0"), 0, "{log:?}"); + // Within each task, every group is a left frame, a right frame, then + // `EndGroup`. + for own in rt.per_task.lock().unwrap().iter() { + let own = own.lock().unwrap(); + for group in own + .split(|e| e == "G") + .filter(|g| *g != ["E"] && !g.is_empty()) + { + assert!(group[0].starts_with("B0"), "{own:?}"); + assert!(group.last().unwrap().starts_with("B1"), "{own:?}"); + } + } + // The frames lead with the key: the right side's keys are 0 and 1. + let keys = rt.keys.lock().unwrap().clone(); + let mut right: Vec<&String> = keys + .iter() + .filter(|k| k.starts_with("1:") && !k.ends_with(":0")) + .collect(); + right.sort(); + assert_eq!(right, ["1:0:1", "1:1:1"], "{keys:?}"); +} + +#[test] +fn a_computed_grouping_key_leads_every_frame() { + let env = env(); + // `groupBy(df.a % 2)`: the key is computed, not an input column. + let analyzed = env.analyze("SELECT a, b, a % 2 AS m FROM t"); + let root = analyzed.root(); + let PlanNode::Project { input, exprs } = &root.node else { + panic!("{root}") + }; + let key = match &exprs[2].node { + ExprNode::Resolved(ResolvedExpr::Alias { expr, .. }) => Arc::clone(expr), + other => panic!("{other:?}"), + }; + let rt = Arc::new(Runtime::default()); + let node = LogicalPlan::new(PlanNode::GroupMap { + input: Arc::clone(input), + keys: vec![key], + udf: udf("count", PythonEvalType::GroupedMapPandas), + output: counted(), + }); + let out = run_with(&env, node, Some(Arc::clone(&rt))).unwrap(); + // `a` is 0..9 with 3 and 7 NULL: five even, three odd, two NULL. + assert_eq!(lines(&out), ["2", "3", "5"]); + let mut keys = rt.keys.lock().unwrap().clone(); + keys.sort(); + assert_eq!(keys, ["0:0:5", "0:1:3", "0:NULL:2"]); +} + +/// `groupBy(k)` of a `UTF8_LCASE` key: `K` and `k` are one group (Spark's +/// `GroupedIterator` compares keys by the collation), each row keeping its +/// own value. +#[test] +fn a_grouped_map_groups_a_collated_key_by_its_collation() { + let env = env(); + let (input, cols) = input( + &env, + "SELECT CASE WHEN b < 50 THEN 'K' ELSE 'k' END COLLATE UTF8_LCASE AS k, b FROM t", + ); + let rt = Arc::new(Runtime::default()); + let node = LogicalPlan::new(PlanNode::GroupMap { + input, + keys: vec![col(&cols[0])], + udf: udf("count", PythonEvalType::GroupedMapPandas), + output: counted(), + }); + let out = run_with(&env, node, Some(Arc::clone(&rt))).unwrap(); + assert_eq!(lines(&out), ["10"]); + // The frames carry the rows' own key values, not the collation key. + let keys = rt.keys.lock().unwrap().clone(); + assert!( + keys.iter() + .all(|k| k.starts_with("0:K:") || k.starts_with("0:k:")), + "{keys:?}" + ); +} + +/// A cogroup on `UTF8_LCASE` keys: the left side's `K` and the right side's +/// `k` are one key (Spark's `CoGroupedIterator` compares them by the +/// collation), so the UDF gets one call with both sides. +#[test] +fn a_cogroup_matches_collated_keys_by_their_collation() { + let env = env(); + let (left, lc) = input( + &env, + "SELECT 'K' COLLATE UTF8_LCASE AS k, b FROM t WHERE b < 30", + ); + let (right, rc) = input( + &env, + "SELECT 'k' COLLATE UTF8_LCASE AS c, s FROM t WHERE b >= 80", + ); + let rt = Arc::new(Runtime::default()); + let node = LogicalPlan::new(PlanNode::CoGroupMap { + left, + right, + left_keys: vec![col(&lc[0])], + right_keys: vec![col(&rc[0])], + udf: udf("count", PythonEvalType::CogroupedMapArrow), + output: counted(), + }); + let out = run_with(&env, node, Some(Arc::clone(&rt))).unwrap(); + assert_eq!(lines(&out), ["5"]); + let log = log_of(&rt); + assert_eq!(log.iter().filter(|e| *e == "G").count(), 1, "{log:?}"); + let mut keys = rt.keys.lock().unwrap().clone(); + keys.sort(); + assert_eq!(keys, ["0:K:3", "1:k:2"], "{keys:?}"); +} + +/// `groupBy(plus1(b)).applyInPandas`: the scalar call in the key runs in a +/// `PythonEval` below the node (Spark evaluates it in a `Project` under +/// `FlatMapGroupsInPandas`, which `ExtractPythonUDFs` at v4.2.0 turns into +/// an `ArrowEvalPython`); the frames lead with its values. +#[test] +fn a_scalar_python_call_in_a_grouped_maps_key_runs_below_it() { + let env = env(); + let (input, cols) = input(&env, "SELECT a, b FROM t"); + let rt = Arc::new(Runtime::default()); + let node = LogicalPlan::new(PlanNode::GroupMap { + input, + keys: vec![plus1(col(&cols[1]))], + udf: udf("count", PythonEvalType::GroupedMapPandas), + output: counted(), + }); + let out = run_with(&env, node, Some(Arc::clone(&rt))).unwrap(); + // `b` is 0, 10, ..., 90: ten keys of one row each. + assert_eq!(lines(&out), vec!["1"; 10]); + let keys = rt.keys.lock().unwrap().clone(); + assert!(keys.contains(&"0:1:1".to_owned()), "{keys:?}"); + assert!(keys.contains(&"0:91:1".to_owned()), "{keys:?}"); +} + +/// A cogroup keyed by `plus1(b)` on both sides: each side's call runs +/// below its branch, and equal results meet. +#[test] +fn a_scalar_python_call_in_a_cogroups_keys_runs_below_it() { + let env = env(); + let (left, lc) = input(&env, "SELECT a, b FROM t"); + let (right, rc) = input(&env, "SELECT b AS c, s FROM t WHERE b < 20"); + let rt = Arc::new(Runtime::default()); + let node = LogicalPlan::new(PlanNode::CoGroupMap { + left, + right, + left_keys: vec![plus1(col(&lc[1]))], + right_keys: vec![plus1(col(&rc[0]))], + udf: udf("count", PythonEvalType::CogroupedMapArrow), + output: counted(), + }); + let out = run_with(&env, node, Some(Arc::clone(&rt))).unwrap(); + // Keys 1 and 11 have a row on each side; the other eight the left one. + let rows = lines(&out); + assert_eq!(rows.len(), 10, "{rows:?}"); + assert_eq!(rows.iter().filter(|r| *r == "2").count(), 2, "{rows:?}"); + let keys = rt.keys.lock().unwrap().clone(); + let mut right: Vec<&String> = keys + .iter() + .filter(|k| k.starts_with("1:") && !k.ends_with(":0")) + .collect(); + right.sort(); + assert_eq!(right, ["1:11:1", "1:1:1"], "{keys:?}"); +} + +/// A grouped-aggregate Python UDF call over `arg`. +fn udaf(arg: &BoundColumn) -> Arc { + udaf_named("count", arg) +} + +fn udaf_named(name: &str, arg: &BoundColumn) -> Arc { + Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Function( + ResolvedFunction { + id: udf(name, PythonEvalType::GroupedAggArrow).function_id(), + kind: FunctionKind::Aggregate, + args: vec![col(arg)], + distinct: false, + filter: None, + within_group: Vec::new(), + data_type: DataType::Long, + nullable: true, + deterministic: true, + }, + )))) +} + +fn alias(e: Arc, id: u64, name: &str) -> Arc { + Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Alias { + expr: e, + output: bound(id, name, DataType::Long), + }))) +} + +#[test] +fn a_python_aggregate_yields_one_row_per_group() { + let env = env(); + let (input, cols) = input(&env, "SELECT a, b FROM t"); + let node = LogicalPlan::new(PlanNode::Aggregate { + input, + grouping: Grouping::GroupBy(vec![col(&cols[0])]), + aggregates: vec![col(&cols[0]), alias(udaf(&cols[1]), 950, "n")], + }); + let out = run_with(&env, node, Some(Arc::new(Runtime::default()))).unwrap(); + let rows = lines(&out); + assert_eq!(rows.len(), 9, "{rows:?}"); + assert!(rows.contains(&"NULL|2".to_owned()), "{rows:?}"); + assert!(rows.contains(&"5|1".to_owned()), "{rows:?}"); +} + +/// `SELECT k, udaf(b) ... GROUP BY k` of a `UTF8_LCASE` key: `K` and `k` +/// are one group (Spark's `ArrowAggregatePythonExec` cuts groups with a +/// `GroupedIterator` comparing keys by the collation, v4.2.0), and the key +/// shown is a value of the group's own (its first row's), not the +/// comparison key. The collation rewrite leaves the aggregate to the Python +/// aggregate's planner, so no `any_value(k)` sits beside the UDAF. +#[test] +fn a_python_aggregate_groups_a_collated_key_by_its_collation() { + let env = env(); + let (input, cols) = input( + &env, + "SELECT CASE WHEN b < 50 THEN 'K' ELSE 'k' END COLLATE UTF8_LCASE AS k, b FROM t", + ); + let node = LogicalPlan::new(PlanNode::Aggregate { + input, + grouping: Grouping::GroupBy(vec![col(&cols[0])]), + aggregates: vec![col(&cols[0]), alias(udaf(&cols[1]), 950, "n")], + }); + let out = run_with(&env, node, Some(Arc::new(Runtime::default()))).unwrap(); + let rows = lines(&out); + assert_eq!(rows.len(), 1, "{rows:?}"); + assert!(rows[0] == "K|10" || rows[0] == "k|10", "{rows:?}"); + let key_type = out[0].schema().field(0).data_type().clone(); + assert!( + matches!(key_type, ArrowType::Utf8 | ArrowType::Utf8View), + "the key column is the key's own, not its comparison key: {key_type}" + ); +} + +/// A scalar Python UDF call `plus1(arg)`. +fn plus1(arg: Arc) -> Arc { + Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Function( + ResolvedFunction { + id: udf("plus1", PythonEvalType::ScalarArrow).function_id(), + kind: FunctionKind::Scalar, + args: vec![arg], + distinct: false, + filter: None, + within_group: Vec::new(), + data_type: DataType::Long, + nullable: true, + deterministic: true, + }, + )))) +} + +/// `udaf(arg)` with `arg` any expression. +fn udaf_of(arg: Arc) -> Arc { + let ExprNode::Resolved(ResolvedExpr::Function(f)) = + &udaf_named("count", &bound(0, "x", DataType::Long)).node + else { + unreachable!() + }; + Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Function( + ResolvedFunction { + args: vec![arg], + ..f.clone() + }, + )))) +} + +/// `udaf(plus1(b))`: the scalar call in the aggregate's argument runs +/// below it (Spark's `ExtractPythonUDFs` at v4.2.0 evaluates it in an +/// `ArrowEvalPython` under the aggregate). +#[test] +fn a_scalar_python_call_in_a_python_aggregates_argument_runs_below_it() { + let env = env(); + let (input, cols) = input(&env, "SELECT a, b FROM t"); + let node = LogicalPlan::new(PlanNode::Aggregate { + input, + grouping: Grouping::GroupBy(vec![col(&cols[0])]), + aggregates: vec![ + col(&cols[0]), + alias(udaf_of(plus1(col(&cols[1]))), 950, "n"), + ], + }); + let out = run_with(&env, node, Some(Arc::new(Runtime::default()))).unwrap(); + let rows = lines(&out); + assert_eq!(rows.len(), 9, "{rows:?}"); + assert!(rows.contains(&"NULL|2".to_owned()), "{rows:?}"); +} + +/// `GROUP BY plus1(b)`: the scalar call in the key runs below the aggregate +/// (Spark's `ExtractGroupingPythonUDFFromAggregate` at v4.2.0), and the +/// output reads the key's column. +#[test] +fn a_scalar_python_call_in_a_python_aggregates_key_runs_below_it() { + let env = env(); + let (input, cols) = input(&env, "SELECT a, b FROM t"); + let key = plus1(col(&cols[1])); + let node = LogicalPlan::new(PlanNode::Aggregate { + input, + grouping: Grouping::GroupBy(vec![Arc::clone(&key)]), + aggregates: vec![alias(key, 949, "k"), alias(udaf(&cols[0]), 950, "n")], + }); + let out = run_with(&env, node, Some(Arc::new(Runtime::default()))).unwrap(); + let rows = lines(&out); + // `b` is 0, 10, ..., 90: ten keys of one row each. + assert_eq!(rows.len(), 10, "{rows:?}"); + assert!(rows.contains(&"1|1".to_owned()), "{rows:?}"); + assert!(rows.contains(&"91|1".to_owned()), "{rows:?}"); +} + +/// `plus1(udaf(b))`: the scalar call over the aggregate's result runs above +/// it (Spark's `ExtractPythonUDFFromAggregate` at v4.2.0). +#[test] +fn a_scalar_python_call_over_a_python_aggregate_runs_above_it() { + let env = env(); + let (input, cols) = input(&env, "SELECT a, b FROM t"); + let node = LogicalPlan::new(PlanNode::Aggregate { + input, + grouping: Grouping::GroupBy(vec![col(&cols[0])]), + aggregates: vec![col(&cols[0]), alias(plus1(udaf(&cols[1])), 950, "n")], + }); + let out = run_with(&env, node, Some(Arc::new(Runtime::default()))).unwrap(); + let rows = lines(&out); + // Eight keys of one row (1 + 1), NULL of two (2 + 1). + assert_eq!(rows.len(), 9, "{rows:?}"); + assert!(rows.contains(&"NULL|3".to_owned()), "{rows:?}"); + assert!(rows.contains(&"5|2".to_owned()), "{rows:?}"); +} + +/// A keyless Python aggregate over no rows yields no row: Spark's +/// `ArrowAggregatePythonExec` (L156-157 at v4.2.0) returns an empty +/// partition as it is, without calling the UDF, unlike a built-in global +/// aggregate's one row. +#[test] +fn a_keyless_python_aggregate_over_no_rows_yields_no_row() { + let env = env(); + let (input, cols) = input(&env, "SELECT a, b FROM t WHERE b < 0"); + let rt = Arc::new(Runtime::default()); + let node = LogicalPlan::new(PlanNode::Aggregate { + input, + grouping: Grouping::GroupBy(Vec::new()), + aggregates: vec![alias(udaf(&cols[1]), 950, "n")], + }); + let out = run_with(&env, node, Some(Arc::clone(&rt))).unwrap(); + assert!(lines(&out).is_empty(), "{:?}", lines(&out)); + assert!(log_of(&rt).is_empty(), "{:?}", log_of(&rt)); +} + +/// Two Python aggregates, one slow: each key's values pair up. +#[test] +fn two_python_aggregates_pair_their_values_per_key() { + let env = env(); + let (input, cols) = input(&env, "SELECT a, b FROM t"); + let node = LogicalPlan::new(PlanNode::Aggregate { + input, + grouping: Grouping::GroupBy(vec![col(&cols[0])]), + aggregates: vec![ + col(&cols[0]), + alias(udaf(&cols[1]), 950, "n"), + alias(udaf_named("slow_count", &cols[0]), 951, "m"), + ], + }); + let out = run_with(&env, node, Some(Arc::new(Runtime::default()))).unwrap(); + let rows = lines(&out); + assert_eq!(rows.len(), 9, "{rows:?}"); + for row in &rows { + let cells: Vec<&str> = row.split('|').collect(); + assert_eq!(cells[1], cells[2], "{rows:?}"); + } + assert!(rows.contains(&"NULL|2|2".to_owned()), "{rows:?}"); + assert!(rows.contains(&"5|1|1".to_owned()), "{rows:?}"); +} + +#[test] +fn a_python_aggregate_beside_a_builtin_is_refused() { + let env = env(); + let analyzed = env.analyze("SELECT a, sum(b) AS s FROM t GROUP BY a"); + let root = analyzed.root(); + let PlanNode::Aggregate { + input, + grouping, + aggregates, + } = &root.node + else { + panic!("{root}") + }; + let cols = input.resolved_output().unwrap(); + let mut aggregates = aggregates.clone(); + aggregates.push(alias(udaf(&cols[1]), 950, "n")); + let node = LogicalPlan::new(PlanNode::Aggregate { + input: Arc::clone(input), + grouping: grouping.clone(), + aggregates, + }); + // Raised at physical planning, as Spark does (`SparkStrategies`): the + // plan itself is valid and optimizes. + let plan = krishiv_optimizer::optimize(&Plan::try_new(node).unwrap()).unwrap(); + let err = krishiv_exec::planner::PhysicalPlanner::plan( + &krishiv_exec::planner::DefaultPlanner, + &plan, + &env.cx(), + ) + .unwrap_err(); + assert_eq!( + err.error_class(), + ErrorClass::InvalidPandasUdfPlacement, + "{err}" + ); + assert!(err.to_string().contains("`count`"), "{err}"); +} + +/// A Python UDAF call with `DISTINCT`, `FILTER (WHERE ...)` or `WITHIN +/// GROUP` would run over every row of its group: refused at planning. +#[test] +fn a_python_aggregate_with_distinct_filter_or_within_group_is_refused() { + let env = env(); + let (input, cols) = input(&env, "SELECT a, b FROM t"); + let changed = |change: &dyn Fn(&mut ResolvedFunction)| { + let ExprNode::Resolved(ResolvedExpr::Function(f)) = &udaf(&cols[1]).node else { + unreachable!() + }; + let mut f = f.clone(); + change(&mut f); + Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Function(f)))) + }; + let cases: [(&str, Arc); 3] = [ + ("DISTINCT", changed(&|f| f.distinct = true)), + ("FILTER", changed(&|f| f.filter = Some(col(&cols[0])))), + ( + "WITHIN GROUP", + changed(&|f| f.within_group = vec![krishiv_plan::SortOrder::asc(col(&cols[0]))]), + ), + ]; + for (what, call) in cases { + let node = LogicalPlan::new(PlanNode::Aggregate { + input: Arc::clone(&input), + grouping: Grouping::GroupBy(vec![col(&cols[0])]), + aggregates: vec![col(&cols[0]), alias(call, 950, "n")], + }); + let err = run_with(&env, node, Some(Arc::new(Runtime::default()))).unwrap_err(); + assert_eq!( + err.error_class(), + ErrorClass::KrishivUnsupportedFeature, + "{what}: {err}" + ); + assert!(err.to_string().contains(what), "{what}: {err}"); + } +} + +#[test] +fn barrier_mode_is_refused_at_planning() { + let env = env(); + let (input, cols) = input(&env, "SELECT a, b FROM t"); + let node = LogicalPlan::new(PlanNode::MapPartitions { + input, + udf: udf("echo", PythonEvalType::MapArrowIter), + output: vec![bound(900, "x", cols[0].data_type.clone())], + barrier: true, + }); + let err = run_with(&env, node, Some(Arc::new(Runtime::default()))).unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::KrishivUnsupportedFeature); + assert!(err.to_string().contains("isBarrier"), "{err}"); +} + +#[test] +fn python_udtfs_and_window_aggregates_are_refused_until_pr_26c() { + let env = env(); + let (input, cols) = input(&env, "SELECT a, b FROM t"); + let call = |eval_type, kind| { + Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Function( + ResolvedFunction { + id: udf("f", eval_type).function_id(), + kind, + args: vec![col(&cols[1])], + distinct: false, + filter: None, + within_group: Vec::new(), + data_type: DataType::Long, + nullable: true, + deterministic: true, + }, + )))) + }; + let lateral = LogicalPlan::new(PlanNode::Generate { + input: Arc::clone(&input), + kind: GeneratorKind::PythonUdtf, + outer: false, + arg: call(PythonEvalType::TableUdf, FunctionKind::Generator), + output: counted(), + names: Vec::new(), + qualifier: None, + }); + let table = LogicalPlan::new(PlanNode::PythonTableFunction { + args: Vec::new(), + table_arg: None, + udf: udf("f", PythonEvalType::ArrowUdtf), + output: counted(), + }); + let window = LogicalPlan::new(PlanNode::Window { + input, + exprs: vec![alias( + call(PythonEvalType::WindowAggArrow, FunctionKind::Aggregate), + 950, + "w", + )], + partition_by: vec![col(&cols[0])], + order_by: Vec::new(), + }); + for node in [lateral, table, window] { + let text = node.to_string(); + let err = run_with(&env, node, Some(Arc::new(Runtime::default()))).unwrap_err(); + assert_eq!( + err.error_class(), + ErrorClass::KrishivUnsupportedFeature, + "{text}: {err}" + ); + assert!(err.to_string().contains("PR-26c"), "{text}: {err}"); + } +} + +#[test] +fn a_scalar_python_udf_in_a_window_is_not_named_a_window_aggregate() { + let env = env(); + let (input, cols) = input(&env, "SELECT a, b FROM t"); + let scalar = Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Function( + ResolvedFunction { + id: udf("f", PythonEvalType::ScalarArrow).function_id(), + kind: FunctionKind::Scalar, + args: vec![col(&cols[1])], + distinct: false, + filter: None, + within_group: Vec::new(), + data_type: DataType::Long, + nullable: true, + deterministic: true, + }, + )))); + let window = LogicalPlan::new(PlanNode::Window { + input, + exprs: vec![alias(scalar, 950, "w")], + partition_by: vec![col(&cols[0])], + order_by: Vec::new(), + }); + let err = run_with(&env, window, Some(Arc::new(Runtime::default()))).unwrap_err(); + assert!(!err.to_string().contains("PR-26c"), "{err}"); +} + +#[test] +fn a_worker_without_a_runtime_refuses_the_plan() { + let env = env(); + let (input, _) = input(&env, "SELECT a, b FROM t"); + let node = LogicalPlan::new(PlanNode::GroupMap { + input, + keys: Vec::new(), + udf: udf("count", PythonEvalType::GroupedMapArrow), + output: counted(), + }); + let err = run_with(&env, node, None).unwrap_err(); + assert_eq!( + err.error_class(), + ErrorClass::KrishivUdfNotAvailable, + "{err}" + ); +} + +/// A table `big (k INT, v BIGINT)` of `rows` rows keyed by `key`, `node` +/// over its scan planned and run under a 4 MiB limit with a spill +/// directory: the results, and the spill files written. +fn run_big( + rows: i32, + key: impl Fn(i32) -> i32, + node: impl FnOnce(Arc, &BoundColumn) -> LogicalPlan, +) -> (Vec, u64) { + use krishiv_exec::planner::{MemoryTable, MemoryTables}; + use krishiv_memory::{MemoryConfig, SpillConfig, SpillManager, TreePool}; + use krishiv_types::batch::Batch; + + let schema = Arc::new(Schema::new(vec![ + Field::new("k", ArrowType::Int32, false), + Field::new("v", ArrowType::Int64, false), + ])); + let batches: Vec = (0..rows) + .step_by(8192) + .map(|start| { + let end = (start + 8192).min(rows); + let k: Int32Array = (start..end).map(&key).collect(); + let v: Int64Array = (start..end).map(i64::from).collect(); + Batch::new( + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(k), Arc::new(v)]).unwrap(), + ) + }) + .collect(); + let tables = Arc::new(MemoryTables::new().with_table( + &["spark_catalog", "default", "big"], + Arc::new(MemoryTable::new(Arc::clone(&schema), batches)), + )); + let k = bound(1, "k", DataType::Integer); + let v = bound(2, "v", DataType::Long); + let scan = LogicalPlan::new(PlanNode::Scan(krishiv_plan::Scan { + source: krishiv_plan::ReadSource::Table { + name: vec![ + "spark_catalog".to_owned(), + "default".to_owned(), + "big".to_owned(), + ], + options: std::collections::BTreeMap::new(), + }, + output: vec![k.clone(), v], + version: None, + })); + let plan = Plan::try_new(node(Arc::new(scan), &k)).unwrap(); + let cx = krishiv_exec::planner::PlanningContext::new( + Arc::clone(&tables) as _, + Arc::new(krishiv_common::config::Config::new()), + ); + let physical = krishiv_exec::planner::PhysicalPlanner::plan( + &krishiv_exec::planner::DefaultPlanner, + &plan, + &cx, + ) + .unwrap(); + let dir = tempfile::tempdir().unwrap(); + let scope = SpillManager::open(&SpillConfig { + dir: dir.path().join("spill"), + ..SpillConfig::default() + }) + .unwrap() + .scope() + .unwrap(); + let runnable = Runnable::new(physical, tables as _) + .with_spill(scope) + .with_worker(|w| w.with_python_udfs(Arc::new(Runtime::default()))); + let root = TreePool::new(&MemoryConfig { + limit: Some(4 << 20), + ..MemoryConfig::default() + }); + let (out, query) = super::run_in(&runnable, &root); + let out = out.unwrap(); + let files: u64 = query + .metrics() + .operators + .values() + .map(|o| o.metrics.spill_files) + .sum(); + drop(query); + assert_eq!(krishiv_memory::MemoryPool::reserved(&*root), 0); + (out, files) +} + +#[test] +fn a_skewed_group_runs_through_the_spilling_sort_in_bounded_memory() { + const ROWS: i32 = 300_000; + // One key of 270k rows (90%), the rest in keys of 100 rows; the + // brief's 1M/900k shape shrunk to stay near 10 s in debug builds. + let (out, files) = run_big( + ROWS, + |i| if i < 270_000 { 0 } else { 1 + i / 100 }, + |scan, k| { + LogicalPlan::new(PlanNode::GroupMap { + input: scan, + keys: vec![col(k)], + udf: udf("count", PythonEvalType::GroupedMapArrow), + output: counted(), + }) + }, + ); + let mut counts: Vec = out + .iter() + .flat_map(|b| { + b.column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .to_vec() + }) + .collect(); + counts.sort_unstable(); + assert_eq!(counts.len(), 301); + assert_eq!(counts.last(), Some(&270_000)); + assert_eq!(counts.iter().sum::(), i64::from(ROWS)); + assert!(files > 0, "the sort did not spill"); +} + +/// A task yielding eight times its input under a 4 MiB limit: its results +/// are reserved as they arrive, so the store spills them and the task waits +/// for memory instead of piling up results nobody accounts for. +#[test] +fn a_task_yielding_far_more_than_its_input_runs_in_bounded_memory() { + const ROWS: i32 = 100_000; + let (out, files) = run_big( + ROWS, + |i| i % 7, + |scan, k| { + LogicalPlan::new(PlanNode::MapPartitions { + input: scan, + udf: udf("burst", PythonEvalType::MapArrowIter), + output: vec![ + bound(900, "k", k.data_type.clone()), + bound(901, "v", DataType::Long), + ], + barrier: false, + }) + }, + ); + let rows: usize = out.iter().map(RecordBatch::num_rows).sum(); + assert_eq!(rows, 8 * ROWS as usize); + assert!(files > 0, "the results did not spill"); +} diff --git a/crates/krishiv-optimizer/src/collation_keys.rs b/crates/krishiv-optimizer/src/collation_keys.rs index 233e5706b..f5b2566c1 100644 --- a/crates/krishiv-optimizer/src/collation_keys.rs +++ b/crates/krishiv-optimizer/src/collation_keys.rs @@ -5,7 +5,8 @@ //! key column reaches an output schema; a second run finds nothing to do (a key is `BINARY`). //! //! - **Aggregates** group by the key and output the first value of each group (Spark's -//! `RewriteGroupByCollation`, `First(c)`). +//! `RewriteGroupByCollation`, `First(c)`). An aggregate of grouped-aggregate Python UDFs is left +//! as it is: its physical planner groups by the key itself and shows each group's first value. //! - **`min`/`max`** of a collated string take the key as a second argument, which the executor //! runs as a by-key extremum that keeps the first of equal keys (as Spark's `Max` keeps its //! buffer on a tie; `max_by`, which takes the last, keeps its own form with a keyed ordering). @@ -24,7 +25,8 @@ use krishiv_functions::registry::{Lookup, builtins}; use krishiv_functions::signature::SparkProfile; use krishiv_plan::{ BoundColumn, Expr, ExprNode, FunctionId, FunctionKind, Grouping, Literal, LogicalPlan, - PlanNode, ResolvedExpr, ResolvedFunction, SortOrder, WindowExpr, WindowSpec, WindowSpecRef, + PlanNode, PythonUdfKind, ResolvedExpr, ResolvedFunction, SortOrder, WindowExpr, WindowSpec, + WindowSpecRef, }; use krishiv_types::DataType; @@ -91,7 +93,13 @@ fn key(e: &Arc) -> Res> { )) } -fn key_if_collated(e: &Arc) -> Res> { +/// `e`'s comparison key if it is collated (`__collation_key(e, 'type of e')`), else `e`: what an +/// operator outside this rewrite that compares `e` by bytes groups, sorts and hashes by (the +/// Python UDF relations' grouping keys, whose frames keep `e` itself). +/// +/// # Errors +/// `INTERNAL_ERROR` if `__collation_key` is not registered. +pub fn key_if_collated(e: &Arc) -> Res> { if collated(e) { key(e) } else { @@ -239,6 +247,14 @@ fn calls<'a>(e: &'a Expr, out: &mut Vec<&'a ResolvedFunction>) { e.for_each_child(&mut |c| calls(c, out)); } +/// Whether `f` calls a grouped-aggregate Python UDF. +fn python_udaf(f: &ResolvedFunction) -> bool { + matches!( + &f.id, + FunctionId::PythonUdf { eval_type, .. } if eval_type.kind() == PythonUdfKind::GroupedAgg + ) +} + fn alias(expr: Arc, output: BoundColumn) -> Arc { let origin = expr.origin; Arc::new( @@ -436,12 +452,18 @@ fn aggregate(p: &Arc, next: &mut u64) -> Res> { else { return Ok(Arc::clone(p)); }; + let mut found = Vec::new(); + aggregates.iter().for_each(|a| calls(a, &mut found)); + // A grouped-aggregate Python UDF's aggregate keys a collated key itself + // (the physical planner's `python_aggregate`): an `any_value(k)` beside + // the UDF would be a native aggregate it cannot run with. + if found.iter().any(|f| python_udaf(f)) { + return Ok(Arc::clone(p)); + } if let Some(done) = distinct_over_keys(p, input, keys, aggregates, next)? { return Ok(done); } let held: Vec> = keys.iter().filter(|k| collated(k)).cloned().collect(); - let mut found = Vec::new(); - aggregates.iter().for_each(|a| calls(a, &mut found)); let mut changes = !held.is_empty(); for f in found { changes |= keyed_call(f, p.origin)?.is_some(); diff --git a/crates/krishiv-optimizer/src/inline_cte.rs b/crates/krishiv-optimizer/src/inline_cte.rs index c9acec904..be719ba01 100644 --- a/crates/krishiv-optimizer/src/inline_cte.rs +++ b/crates/krishiv-optimizer/src/inline_cte.rs @@ -172,8 +172,8 @@ fn has_outer(plan: &Arc) -> bool { found } -/// Whether every expression of `plan` (subquery plans included) is -/// deterministic. +/// Whether every expression of `plan` (subquery plans included) and every +/// Python relation UDF it runs is deterministic. pub(crate) fn deterministic(plan: &Arc) -> bool { let mut seen = HashSet::new(); let mut stack = vec![Arc::clone(plan)]; @@ -185,6 +185,7 @@ pub(crate) fn deterministic(plan: &Arc) -> bool { if !seen.insert(Arc::as_ptr(&p)) { continue; } + ok &= krishiv_plan::python_relation_is_deterministic(&p.node); p.for_each_expr(&mut |e| ok &= is_deterministic(e)); p.for_each_child(&mut |c| stack.push(Arc::clone(c))); } diff --git a/crates/krishiv-optimizer/src/lib.rs b/crates/krishiv-optimizer/src/lib.rs index 54378779d..7d161af46 100644 --- a/crates/krishiv-optimizer/src/lib.rs +++ b/crates/krishiv-optimizer/src/lib.rs @@ -91,6 +91,7 @@ use krishiv_plan::{ }; use krishiv_udf_sdk::UdfLibrary; +pub use collation_keys::key_if_collated as collation_key; pub use group_code::{could_code_group_keys, is_group_code_dictionary}; pub use group_limit::{ DEFAULT_WINDOW_GROUP_LIMIT_THRESHOLD, RankLikeFunction, WindowGroupLimit, @@ -908,7 +909,8 @@ pub(crate) fn groups_by_columns(g: &Grouping) -> bool { } } -/// Whether every expression of `plan` (subquery plans included) is deterministic: a shared plan that is not is +/// Whether every expression of `plan` (subquery plans included) and every Python relation UDF it runs is +/// deterministic: a shared plan that is not is /// run again, with its readers, when its output is lost (P2-T198). #[must_use] pub fn plan_is_deterministic(plan: &Arc) -> bool { diff --git a/crates/krishiv-optimizer/src/prune.rs b/crates/krishiv-optimizer/src/prune.rs index d70796205..f40f6bbe9 100644 --- a/crates/krishiv-optimizer/src/prune.rs +++ b/crates/krishiv-optimizer/src/prune.rs @@ -156,6 +156,18 @@ impl Graph { child(recursive, None); } PlanNode::Window { input, .. } => child(input, Some(window_need(p, req))), + // A Python UDF relation sees every column of its inputs + // (`group_ops.py` L305 at v4.2.0; design P3-T085 §3.1). + PlanNode::MapPartitions { input, .. } | PlanNode::GroupMap { input, .. } => { + child(input, None); + } + PlanNode::CoGroupMap { left, right, .. } => { + child(left, None); + child(right, None); + } + PlanNode::PythonTableFunction { + table_arg: Some(t), .. + } => child(&t.input, None), _ => {} } subqueries(p, &mut |kind, s| { diff --git a/crates/krishiv-optimizer/tests/it/main.rs b/crates/krishiv-optimizer/tests/it/main.rs index 0d35fe490..522951ad3 100644 --- a/crates/krishiv-optimizer/tests/it/main.rs +++ b/crates/krishiv-optimizer/tests/it/main.rs @@ -16,6 +16,7 @@ mod join_order; mod oracle; mod plan_stats; mod props; +mod python_relations; mod qerror; mod rules; mod share_subtree; diff --git a/crates/krishiv-optimizer/tests/it/python_relations.rs b/crates/krishiv-optimizer/tests/it/python_relations.rs new file mode 100644 index 000000000..65e3bf119 --- /dev/null +++ b/crates/krishiv-optimizer/tests/it/python_relations.rs @@ -0,0 +1,198 @@ +// SPDX-License-Identifier: Apache-2.0 +//! The Python UDF relation nodes are opaque to the optimizer (design P3-T085 +//! §3.1, ADR 0022 §5): no filter moves through them and their inputs keep +//! every column (the UDF sees them all). + +use std::sync::Arc; + +use krishiv_optimizer::optimize; +use krishiv_plan::{ + AttrId, BoundColumn, Builtin, Expr, ExprNode, FunctionId, FunctionKind, Literal, LogicalPlan, + Plan, PlanNode, PythonEvalType, PythonUdfRef, ResolvedExpr, ResolvedFunction, +}; +use krishiv_types::DataType; + +use super::support::{ctx, resolve}; + +fn udf(eval_type: PythonEvalType) -> PythonUdfRef { + PythonUdfRef { + name: "f".to_owned(), + version: 1, + eval_type, + thread_safe: true, + deterministic: true, + } +} + +fn column(c: &BoundColumn) -> Arc { + Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Column( + c.clone(), + )))) +} + +/// `c > 1` over a `BIGINT` column. +fn above_one(c: &BoundColumn) -> Arc { + Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Function( + ResolvedFunction { + id: FunctionId::Builtin(Builtin::Gt), + kind: FunctionKind::Scalar, + args: vec![column(c), Arc::new(Expr::from(Literal::Long(1)))], + distinct: false, + filter: None, + within_group: Vec::new(), + data_type: DataType::Boolean, + nullable: true, + deterministic: true, + }, + )))) +} + +/// The scan below `p` (through any single-input nodes). +fn scan_width(p: &LogicalPlan) -> usize { + let mut found = 0; + p.for_each_child(&mut |c| found = found.max(scan_width(c))); + match &p.node { + PlanNode::Scan(s) => s.output.len(), + _ => found, + } +} + +/// The relation nodes over `input` (columns `cols`), each declaring one +/// `BIGINT` column `x` with id 1000. +fn relations(input: &Arc, cols: &[BoundColumn]) -> Vec { + let x = vec![BoundColumn { + id: AttrId(1000), + name: "x".to_owned(), + data_type: DataType::Long, + nullable: true, + }]; + vec![ + LogicalPlan::new(PlanNode::MapPartitions { + input: Arc::clone(input), + udf: udf(PythonEvalType::MapArrowIter), + output: x.clone(), + barrier: false, + }), + LogicalPlan::new(PlanNode::GroupMap { + input: Arc::clone(input), + keys: vec![column(&cols[0])], + udf: udf(PythonEvalType::GroupedMapPandas), + output: x.clone(), + }), + LogicalPlan::new(PlanNode::CoGroupMap { + left: Arc::clone(input), + right: Arc::clone(input), + left_keys: vec![column(&cols[0])], + right_keys: vec![column(&cols[0])], + udf: udf(PythonEvalType::CogroupedMapArrow), + output: x, + }), + ] +} + +#[test] +fn a_filter_stays_above_a_python_relation_whose_input_keeps_every_column() { + let base = resolve("SELECT * FROM nation", &ctx()); + let input = Arc::clone(base.root()); + let cols = input.resolved_output().unwrap(); + let width = cols.len(); + assert!(width > 2, "nation has several columns"); + let x = BoundColumn { + id: AttrId(1000), + name: "x".to_owned(), + data_type: DataType::Long, + nullable: true, + }; + for node in relations(&input, &cols) { + let kind = node + .to_string() + .split_whitespace() + .next() + .unwrap() + .to_owned(); + let plan = Plan::try_new(LogicalPlan::new(PlanNode::Filter { + input: Arc::new(node), + predicate: above_one(&x), + })) + .unwrap(); + let out = optimize(&plan).unwrap(); + let root = out.root(); + let PlanNode::Filter { input, .. } = &root.node else { + panic!("{kind}: the filter moved:\n{root}"); + }; + assert!( + input.to_string().starts_with(&kind), + "{kind}: not directly above it:\n{root}" + ); + assert_eq!( + scan_width(root), + width, + "{kind}: an input column was pruned:\n{root}" + ); + } +} + +/// The `CteRef`s left in `p`. +fn cte_refs(p: &LogicalPlan) -> usize { + let mut n = usize::from(matches!(p.node, PlanNode::CteRef(_))); + p.for_each_child(&mut |c| n += cte_refs(c)); + n +} + +/// `WITH c AS () SELECT * +/// FROM c UNION ALL SELECT * FROM c`, optimized. +fn twice_referenced_map(deterministic: bool) -> Arc { + let base = resolve("SELECT * FROM nation", &ctx()); + let x = |id: u64| BoundColumn { + id: AttrId(id), + name: "x".to_owned(), + data_type: DataType::Long, + nullable: true, + }; + let def = Arc::new(LogicalPlan::new(PlanNode::MapPartitions { + input: Arc::clone(base.root()), + udf: PythonUdfRef { + deterministic, + ..udf(PythonEvalType::MapArrowIter) + }, + output: vec![x(1000)], + barrier: false, + })); + let reference = |id: u64| { + Arc::new(LogicalPlan::new(PlanNode::CteRef(krishiv_plan::CteRef { + name: "c".to_owned(), + source: vec![AttrId(1000)], + output: vec![x(id)], + }))) + }; + let plan = Plan::try_new(LogicalPlan::new(PlanNode::WithCte { + input: Arc::new(LogicalPlan::new(PlanNode::Union { + inputs: vec![reference(2000), reference(3000)], + by_name: false, + allow_missing_columns: false, + })), + ctes: vec![krishiv_plan::Cte { + name: "c".to_owned(), + plan: Arc::clone(&def), + }], + })) + .unwrap(); + assert_eq!( + krishiv_optimizer::plan_is_deterministic(&def), + deterministic + ); + Arc::clone(optimize(&plan).unwrap().root()) +} + +/// A map over a nondeterministic UDF (Spark's `MapInPandas` allows one, +/// `CheckAnalysis` at v4.2.0) makes its CTE nondeterministic: referenced +/// twice, it is computed once and shared, not inlined per reference +/// (Spark's `InlineCTE.shouldInline`, `cteDef.deterministic`, which sees +/// the node's `PythonUDF`). A deterministic one is inlined. +#[test] +fn a_cte_over_a_nondeterministic_python_map_is_shared_not_inlined() { + let shared = twice_referenced_map(false); + assert_eq!(cte_refs(&shared), 2, "{shared}"); + let inlined = twice_referenced_map(true); + assert_eq!(cte_refs(&inlined), 0, "{inlined}"); +} diff --git a/crates/krishiv-plan-codec/src/decode.rs b/crates/krishiv-plan-codec/src/decode.rs index 494c61d5a..bef2b1c70 100644 --- a/crates/krishiv-plan-codec/src/decode.rs +++ b/crates/krishiv-plan-codec/src/decode.rs @@ -573,6 +573,37 @@ impl Decoder { }) .collect::>()?, }, + Kind::PythonMap(p) => PhysNode::PythonMap { + input: self.boxed(p.input)?, + udf: python_relation(required(&p.udf, "a Python relation UDF")?)?, + }, + Kind::PythonGroupMap(p) => PhysNode::PythonGroupMap { + input: self.boxed(p.input)?, + keys: usizes(&p.keys)?, + udf: python_relation(required(&p.udf, "a Python relation UDF")?)?, + }, + Kind::PythonCoGroupMap(p) => PhysNode::PythonCoGroupMap { + input: self.boxed(p.input)?, + keys: usizes(&p.keys)?, + side: usize_of(p.side)?, + udf: python_relation(required(&p.udf, "a Python relation UDF")?)?, + }, + Kind::PythonAggregate(p) => PhysNode::PythonAggregate { + input: self.boxed(p.input)?, + layout: layout(&p.layout), + keys: usizes(&p.keys)?, + key_values: usizes(&p.key_values)?, + calls: p + .calls + .iter() + .map(|c| { + Ok(PythonCallPlan { + call: self.expr_ref(c.call)?, + output: AttrId(c.output), + }) + }) + .collect::>()?, + }, Kind::Generate(g) => PhysNode::Generate { input: self.boxed(g.input)?, layout: layout(&g.layout), @@ -581,6 +612,7 @@ impl Decoder { PosExplode => krishiv_plan::GeneratorKind::PosExplode, Inline => krishiv_plan::GeneratorKind::Inline, VariantExplode => krishiv_plan::GeneratorKind::VariantExplode, + PythonUdtf => krishiv_plan::GeneratorKind::PythonUdtf, })?, outer: g.outer, arg: self.expr_ref(g.arg)?, @@ -1711,3 +1743,43 @@ fn day_time_field(v: i32) -> Result { Second => DayTimeField::Second, }) } + +/// A Python relation node's UDF, frame sides and declared result. +fn python_relation(p: &pb::PythonRelation) -> Result { + let u = required(&p.udf, "a Python relation UDF identity")?; + let udf = krishiv_plan::PythonUdfRef { + name: u.name.clone(), + version: u32::try_from(u.version) + .map_err(|_| corrupt("a Python UDF version out of range"))?, + eval_type: PythonEvalType::from_code(u.eval_type) + .ok_or_else(|| corrupt("an unknown Python UDF eval type"))?, + thread_safe: u.thread_safe, + deterministic: p.deterministic, + }; + let sides = p + .sides + .iter() + .map(|s| { + if s.names.len() != s.columns.len() || s.types.len() != s.columns.len() { + return Err(corrupt("a Python frame side of uneven lengths")); + } + Ok(krishiv_exec::physical::PythonFrameSide { + columns: usizes(&s.columns)?, + names: s.names.clone(), + types: s.types.iter().map(spark_type).collect::>()?, + }) + }) + .collect::>()?; + let output = p + .output + .iter() + .map(|c| { + Ok(krishiv_exec::physical::PythonOutputColumn { + name: c.name.clone(), + data_type: spark_type(required(&c.data_type, "a Python result type")?)?, + nullable: c.nullable, + }) + }) + .collect::>()?; + Ok(krishiv_exec::physical::PythonRelationPlan { udf, sides, output }) +} diff --git a/crates/krishiv-plan-codec/src/encode.rs b/crates/krishiv-plan-codec/src/encode.rs index 1088ad88c..27c2e7e86 100644 --- a/crates/krishiv-plan-codec/src/encode.rs +++ b/crates/krishiv-plan-codec/src/encode.rs @@ -345,6 +345,51 @@ impl Encoder { }) .collect::>()?, }), + PhysNode::PythonMap { input, udf } => Kind::PythonMap(pb::PythonMapNode { + input: self.node_index(input)?, + udf: Some(self.python_relation(udf)?), + }), + PhysNode::PythonGroupMap { input, keys, udf } => { + Kind::PythonGroupMap(pb::PythonGroupMapNode { + input: self.node_index(input)?, + keys: indexes(keys), + udf: Some(self.python_relation(udf)?), + }) + } + PhysNode::PythonCoGroupMap { + input, + keys, + side, + udf, + } => Kind::PythonCoGroupMap(pb::PythonCoGroupMapNode { + input: self.node_index(input)?, + keys: indexes(keys), + side: u64::try_from(*side).map_err(|_| { + CodecError::Unsupported("a side column out of range".to_owned()) + })?, + udf: Some(self.python_relation(udf)?), + }), + PhysNode::PythonAggregate { + input, + layout: l, + keys, + key_values, + calls, + } => Kind::PythonAggregate(pb::PythonAggregateNode { + input: self.node_index(input)?, + layout: layout(l), + keys: indexes(keys), + key_values: indexes(key_values), + calls: calls + .iter() + .map(|c| { + Ok(pb::PythonCall { + call: self.expr(&c.call)?, + output: c.output.0, + }) + }) + .collect::>()?, + }), PhysNode::Generate { input, layout: l, @@ -362,6 +407,7 @@ impl Encoder { krishiv_plan::GeneratorKind::VariantExplode => { pb::GeneratorKind::VariantExplode } + krishiv_plan::GeneratorKind::PythonUdtf => pb::GeneratorKind::PythonUdtf, } .into(), outer: *outer, @@ -754,6 +800,48 @@ impl Encoder { }) } + /// A Python relation node's UDF, frame sides and declared result. + fn python_relation( + &self, + r: &krishiv_exec::physical::PythonRelationPlan, + ) -> Result { + Ok(pb::PythonRelation { + udf: Some(pb::PythonUdfFunction { + name: r.udf.name.clone(), + version: u64::from(r.udf.version), + eval_type: r.udf.eval_type.code(), + thread_safe: r.udf.thread_safe, + }), + deterministic: r.udf.deterministic, + sides: r + .sides + .iter() + .map(|s| { + Ok(pb::PythonFrameSide { + columns: indexes(&s.columns), + names: s.names.clone(), + types: s + .types + .iter() + .map(|t| spark_type(t, &self.nested)) + .collect::>()?, + }) + }) + .collect::>()?, + output: r + .output + .iter() + .map(|c| { + Ok(pb::PythonOutputColumn { + name: c.name.clone(), + data_type: Some(spark_type(&c.data_type, &self.nested)?), + nullable: c.nullable, + }) + }) + .collect::>()?, + }) + } + fn function(&self, f: &ResolvedFunction) -> Result { let id = match &f.id { FunctionId::Builtin(b) => pb::function::Id::Builtin(b.name().to_owned()), diff --git a/crates/krishiv-plan-codec/src/features.rs b/crates/krishiv-plan-codec/src/features.rs index f811db498..1bef98a48 100644 --- a/crates/krishiv-plan-codec/src/features.rs +++ b/crates/krishiv-plan-codec/src/features.rs @@ -154,6 +154,7 @@ registry! { (partitioned_writes, PartitionedWrites, POST_MIGRATION_READER_VERSION, Semantic, "A `WritePlan` laid out by partition columns (Spark's directory layout)."), (percentile_aggregates, PercentileAggregates, POST_MIGRATION_READER_VERSION, Semantic, "The percentile family and `approx_count_distinct` aggregates (`percentile`, `percentile_approx`, `percentile_disc`, `approx_count_distinct` with their folded parameters)."), (python_udf, PythonUdf, POST_MIGRATION_READER_VERSION, Semantic, "Python UDF calls (`PythonUdfFunction`, by name and registration version) and the `PythonEval` node evaluating them; a worker without a UDF runtime refuses the plan."), + (python_udf_relations, PythonUdfRelations, POST_MIGRATION_READER_VERSION, Semantic, "The Python UDF relation kinds (design P3-T085): the `PythonMap`, `PythonGroupMap`, `PythonCoGroupMap` and `PythonAggregate` nodes, the Python UDTF generator and the eval types beyond the scalar ones (Spark's codes 201-217, 251-254, 300-302); a worker without a runtime for them refuses the plan."), (range, Range, POST_MIGRATION_READER_VERSION, Semantic, "The `Range` node (`range(start, end, step, numSlices)`)."), (rust_udaf, RustUdaf, POST_MIGRATION_READER_VERSION, Semantic, "A call of a session-registered native Rust aggregate UDF (`UdafCall`, by name and version, with its argument and state columns)."), (rust_udf, RustUdf, POST_MIGRATION_READER_VERSION, Semantic, "A call of a session-registered native Rust UDF (`UdfFunction`, by name and version)."), @@ -240,6 +241,8 @@ pub enum Construct<'a> { Generator(krishiv_plan::GeneratorKind), /// A hash join's probe selection (`HashJoinPlan.select`). ProbeSelect, + /// A Python UDF's eval type (of a call or a relation node's UDF). + PythonEvalType(krishiv_plan::PythonEvalType), } /// The feature a type needs on its own. The walk reports a type inside an @@ -414,6 +417,10 @@ pub fn feature_of(construct: Construct<'_>) -> Option { PhysNode::WindowGroupLimit { .. } => Some(Feature::WindowGroupLimit), PhysNode::PythonEval { .. } => Some(Feature::PythonUdf), PhysNode::Sample { .. } => Some(Feature::Sample), + PhysNode::PythonMap { .. } + | PhysNode::PythonGroupMap { .. } + | PhysNode::PythonCoGroupMap { .. } + | PhysNode::PythonAggregate { .. } => Some(Feature::PythonUdfRelations), // A partitioned write needs a reader that // lays files out by partition (the older fields would be // ignored and the rows written flat). @@ -481,8 +488,11 @@ pub fn feature_of(construct: Construct<'_>) -> Option { Construct::EvalMode => Some(Feature::CapturedEvalMode), Construct::WriteColumnTypes => Some(Feature::WriteColumnTypes), Construct::ProbeSelect => Some(Feature::FusedProbe), + // The scalar kinds came with `python_udf`; every other kind is newer. + Construct::PythonEvalType(t) => (!t.is_scalar()).then_some(Feature::PythonUdfRelations), Construct::Generator(kind) => match kind { krishiv_plan::GeneratorKind::VariantExplode => Some(Feature::VariantExplode), + krishiv_plan::GeneratorKind::PythonUdtf => Some(Feature::PythonUdfRelations), krishiv_plan::GeneratorKind::Explode | krishiv_plan::GeneratorKind::PosExplode | krishiv_plan::GeneratorKind::Inline => None, @@ -711,7 +721,13 @@ fn constructs_of<'a>(node: &'a PhysNode, f: &mut impl FnMut(Construct<'a>)) { | PhysNode::TopN { .. } | PhysNode::WindowGroupLimit { .. } | PhysNode::PythonEval { .. } + | PhysNode::PythonAggregate { .. } | PhysNode::Union { .. } => {} + PhysNode::PythonMap { udf, .. } + | PhysNode::PythonGroupMap { udf, .. } + | PhysNode::PythonCoGroupMap { udf, .. } => { + f(Construct::PythonEvalType(udf.udf.eval_type)); + } PhysNode::WriteFiles { write, .. } => { if write.as_ref().is_some_and(|w| !w.column_types.is_empty()) { f(Construct::WriteColumnTypes); @@ -807,6 +823,13 @@ fn types_of<'a>(node: &'a PhysNode) -> Vec<&'a DataType> { if let PhysNode::WriteFiles { write: Some(w), .. } = node { out.extend(w.column_types.values()); } + if let PhysNode::PythonMap { udf, .. } + | PhysNode::PythonGroupMap { udf, .. } + | PhysNode::PythonCoGroupMap { udf, .. } = node + { + out.extend(udf.sides.iter().flat_map(|s| &s.types)); + out.extend(udf.output.iter().map(|c| &c.data_type)); + } out } @@ -830,7 +853,9 @@ fn exprs_of<'a>(node: &'a PhysNode) -> Vec<&'a Arc> { } PhysNode::ReplicateRows { count, .. } => roots.push(count), PhysNode::Generate { arg, .. } => roots.push(arg), - PhysNode::PythonEval { calls, .. } => roots.extend(calls.iter().map(|c| &c.call)), + PhysNode::PythonEval { calls, .. } | PhysNode::PythonAggregate { calls, .. } => { + roots.extend(calls.iter().map(|c| &c.call)); + } PhysNode::HashJoin(j) => { roots.extend(j.residual.iter().flat_map(|r| &r.conjuncts)); roots.extend( @@ -858,7 +883,10 @@ fn exprs_of<'a>(node: &'a PhysNode) -> Vec<&'a Arc> { | PhysNode::Union { .. } | PhysNode::WriteFiles { .. } | PhysNode::Exchange(_) - | PhysNode::ExchangeSource { .. } => {} + | PhysNode::ExchangeSource { .. } + | PhysNode::PythonMap { .. } + | PhysNode::PythonGroupMap { .. } + | PhysNode::PythonCoGroupMap { .. } => {} } // Expressions share subtrees: each is visited once. let mut seen = std::collections::HashSet::new(); @@ -885,6 +913,9 @@ fn exprs_of<'a>(node: &'a PhysNode) -> Vec<&'a Arc> { fn functions_in<'a>(e: &'a Arc, f: &mut impl FnMut(Construct<'a>)) { if let ExprNode::Resolved(ResolvedExpr::Function(c)) = &e.node { f(Construct::Function(&c.id)); + if let FunctionId::PythonUdf { eval_type, .. } = &c.id { + f(Construct::PythonEvalType(*eval_type)); + } } } diff --git a/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs b/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs index df1dcdee4..96c700b2d 100644 --- a/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs +++ b/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs @@ -145,7 +145,7 @@ pub struct ExplainNode { pub struct Node { #[prost( oneof = "node::Kind", - tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 22, 23, 21, 24, 25, 26, 27" + tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 22, 23, 21, 24, 25, 26, 27, 28, 29, 30, 31" )] pub kind: ::core::option::Option, } @@ -220,6 +220,18 @@ pub mod node { /// feature: sample #[prost(message, tag = "27")] Sample(super::SampleNode), + /// feature: python_udf_relations + #[prost(message, tag = "28")] + PythonMap(super::PythonMapNode), + /// feature: python_udf_relations + #[prost(message, tag = "29")] + PythonGroupMap(super::PythonGroupMapNode), + /// feature: python_udf_relations + #[prost(message, tag = "30")] + PythonCoGroupMap(super::PythonCoGroupMapNode), + /// feature: python_udf_relations + #[prost(message, tag = "31")] + PythonAggregate(super::PythonAggregateNode), } } /// feature: sample @@ -238,6 +250,93 @@ pub struct SampleNode { #[prost(int64, tag = "4")] pub seed: i64, } +/// feature: python_udf_relations +/// `mapInPandas` / `mapInArrow`: each partition of `input` sent to a Python +/// UDF task. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct PythonMapNode { + #[prost(uint32, tag = "1")] + pub input: u32, + #[prost(message, optional, tag = "2")] + pub udf: ::core::option::Option, +} +/// feature: python_udf_relations +/// A grouped map over `input` clustered and sorted by the key columns. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct PythonGroupMapNode { + #[prost(uint32, tag = "1")] + pub input: u32, + #[prost(uint64, repeated, tag = "2")] + pub keys: ::prost::alloc::vec::Vec, + #[prost(message, optional, tag = "3")] + pub udf: ::core::option::Option, +} +/// feature: python_udf_relations +/// A cogrouped map over both relations in one stream, sorted by the key +/// columns then the side column. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct PythonCoGroupMapNode { + #[prost(uint32, tag = "1")] + pub input: u32, + #[prost(uint64, repeated, tag = "2")] + pub keys: ::prost::alloc::vec::Vec, + #[prost(uint64, tag = "3")] + pub side: u64, + #[prost(message, optional, tag = "4")] + pub udf: ::core::option::Option, +} +/// feature: python_udf_relations +/// An aggregate of grouped-aggregate Python UDF calls over `input` clustered +/// and sorted by the key columns (`keys`: a collated key's comparison key); +/// each group shows its first row's `key_values`. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct PythonAggregateNode { + #[prost(uint32, tag = "1")] + pub input: u32, + #[prost(uint64, repeated, tag = "2")] + pub layout: ::prost::alloc::vec::Vec, + #[prost(uint64, repeated, tag = "3")] + pub keys: ::prost::alloc::vec::Vec, + #[prost(message, repeated, tag = "4")] + pub calls: ::prost::alloc::vec::Vec, + #[prost(uint64, repeated, tag = "5")] + pub key_values: ::prost::alloc::vec::Vec, +} +/// feature: python_udf_relations +/// The Python UDF of a relation node: by name and version, the columns of its +/// frames per side, and its declared result. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct PythonRelation { + #[prost(message, optional, tag = "1")] + pub udf: ::core::option::Option, + #[prost(message, repeated, tag = "2")] + pub sides: ::prost::alloc::vec::Vec, + #[prost(message, repeated, tag = "3")] + pub output: ::prost::alloc::vec::Vec, + /// Whether the UDF is deterministic (Spark's `udfDeterministic`). + #[prost(bool, tag = "4")] + pub deterministic: bool, +} +/// feature: python_udf_relations +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct PythonFrameSide { + #[prost(uint64, repeated, tag = "1")] + pub columns: ::prost::alloc::vec::Vec, + #[prost(string, repeated, tag = "2")] + pub names: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, + #[prost(message, repeated, tag = "3")] + pub types: ::prost::alloc::vec::Vec, +} +/// feature: python_udf_relations +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct PythonOutputColumn { + #[prost(string, tag = "1")] + pub name: ::prost::alloc::string::String, + #[prost(message, optional, tag = "2")] + pub data_type: ::core::option::Option, + #[prost(bool, tag = "3")] + pub nullable: bool, +} /// feature: python_udf /// The rows of `input`, then one column per Python UDF call over every row. #[derive(Clone, PartialEq, ::prost::Message)] @@ -1676,6 +1775,8 @@ pub enum GeneratorKind { Inline = 3, /// feature: variant_explode VariantExplode = 4, + /// feature: python_udf_relations + PythonUdtf = 5, } impl GeneratorKind { /// String value of the enum field names used in the ProtoBuf definition. @@ -1689,6 +1790,7 @@ impl GeneratorKind { Self::PosExplode => "GENERATOR_KIND_POS_EXPLODE", Self::Inline => "GENERATOR_KIND_INLINE", Self::VariantExplode => "GENERATOR_KIND_VARIANT_EXPLODE", + Self::PythonUdtf => "GENERATOR_KIND_PYTHON_UDTF", } } /// Creates an enum from field names used in the ProtoBuf definition. @@ -1699,6 +1801,7 @@ impl GeneratorKind { "GENERATOR_KIND_POS_EXPLODE" => Some(Self::PosExplode), "GENERATOR_KIND_INLINE" => Some(Self::Inline), "GENERATOR_KIND_VARIANT_EXPLODE" => Some(Self::VariantExplode), + "GENERATOR_KIND_PYTHON_UDTF" => Some(Self::PythonUdtf), _ => None, } } diff --git a/crates/krishiv-plan-codec/tests/common/inventory.rs b/crates/krishiv-plan-codec/tests/common/inventory.rs index bb4a262b8..655dfa001 100644 --- a/crates/krishiv-plan-codec/tests/common/inventory.rs +++ b/crates/krishiv-plan-codec/tests/common/inventory.rs @@ -1743,6 +1743,169 @@ pub(crate) fn python_eval() -> PhysNode { } } +/// A Python UDF reference of a relation node (`python_udf_relations`). +fn python_udf_ref(eval_type: krishiv_plan::PythonEvalType) -> krishiv_plan::PythonUdfRef { + krishiv_plan::PythonUdfRef { + name: "py_rel".to_owned(), + version: 4, + eval_type, + thread_safe: true, + deterministic: true, + } +} + +/// A relation node's UDF: frames of columns 0 and 1 per side, a declared +/// `n BIGINT, tags ARRAY` result. +fn python_relation( + eval_type: krishiv_plan::PythonEvalType, + sides: usize, +) -> krishiv_exec::physical::PythonRelationPlan { + use krishiv_exec::physical::{PythonFrameSide, PythonOutputColumn, PythonRelationPlan}; + PythonRelationPlan { + udf: python_udf_ref(eval_type), + sides: (0..sides) + .map(|s| PythonFrameSide { + columns: vec![2 * s, 2 * s + 1], + names: vec!["a".to_owned(), "b".to_owned()], + types: vec![DataType::Integer, DataType::Long], + }) + .collect(), + output: vec![ + PythonOutputColumn { + name: "n".to_owned(), + data_type: DataType::Long, + nullable: false, + }, + PythonOutputColumn { + name: "tags".to_owned(), + data_type: DataType::Array( + krishiv_types::ArrayType::new(DataType::String, true).into(), + ), + nullable: true, + }, + ], + } +} + +/// The Python UDF relation nodes (`python_udf_relations`, design P3-T085 +/// §3.1), each over a scan. +pub(crate) fn python_relations() -> Vec<(&'static str, PhysNode)> { + use krishiv_plan::PythonEvalType as T; + let call = |eval_type: T| { + Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Function( + ResolvedFunction { + within_group: Vec::new(), + id: python_udf_ref(eval_type).function_id(), + kind: FunctionKind::Aggregate, + args: vec![col(2, "b", DataType::Long)], + distinct: false, + filter: None, + data_type: DataType::Double, + nullable: true, + deterministic: true, + }, + )))) + }; + vec![ + ( + "python_map", + PhysNode::PythonMap { + input: boxed(scan()), + udf: python_relation(T::MapArrowIter, 1), + }, + ), + ( + "python_group_map", + PhysNode::PythonGroupMap { + input: boxed(scan()), + keys: vec![0], + udf: python_relation(T::GroupedMapPandasIter, 1), + }, + ), + ( + "python_cogroup_map", + PhysNode::PythonCoGroupMap { + input: boxed(scan()), + keys: vec![0], + side: 1, + udf: python_relation(T::CogroupedMapPandas, 2), + }, + ), + ( + "python_aggregate", + PhysNode::PythonAggregate { + input: boxed(scan()), + layout: vec![AttrId(1), AttrId(2)], + keys: vec![0], + key_values: vec![1], + calls: vec![ + krishiv_exec::physical::PythonCallPlan { + call: call(T::GroupedAggArrow), + output: AttrId(9), + }, + krishiv_exec::physical::PythonCallPlan { + call: call(T::GroupedAggPandasIter), + output: AttrId(10), + }, + ], + }, + ), + ( + "python_udtf_generate", + PhysNode::Generate { + input: boxed(scan()), + layout: vec![AttrId(1), AttrId(2)], + kind: krishiv_plan::GeneratorKind::PythonUdtf, + outer: false, + arg: Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Function( + ResolvedFunction { + within_group: Vec::new(), + id: python_udf_ref(T::ArrowUdtf).function_id(), + kind: FunctionKind::Generator, + args: vec![col(1, "a", DataType::Integer)], + distinct: false, + filter: None, + data_type: DataType::Long, + nullable: true, + deterministic: true, + }, + )))), + subqueries: Vec::new(), + }, + ), + ] +} + +/// Every eval type, each in a `PythonEval` call: the scalar ones need only +/// `python_udf`, every other one `python_udf_relations` too. +pub(crate) fn python_eval_types() -> PhysNode { + let calls = krishiv_plan::PythonEvalType::ALL + .iter() + .enumerate() + .map(|(i, t)| krishiv_exec::physical::PythonCallPlan { + call: Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Function( + ResolvedFunction { + within_group: Vec::new(), + id: python_udf_ref(*t).function_id(), + kind: FunctionKind::Scalar, + args: vec![col(1, "a", DataType::Integer)], + distinct: false, + filter: None, + data_type: DataType::Long, + nullable: true, + deterministic: true, + }, + )))), + output: AttrId(100 + i as u64), + }) + .collect(); + PhysNode::PythonEval { + input: boxed(scan()), + layout: vec![AttrId(1), AttrId(2)], + calls, + } +} + /// A call of a session-registered native Rust aggregate UDF /// (`rust_udaf`): by name and version, with the argument and state /// columns the planner fixed, in a final aggregate and the merge of its @@ -1925,6 +2088,12 @@ pub(crate) fn cases() -> Vec<(String, PhysicalPlan)> { out.push(("expressions".to_owned(), plan_of(expressions()))); out.push(("rust_udf".to_owned(), plan_of(udf_call()))); out.push(("python_udf".to_owned(), plan_of(python_eval()))); + out.extend( + python_relations() + .into_iter() + .map(|(name, node)| (name.to_owned(), plan_of(node))), + ); + out.push(("python_eval_types".to_owned(), plan_of(python_eval_types()))); out.push(("rust_udaf".to_owned(), plan_of(udaf_call()))); // The expression environment flag the `legacy_statistical_aggregate` // feature is about: the same root as `one_row`, under a session with diff --git a/crates/krishiv-plan-codec/tests/features.rs b/crates/krishiv-plan-codec/tests/features.rs index 42745b98b..0f3c8abaf 100644 --- a/crates/krishiv-plan-codec/tests/features.rs +++ b/crates/krishiv-plan-codec/tests/features.rs @@ -2101,6 +2101,58 @@ fn a_python_udf_needs_its_feature() { ); } +/// The Python UDF relation nodes, the Python UDTF generator and every eval +/// type past the scalar ones list `python_udf_relations` (semantic: no plan +/// without them computes the same rows); a reader without it refuses each +/// plan naming it, and a scalar-only plan does not list it. +#[test] +fn the_python_udf_relation_kinds_need_their_feature() { + assert_eq!(Feature::PythonUdfRelations.kind(), FeatureKind::Semantic); + // The declared `ARRAY` result is a nested type. + for name in ["python_map", "python_group_map", "python_cogroup_map"] { + assert_eq!( + message(name).features, + ["nested_types", "python_udf_relations"], + "{name}" + ); + } + assert_eq!( + message("python_udtf_generate").features, + ["generate", "python_udf", "python_udf_relations"] + ); + for name in ["python_aggregate", "python_eval_types"] { + assert_eq!( + message(name).features, + ["python_udf", "python_udf_relations"], + "{name}" + ); + } + assert_eq!(message("python_udf").features, ["python_udf"]); + let older = ReaderFeatures::only( + Feature::ALL + .iter() + .copied() + .filter(|f| *f != Feature::PythonUdfRelations), + ); + for name in [ + "python_map", + "python_group_map", + "python_cogroup_map", + "python_aggregate", + "python_udtf_generate", + "python_eval_types", + ] { + let bytes = message(name).encode_to_vec(); + assert_eq!( + decode_with(&bytes, &Limits::default(), &older).unwrap_err(), + CodecError::UnsupportedFeatures { + missing: vec!["python_udf_relations".to_owned()] + }, + "{name}" + ); + } +} + /// A `TABLESAMPLE` plan lists `sample` (a semantic feature: no plan without /// the node keeps the same rows), and a reader without it refuses the plan /// naming the feature. diff --git a/crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_aggregate.bin b/crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_aggregate.bin new file mode 100644 index 000000000..6c793805a Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_aggregate.bin differ diff --git a/crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_cogroup_map.bin b/crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_cogroup_map.bin new file mode 100644 index 000000000..6054d2571 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_cogroup_map.bin differ diff --git a/crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_eval_types.bin b/crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_eval_types.bin new file mode 100644 index 000000000..23b0e1bca Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_eval_types.bin differ diff --git a/crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_group_map.bin b/crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_group_map.bin new file mode 100644 index 000000000..c2b5b2e5d Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_group_map.bin differ diff --git a/crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_map.bin b/crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_map.bin new file mode 100644 index 000000000..e18f471c3 Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_map.bin differ diff --git a/crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_udtf_generate.bin b/crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_udtf_generate.bin new file mode 100644 index 000000000..c236a5e5d Binary files /dev/null and b/crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_udtf_generate.bin differ diff --git a/crates/krishiv-plan-codec/tests/roundtrip.rs b/crates/krishiv-plan-codec/tests/roundtrip.rs index 8fe380ce2..97eb20775 100644 --- a/crates/krishiv-plan-codec/tests/roundtrip.rs +++ b/crates/krishiv-plan-codec/tests/roundtrip.rs @@ -45,6 +45,10 @@ variants!( ReplicateRows, Generate, PythonEval, + PythonMap, + PythonGroupMap, + PythonCoGroupMap, + PythonAggregate, Expand, Sort, TopN, @@ -76,6 +80,10 @@ fn variant_of(node: &PhysNode) -> Variant { PhysNode::ReplicateRows { .. } => Variant::ReplicateRows, PhysNode::Generate { .. } => Variant::Generate, PhysNode::PythonEval { .. } => Variant::PythonEval, + PhysNode::PythonMap { .. } => Variant::PythonMap, + PhysNode::PythonGroupMap { .. } => Variant::PythonGroupMap, + PhysNode::PythonCoGroupMap { .. } => Variant::PythonCoGroupMap, + PhysNode::PythonAggregate { .. } => Variant::PythonAggregate, PhysNode::UnionLoop { .. } => Variant::UnionLoop, PhysNode::LoopRef { .. } => Variant::LoopRef, PhysNode::Expand { .. } => Variant::Expand, @@ -612,6 +620,12 @@ fn readers_that_predate_features_still_refuse_what_they_cannot_read() { | "rust_udf" | "python_udf" | "sample" + | "python_map" + | "python_group_map" + | "python_cogroup_map" + | "python_aggregate" + | "python_udtf_generate" + | "python_eval_types" | "rust_udaf" | "nested_types" | "struct_field" diff --git a/crates/krishiv-plan/src/analyzer/resolve.rs b/crates/krishiv-plan/src/analyzer/resolve.rs index ccff9b9e7..3f3710ec5 100644 --- a/crates/krishiv-plan/src/analyzer/resolve.rs +++ b/crates/krishiv-plan/src/analyzer/resolve.rs @@ -954,6 +954,10 @@ impl Resolver<'_> { | PlanNode::Sample { .. } | PlanNode::Expand { .. } | PlanNode::Range { .. } + | PlanNode::MapPartitions { .. } + | PlanNode::GroupMap { .. } + | PlanNode::CoGroupMap { .. } + | PlanNode::PythonTableFunction { .. } | PlanNode::JoinHint { .. } => self.leaf_node(p, env), } } @@ -1149,6 +1153,15 @@ impl Resolver<'_> { | PlanNode::UnionLoop { .. } | PlanNode::UnionLoopRef { .. } | PlanNode::Window { .. } => Err(unsupported("a plan shape", o)), + // Resolved-only nodes: the front ends that build them (Connect, + // design P3-T014 PR-26d) resolve their inputs first. + PlanNode::MapPartitions { .. } + | PlanNode::GroupMap { .. } + | PlanNode::CoGroupMap { .. } + | PlanNode::PythonTableFunction { .. } => Err(unsupported( + "a Python UDF relation in an unresolved plan", + o, + )), } } diff --git a/crates/krishiv-plan/src/analyzer/subquery/support.rs b/crates/krishiv-plan/src/analyzer/subquery/support.rs index 6bd879615..49eed811d 100644 --- a/crates/krishiv-plan/src/analyzer/subquery/support.rs +++ b/crates/krishiv-plan/src/analyzer/subquery/support.rs @@ -109,9 +109,13 @@ pub(super) fn one_row_select(p: &Arc) -> Option>> { } } -/// Whether no expression of the plan `p` is nondeterministic. +/// Whether no expression of the plan `p`, and no Python relation UDF it +/// runs, is nondeterministic. pub(in crate::analyzer) fn plan_is_deterministic(p: &Arc) -> bool { visit_distinct(Node::Plan(p), |n| match n { + Node::Plan(q) if !crate::python_relation_is_deterministic(&q.node) => { + ControlFlow::Break(()) + } Node::Expr(e) if matches!( &e.node, diff --git a/crates/krishiv-plan/src/cache.rs b/crates/krishiv-plan/src/cache.rs index 404a4a1e5..ce85511ea 100644 --- a/crates/krishiv-plan/src/cache.rs +++ b/crates/krishiv-plan/src/cache.rs @@ -612,6 +612,15 @@ impl Encoder { } } + /// A relation node's Python UDF (by identity) and declared columns. + fn python_udf(&mut self, udf: &crate::PythonUdfRef, output: &[BoundColumn], bias: u64) { + self.debug(udf); + self.num(output.len() as u64); + for c in output { + self.bound(c, bias); + } + } + /// Writes `plan` as the digest of its encoding, computed once per node. fn plan(&mut self, plan: &LogicalPlan, bias: u64) { if self.semantic @@ -899,6 +908,63 @@ impl Encoder { } self.plan(input, bias); } + PlanNode::MapPartitions { + input, + udf, + output, + barrier, + } => { + self.tag("MapPartitions"); + self.python_udf(udf, output, bias); + self.flag(*barrier); + self.plan(input, bias); + } + PlanNode::GroupMap { + input, + keys, + udf, + output, + } => { + self.tag("GroupMap"); + self.exprs(keys, bias); + self.python_udf(udf, output, bias); + self.plan(input, bias); + } + PlanNode::CoGroupMap { + left, + right, + left_keys, + right_keys, + udf, + output, + } => { + self.tag("CoGroupMap"); + self.exprs(left_keys, bias); + self.exprs(right_keys, bias); + self.python_udf(udf, output, bias); + self.plan(left, bias); + self.plan(right, bias); + } + PlanNode::PythonTableFunction { + args, + table_arg, + udf, + output, + } => { + self.tag("PythonTableFunction"); + self.exprs(args, bias); + self.python_udf(udf, output, bias); + self.flag(table_arg.is_some()); + if let Some(t) = table_arg { + self.exprs(&t.partition_by, bias); + self.num(t.order_by.len() as u64); + for o in &t.order_by { + self.order(o, bias); + } + self.flag(t.single_partition); + self.plan(&t.input, bias); + } + } PlanNode::ReplicateRows { input, count } => { self.tag("ReplicateRows"); self.expr(count, bias); @@ -1753,6 +1819,58 @@ mod tests { ]), })), ), + ( + "map-partitions", + node(PlanNode::MapPartitions { + input: base(), + udf: py(crate::PythonEvalType::MapArrowIter), + output: vec![bound(2, "x")], + barrier: true, + }), + ), + ( + "group-map", + node(PlanNode::GroupMap { + input: base(), + keys: vec![a()], + udf: py(crate::PythonEvalType::GroupedMapArrow), + output: vec![bound(2, "x")], + }), + ), + ( + "cogroup-map", + node(PlanNode::CoGroupMap { + left: base(), + right: base(), + left_keys: vec![a()], + right_keys: vec![a()], + udf: py(crate::PythonEvalType::CogroupedMapArrow), + output: vec![bound(2, "x")], + }), + ), + ( + "python-table-function", + node(PlanNode::PythonTableFunction { + args: vec![lit(1)], + table_arg: Some(crate::TableArg { + input: base(), + partition_by: vec![a()], + order_by: vec![crate::SortOrder::asc(a())], + single_partition: false, + }), + udf: py(crate::PythonEvalType::ArrowUdtf), + output: vec![bound(2, "x")], + }), + ), + ( + "python-table-function-no-table", + node(PlanNode::PythonTableFunction { + args: vec![lit(1)], + table_arg: None, + udf: py(crate::PythonEvalType::TableUdf), + output: vec![bound(2, "x")], + }), + ), ( "read-table-options", node(PlanNode::Read(ReadSource::Table { @@ -1778,6 +1896,17 @@ mod tests { } } + /// A Python UDF reference of a relation node. + fn py(eval_type: crate::PythonEvalType) -> crate::PythonUdfRef { + crate::PythonUdfRef { + name: "f".into(), + version: 1, + eval_type, + thread_safe: true, + deterministic: true, + } + } + #[test] fn subtree_hash_is_stable_and_differs_from_whole() { let sub = filter(scan(&[3], &["a"]), gt(rcol(3, "a"), 1)); diff --git a/crates/krishiv-plan/src/display.rs b/crates/krishiv-plan/src/display.rs index ed7d4c639..d7ed0df61 100644 --- a/crates/krishiv-plan/src/display.rs +++ b/crates/krishiv-plan/src/display.rs @@ -223,6 +223,25 @@ fn subqueries(plan: &LogicalPlan) -> Vec<&LogicalPlan> { out } +/// A relation node's Python UDF: `name#version(eval type)`. +struct UdfRef<'a>(&'a crate::PythonUdfRef); + +impl Display for UdfRef<'_> { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + let name = crate::function::shown_call_name("", &self.0.name); + write!(f, "{name}#{}({})", self.0.version, self.0.eval_type.code()) + } +} + +/// `[name#id, ...]`. +fn columns(f: &mut Formatter<'_>, output: &[crate::BoundColumn]) -> fmt::Result { + f.write_str("[")?; + for (i, c) in output.iter().enumerate() { + write!(f, "{}{}#{}", sep(i), c.name, c.id.0)?; + } + f.write_str("]") +} + /// One node's line, without its children. struct NodeLine<'a>(&'a LogicalPlan); @@ -520,6 +539,60 @@ impl Display for NodeLine<'_> { PlanNode::TableFunction { name, args } => { write!(f, "TableFunction {name} [{}]", List(args)) } + // Spark's `MapInPandas f(a#1, b#2)#3, [x#4], false`, shortened. + PlanNode::MapPartitions { + udf, + output, + barrier, + .. + } => { + write!(f, "MapPartitions {}, ", UdfRef(udf))?; + columns(f, output)?; + write!(f, ", barrier={barrier}") + } + PlanNode::GroupMap { + keys, udf, output, .. + } => { + write!(f, "GroupMap [{}], {}, ", List(keys), UdfRef(udf))?; + columns(f, output) + } + PlanNode::CoGroupMap { + left_keys, + right_keys, + udf, + output, + .. + } => { + write!( + f, + "CoGroupMap [{}], [{}], {}, ", + List(left_keys), + List(right_keys), + UdfRef(udf) + )?; + columns(f, output) + } + PlanNode::PythonTableFunction { + args, + table_arg, + udf, + output, + } => { + write!(f, "PythonTableFunction {}({}), ", UdfRef(udf), List(args))?; + columns(f, output)?; + if let Some(t) = table_arg { + write!( + f, + ", TABLE PARTITION BY [{}] ORDER BY [", + List(&t.partition_by) + )?; + for (i, o) in t.order_by.iter().enumerate() { + write!(f, "{}{o}", sep(i))?; + } + write!(f, "], single={}", t.single_partition)?; + } + Ok(()) + } PlanNode::Range { start, end, diff --git a/crates/krishiv-plan/src/encode.rs b/crates/krishiv-plan/src/encode.rs index 6d28d5b88..b07dc3e6b 100644 --- a/crates/krishiv-plan/src/encode.rs +++ b/crates/krishiv-plan/src/encode.rs @@ -114,6 +114,10 @@ pub(crate) fn plan_tag(node: &PlanNode) -> u8 { PlanNode::PipeOperator { .. } => 34, PlanNode::Repartition { .. } => 35, PlanNode::Sample { .. } => 36, + PlanNode::MapPartitions { .. } => 37, + PlanNode::GroupMap { .. } => 38, + PlanNode::CoGroupMap { .. } => 39, + PlanNode::PythonTableFunction { .. } => 40, } } @@ -649,6 +653,13 @@ impl<'k, 'a> Encoder<'k, 'a> { } } + /// A relation node's Python UDF and its declared result columns. + fn python_udf(&mut self, udf: &crate::PythonUdfRef, output: &[BoundColumn]) { + self.function_id(&udf.function_id()); + self.len(output.len()); + output.iter().for_each(|c| self.bound(c)); + } + fn bytes(&mut self, b: &[u8]) { self.len(b.len()); self.put(b); @@ -1037,6 +1048,7 @@ impl<'k, 'a> Encoder<'k, 'a> { crate::GeneratorKind::PosExplode => 1, crate::GeneratorKind::Inline => 2, crate::GeneratorKind::VariantExplode => 3, + crate::GeneratorKind::PythonUdtf => 4, }); self.u8(u8::from(*outer)); self.expr(arg); @@ -1055,6 +1067,50 @@ impl<'k, 'a> Encoder<'k, 'a> { self.str(name); self.exprs(args); } + PlanNode::MapPartitions { + udf, + output, + barrier, + .. + } => { + self.python_udf(udf, output); + self.u8(u8::from(*barrier)); + } + PlanNode::GroupMap { + keys, udf, output, .. + } => { + self.exprs(keys); + self.python_udf(udf, output); + } + PlanNode::CoGroupMap { + left_keys, + right_keys, + udf, + output, + .. + } => { + self.exprs(left_keys); + self.exprs(right_keys); + self.python_udf(udf, output); + } + PlanNode::PythonTableFunction { + args, + table_arg, + udf, + output, + } => { + self.exprs(args); + match table_arg { + None => self.u8(0), + Some(t) => { + self.u8(1); + self.exprs(&t.partition_by); + self.orders(&t.order_by); + self.u8(u8::from(t.single_partition)); + } + } + self.python_udf(udf, output); + } PlanNode::Range { start, end, diff --git a/crates/krishiv-plan/src/function.rs b/crates/krishiv-plan/src/function.rs index 28b6f0074..4ee461532 100644 --- a/crates/krishiv-plan/src/function.rs +++ b/crates/krishiv-plan/src/function.rs @@ -355,9 +355,12 @@ pub enum FunctionId { }, } -/// The Python UDF kinds the engine evaluates, with Spark's -/// `PythonEvalType` codes (`pyspark/util.py`, Spark 4.2.0). Grouped, -/// iterator, table and aggregate kinds are later tasks. +/// The Python UDF kinds a plan can name, with Spark's `PythonEvalType` +/// codes (`pyspark/util.py` L655-688, Spark 4.2.0; the same at 4.3.0). +/// The scalar kinds (101, 200, 250) run in a `PythonEval`; the relation +/// kinds (design P3-T085 §2) see a relation, not a column, and run in the +/// relation nodes ([`PythonUdfKind`]). Streaming-state kinds (208, +/// 211-214) are not here (P4). #[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] #[non_exhaustive] pub enum PythonEvalType { @@ -367,14 +370,106 @@ pub enum PythonEvalType { /// A scalar `pandas_udf` (`SQL_SCALAR_PANDAS_UDF`, 200): pandas Series in /// and out. ScalarPandas, + /// `groupBy().applyInPandas` (`SQL_GROUPED_MAP_PANDAS_UDF`, 201). + GroupedMapPandas, + /// A grouped-aggregate `pandas_udf` (`SQL_GROUPED_AGG_PANDAS_UDF`, 202). + GroupedAggPandas, + /// A grouped-aggregate `pandas_udf` over a window + /// (`SQL_WINDOW_AGG_PANDAS_UDF`, 203). + WindowAggPandas, + /// A scalar iterator `pandas_udf` (`SQL_SCALAR_PANDAS_ITER_UDF`, 204). + ScalarPandasIter, + /// `mapInPandas` (`SQL_MAP_PANDAS_ITER_UDF`, 205). + MapPandasIter, + /// `cogroup().applyInPandas` (`SQL_COGROUPED_MAP_PANDAS_UDF`, 206). + CogroupedMapPandas, + /// `mapInArrow` (`SQL_MAP_ARROW_ITER_UDF`, 207). + MapArrowIter, + /// `groupBy().applyInArrow` (`SQL_GROUPED_MAP_ARROW_UDF`, 209). + GroupedMapArrow, + /// `cogroup().applyInArrow` (`SQL_COGROUPED_MAP_ARROW_UDF`, 210). + CogroupedMapArrow, + /// `applyInArrow` with iterator hints (`SQL_GROUPED_MAP_ARROW_ITER_UDF`, + /// 215). + GroupedMapArrowIter, + /// `applyInPandas` with iterator hints + /// (`SQL_GROUPED_MAP_PANDAS_ITER_UDF`, 216). + GroupedMapPandasIter, + /// A grouped-aggregate `pandas_udf` with iterator hints + /// (`SQL_GROUPED_AGG_PANDAS_ITER_UDF`, 217). + GroupedAggPandasIter, /// A scalar `arrow_udf` (`SQL_SCALAR_ARROW_UDF`, 250): `pyarrow.Array` /// in and out, no conversion. ScalarArrow, + /// A scalar iterator `arrow_udf` (`SQL_SCALAR_ARROW_ITER_UDF`, 251). + ScalarArrowIter, + /// A grouped-aggregate `arrow_udf` (`SQL_GROUPED_AGG_ARROW_UDF`, 252). + GroupedAggArrow, + /// A grouped-aggregate `arrow_udf` over a window + /// (`SQL_WINDOW_AGG_ARROW_UDF`, 253). + WindowAggArrow, + /// A grouped-aggregate `arrow_udf` with iterator hints + /// (`SQL_GROUPED_AGG_ARROW_ITER_UDF`, 254). + GroupedAggArrowIter, + /// `@udtf` (`SQL_TABLE_UDF`, 300). + TableUdf, + /// `@udtf(useArrow=True)` (`SQL_ARROW_TABLE_UDF`, 301). + ArrowTableUdf, + /// `@arrow_udtf` (`SQL_ARROW_UDTF`, 302). + ArrowUdtf, +} + +/// What a Python UDF sees and returns, by [`PythonEvalType`] (design P3-T085 +/// §2's table): which plan node runs it. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +#[non_exhaustive] +pub enum PythonUdfKind { + /// A column per argument in, a column out, row for row (`PythonEval`). + Scalar, + /// As [`Self::Scalar`], called once per partition with an iterator of + /// batches. + ScalarIter, + /// A partition's batches in, any batches out (`MapPartitions`). + Map, + /// A group's rows in, any rows out (`GroupMap`). + GroupedMap, + /// A key's rows of two relations in, any rows out (`CoGroupMap`). + CogroupedMap, + /// A group's argument columns in, one value out (an `Aggregate`). + GroupedAgg, + /// A window partition's argument columns in, one value per row out (a + /// `Window`). + WindowAgg, + /// A table function (`PythonTableFunction`, `Generate`). + Udtf, } impl PythonEvalType { /// Every kind, in code order. - pub const ALL: [Self; 3] = [Self::ArrowBatched, Self::ScalarPandas, Self::ScalarArrow]; + pub const ALL: [Self; 22] = [ + Self::ArrowBatched, + Self::ScalarPandas, + Self::GroupedMapPandas, + Self::GroupedAggPandas, + Self::WindowAggPandas, + Self::ScalarPandasIter, + Self::MapPandasIter, + Self::CogroupedMapPandas, + Self::MapArrowIter, + Self::GroupedMapArrow, + Self::CogroupedMapArrow, + Self::GroupedMapArrowIter, + Self::GroupedMapPandasIter, + Self::GroupedAggPandasIter, + Self::ScalarArrow, + Self::ScalarArrowIter, + Self::GroupedAggArrow, + Self::WindowAggArrow, + Self::GroupedAggArrowIter, + Self::TableUdf, + Self::ArrowTableUdf, + Self::ArrowUdtf, + ]; /// Spark's `PythonEvalType` code. #[must_use] @@ -382,15 +477,119 @@ impl PythonEvalType { match self { Self::ArrowBatched => 101, Self::ScalarPandas => 200, + Self::GroupedMapPandas => 201, + Self::GroupedAggPandas => 202, + Self::WindowAggPandas => 203, + Self::ScalarPandasIter => 204, + Self::MapPandasIter => 205, + Self::CogroupedMapPandas => 206, + Self::MapArrowIter => 207, + Self::GroupedMapArrow => 209, + Self::CogroupedMapArrow => 210, + Self::GroupedMapArrowIter => 215, + Self::GroupedMapPandasIter => 216, + Self::GroupedAggPandasIter => 217, Self::ScalarArrow => 250, + Self::ScalarArrowIter => 251, + Self::GroupedAggArrow => 252, + Self::WindowAggArrow => 253, + Self::GroupedAggArrowIter => 254, + Self::TableUdf => 300, + Self::ArrowTableUdf => 301, + Self::ArrowUdtf => 302, } } - /// The kind with Spark's code `code`, if the engine evaluates it. + /// The kind with Spark's code `code`, if a plan can name it. #[must_use] pub fn from_code(code: i32) -> Option { Self::ALL.into_iter().find(|t| t.code() == code) } + + /// What the UDF sees and returns. + #[must_use] + pub const fn kind(self) -> PythonUdfKind { + match self { + Self::ArrowBatched | Self::ScalarPandas | Self::ScalarArrow => PythonUdfKind::Scalar, + Self::ScalarPandasIter | Self::ScalarArrowIter => PythonUdfKind::ScalarIter, + Self::MapPandasIter | Self::MapArrowIter => PythonUdfKind::Map, + Self::GroupedMapPandas + | Self::GroupedMapArrow + | Self::GroupedMapArrowIter + | Self::GroupedMapPandasIter => PythonUdfKind::GroupedMap, + Self::CogroupedMapPandas | Self::CogroupedMapArrow => PythonUdfKind::CogroupedMap, + Self::GroupedAggPandas + | Self::GroupedAggPandasIter + | Self::GroupedAggArrow + | Self::GroupedAggArrowIter => PythonUdfKind::GroupedAgg, + Self::WindowAggPandas | Self::WindowAggArrow => PythonUdfKind::WindowAgg, + Self::TableUdf | Self::ArrowTableUdf | Self::ArrowUdtf => PythonUdfKind::Udtf, + } + } + + /// Whether a `PythonEval` evaluates it (101, 200, 250): the kinds of + /// design P2-T059. Every other kind is a relation kind (design P3-T085). + #[must_use] + pub const fn is_scalar(self) -> bool { + matches!(self.kind(), PythonUdfKind::Scalar) + } +} + +/// A reference to a session-registered Python UDF run by a relation node +/// (design P3-T085 §3.1): the same identity as +/// [`FunctionId::PythonUdf`], never the callable (I13). +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct PythonUdfRef { + /// The canonical (lower-case) name: the identity, with `version`. + pub name: String, + /// The registration version. + pub version: u32, + /// How the worker converts and calls it. + pub eval_type: PythonEvalType, + /// Whether calls may run concurrently; `false` serialises them. + pub thread_safe: bool, + /// Whether the UDF gives the same rows for the same input (Spark's + /// `PythonUDF.udfDeterministic`): a relation over one that does not is + /// nondeterministic ([`crate::python_relation_is_deterministic`]), so a + /// CTE holding it is computed once, not once per reference. Not part of + /// the identity, as a scalar call's determinism is the call's + /// ([`ResolvedFunction`](crate::ResolvedFunction)). + pub deterministic: bool, +} + +impl PythonUdfRef { + /// The function identity it stands for. + #[must_use] + pub fn function_id(&self) -> FunctionId { + FunctionId::PythonUdf { + name: self.name.clone(), + version: self.version, + eval_type: self.eval_type, + thread_safe: self.thread_safe, + } + } + + /// The reference of a [`FunctionId::PythonUdf`] called with + /// `deterministic` (the call's determinism); `None` for any other + /// function. + #[must_use] + pub fn of(id: &FunctionId, deterministic: bool) -> Option { + match id { + FunctionId::PythonUdf { + name, + version, + eval_type, + thread_safe, + } => Some(Self { + name: name.clone(), + version: *version, + eval_type: *eval_type, + thread_safe: *thread_safe, + deterministic, + }), + _ => None, + } + } } impl FunctionId { diff --git a/crates/krishiv-plan/src/lib.rs b/crates/krishiv-plan/src/lib.rs index 69b982b44..4c95eb442 100644 --- a/crates/krishiv-plan/src/lib.rs +++ b/crates/krishiv-plan/src/lib.rs @@ -132,6 +132,7 @@ mod literal; mod origin; mod partial; mod plan; +mod python; mod raises; pub(crate) mod spark_fmt; #[cfg(test)] @@ -149,7 +150,9 @@ pub use expr::{ SortOrder, Star, SubqueryKind, TypedLiteralKind, WindowExpr, WindowFrame, WindowSpec, WindowSpecRef, }; -pub use function::{Builtin, FunctionId, FunctionKind, PythonEvalType, QuantifiedLike}; +pub use function::{ + Builtin, FunctionId, FunctionKind, PythonEvalType, PythonUdfKind, PythonUdfRef, QuantifiedLike, +}; pub use literal::{DayTimeField, Literal, YearMonthField}; pub use origin::{Origin, PlanId, SourceId, Span}; pub use partial::partial_state_columns; @@ -171,8 +174,9 @@ pub fn generate_qualifier(index: usize, aliased: bool) -> String { pub use plan::{ ALIAS_PART_SEP, Cte, CteRef, GeneratorKind, Grouping, IR_VERSION, JoinCondition, JoinKind, JoinStrategy, LogicalPlan, PivotValue, Plan, PlanNode, ReadSource, SampleBounds, Scan, - SetOperation, SqlArgs, TABLE_FUNCTION_NAME_SEP, TableVersion, UnpivotColumnSet, generator_arg, - without_generated_marker, + SetOperation, SqlArgs, TABLE_FUNCTION_NAME_SEP, TableArg, TableVersion, UnpivotColumnSet, + generator_arg, without_generated_marker, }; +pub use python::{check_python_aggregates, is_python_udaf, python_relation_is_deterministic}; pub use raises::{conjuncts, exact_decimal, is_deterministic, may_raise, runs_on_any_row}; pub use walk::{MAX_DEPTH, check_depth}; diff --git a/crates/krishiv-plan/src/plan.rs b/crates/krishiv-plan/src/plan.rs index 5ab5fac78..a46aa7844 100644 --- a/crates/krishiv-plan/src/plan.rs +++ b/crates/krishiv-plan/src/plan.rs @@ -10,7 +10,7 @@ use krishiv_common::secret::OptionValue; use crate::dag::{self, Memo}; use crate::walk::{self, Node}; -use crate::{AttrId, BoundColumn, Expr, Origin, SortOrder}; +use crate::{AttrId, BoundColumn, Expr, Origin, PythonUdfRef, SortOrder}; /// The IR version a [`Plan`] is written in. It changes when a node changes /// meaning; adding a node or field does not change it before 1.0. @@ -514,6 +514,82 @@ pub enum PlanNode { /// The columns (empty until resolved). output: Vec, }, + /// `mapInPandas` / `mapInArrow` (Spark's `MapInPandas`, `MapInArrow`; + /// design P3-T085 §3.1): each partition of `input`, every column, goes to + /// the Python UDF as an iterator of batches; its rows are whatever the + /// function yields, typed as `output` declares. Resolved plans only. + MapPartitions { + /// Input. + input: Arc, + /// The UDF (a map kind: 205, 207). + udf: PythonUdfRef, + /// The declared result columns, with analyzer-assigned ids. + output: Vec, + /// Barrier execution mode (`isBarrier`): representable, refused by + /// the physical planner. + barrier: bool, + }, + /// `groupBy(keys).applyInPandas` / `applyInArrow` (Spark's + /// `FlatMapGroupsInPandas`, `FlatMapGroupsInArrow`): the rows of `input` + /// clustered by `keys` (one group without keys) and sorted by them; the + /// UDF sees each group's rows, every column, and yields any rows, typed as + /// `output` declares. Resolved plans only. + GroupMap { + /// Input. + input: Arc, + /// The grouping expressions over `input`. + keys: Vec>, + /// The UDF (a grouped-map kind: 201, 209, 215, 216). + udf: PythonUdfRef, + /// The declared result columns. + output: Vec, + }, + /// `groupBy(..).cogroup(groupBy(..)).applyInPandas` / `applyInArrow` + /// (Spark's `FlatMapCoGroupsInPandas`, `FlatMapCoGroupsInArrow`): the + /// UDF sees, per key, the left rows then the right rows of that key (an + /// empty side for a key on one side only). Resolved plans only. + CoGroupMap { + /// The left relation. + left: Arc, + /// The right relation. + right: Arc, + /// The grouping expressions over `left`. + left_keys: Vec>, + /// The grouping expressions over `right`, as many as `left_keys`. + right_keys: Vec>, + /// The UDF (a cogrouped-map kind: 206, 210). + udf: PythonUdfRef, + /// The declared result columns. + output: Vec, + }, + /// A Python table function in `FROM` (Spark's resolved `PythonUDTF` + /// call): `args` over no input, and at most one `TABLE` argument with + /// its partitioning. Resolved plans only. + PythonTableFunction { + /// The scalar arguments. + args: Vec>, + /// The `TABLE(...)` argument, if any. + table_arg: Option, + /// The UDF (a table kind: 300, 301, 302). + udf: PythonUdfRef, + /// The declared (or analyzed) result columns. + output: Vec, + }, +} + +/// The `TABLE(t)` argument of a [`PlanNode::PythonTableFunction`] (Spark's +/// `FunctionTableSubqueryArgumentExpression`): the relation and how its rows +/// reach the function. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TableArg { + /// The relation. + pub input: Arc, + /// `PARTITION BY` expressions over `input` (a hash exchange on them). + pub partition_by: Vec>, + /// `ORDER BY` within a partition (a local sort after the partitioning). + pub order_by: Vec, + /// `WITH SINGLE PARTITION`: every row in one partition. + pub single_partition: bool, } /// The attribute a resolved generator reads: its argument is always a @@ -543,6 +619,11 @@ pub enum GeneratorKind { /// a variant array (`pos`, `key`, `value`); any other value, a variant /// null and SQL `NULL` give no rows. VariantExplode, + /// A Python UDTF under `LATERAL` (Spark's `Generate(PythonUDTF)`): the + /// node's `arg` is the call, a [`crate::FunctionId::PythonUdf`] of a + /// table kind with the UDTF's arguments, not a column; the generated + /// columns are its declared result. + PythonUdtf, } impl GeneratorKind { @@ -555,6 +636,7 @@ impl GeneratorKind { Self::PosExplode => "posexplode", Self::Inline => "inline", Self::VariantExplode => "variant_explode", + Self::PythonUdtf => "python_udtf", } } @@ -1022,11 +1104,20 @@ impl LogicalPlan { | PlanNode::Unpivot { input, .. } | PlanNode::Window { input, .. } | PlanNode::Sample { input, .. } - | PlanNode::ReplicateRows { input, .. } => f(input), - PlanNode::Join { left, right, .. } | PlanNode::SetOp { left, right, .. } => { + | PlanNode::ReplicateRows { input, .. } + | PlanNode::MapPartitions { input, .. } + | PlanNode::GroupMap { input, .. } => f(input), + PlanNode::Join { left, right, .. } + | PlanNode::SetOp { left, right, .. } + | PlanNode::CoGroupMap { left, right, .. } => { f(left); f(right); } + PlanNode::PythonTableFunction { table_arg, .. } => { + if let Some(t) = table_arg { + f(&t.input); + } + } PlanNode::Union { inputs, .. } => inputs.iter().for_each(f), PlanNode::WithCte { input, ctes } => { ctes.iter().for_each(|c| f(&c.plan)); @@ -1073,11 +1164,20 @@ impl LogicalPlan { | PlanNode::Unpivot { input, .. } | PlanNode::Window { input, .. } | PlanNode::Sample { input, .. } - | PlanNode::ReplicateRows { input, .. } => f(input), - PlanNode::Join { left, right, .. } | PlanNode::SetOp { left, right, .. } => { + | PlanNode::ReplicateRows { input, .. } + | PlanNode::MapPartitions { input, .. } + | PlanNode::GroupMap { input, .. } => f(input), + PlanNode::Join { left, right, .. } + | PlanNode::SetOp { left, right, .. } + | PlanNode::CoGroupMap { left, right, .. } => { f(left); f(right); } + PlanNode::PythonTableFunction { table_arg, .. } => { + if let Some(t) = table_arg { + f(&mut t.input); + } + } PlanNode::Union { inputs, .. } => inputs.iter_mut().for_each(f), PlanNode::WithCte { input, ctes } => { ctes.iter_mut().for_each(|c| f(&mut c.plan)); @@ -1105,6 +1205,7 @@ impl LogicalPlan { | PlanNode::Union { .. } | PlanNode::SetOp { .. } | PlanNode::Sample { .. } + | PlanNode::MapPartitions { .. } | PlanNode::Distinct { .. } => {} PlanNode::Hint { params, .. } | PlanNode::Repartition { exprs: params, .. } => { params.iter().for_each(f); @@ -1161,6 +1262,21 @@ impl LogicalPlan { .for_each(f), PlanNode::InlineTable { rows, .. } => rows.iter().flatten().for_each(f), PlanNode::TableFunction { args, .. } => args.iter().for_each(f), + PlanNode::GroupMap { keys, .. } => keys.iter().for_each(f), + PlanNode::CoGroupMap { + left_keys, + right_keys, + .. + } => left_keys.iter().chain(right_keys).for_each(f), + PlanNode::PythonTableFunction { + args, table_arg, .. + } => { + args.iter().for_each(&mut *f); + if let Some(t) = table_arg { + t.partition_by.iter().for_each(&mut *f); + t.order_by.iter().for_each(|o| f(&o.expr)); + } + } PlanNode::Window { exprs, partition_by, @@ -1190,6 +1306,7 @@ impl LogicalPlan { | PlanNode::Union { .. } | PlanNode::SetOp { .. } | PlanNode::Sample { .. } + | PlanNode::MapPartitions { .. } | PlanNode::Distinct { .. } => {} PlanNode::Hint { params, .. } | PlanNode::Repartition { exprs: params, .. } => { params.iter_mut().for_each(f); @@ -1258,6 +1375,21 @@ impl LogicalPlan { PlanNode::TableFunction { args, .. } => { args.iter_mut().for_each(f); } + PlanNode::GroupMap { keys, .. } => keys.iter_mut().for_each(f), + PlanNode::CoGroupMap { + left_keys, + right_keys, + .. + } => left_keys.iter_mut().chain(right_keys).for_each(f), + PlanNode::PythonTableFunction { + args, table_arg, .. + } => { + args.iter_mut().for_each(&mut *f); + if let Some(t) = table_arg { + t.partition_by.iter_mut().for_each(&mut *f); + t.order_by.iter_mut().for_each(|o| f(&mut o.expr)); + } + } PlanNode::Window { exprs, partition_by, @@ -1337,7 +1469,12 @@ impl LogicalPlan { return Some(out); } PlanNode::Range { output, .. } => return Some(vec![output.clone()]), - PlanNode::UnionLoop { output, .. } | PlanNode::UnionLoopRef { output, .. } => { + PlanNode::UnionLoop { output, .. } + | PlanNode::UnionLoopRef { output, .. } + | PlanNode::MapPartitions { output, .. } + | PlanNode::GroupMap { output, .. } + | PlanNode::CoGroupMap { output, .. } + | PlanNode::PythonTableFunction { output, .. } => { return Some(output.clone()); } PlanNode::Scan(s) => return Some(s.output.clone()), @@ -1534,11 +1671,28 @@ impl LogicalPlan { impl PartialEq for LogicalPlan { fn eq(&self, other: &Self) -> bool { dag::eq(Node::Plan(self), Node::Plan(other), || { - self.origin == other.origin && self.node == other.node + self.origin == other.origin + && std::mem::discriminant(&self.node) == std::mem::discriminant(&other.node) + && children_eq(self, other) + && self.node == other.node }) } } +/// Whether the children of two nodes of one kind are equal, pair by pair. +/// Compared before the nodes' fields, so the derived `PlanNode` comparison +/// (a large frame in an unoptimised build) finds each child pair already +/// equal in the walk's memo instead of recursing: a chain of +/// [`MAX_DEPTH`](crate::MAX_DEPTH) nodes needs only this small frame per +/// level. +fn children_eq(a: &LogicalPlan, b: &LogicalPlan) -> bool { + let mut left = Vec::new(); + a.for_each_child(&mut |c| left.push(c)); + let mut right = Vec::new(); + b.for_each_child(&mut |c| right.push(c)); + left.len() == right.len() && left.iter().zip(&right).all(|(l, r)| l == r) +} + impl Eq for LogicalPlan {} impl fmt::Debug for LogicalPlan { diff --git a/crates/krishiv-plan/src/python.rs b/crates/krishiv-plan/src/python.rs new file mode 100644 index 000000000..82d0b379a --- /dev/null +++ b/crates/krishiv-plan/src/python.rs @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Python UDF relation kinds in plans (design P3-T085 §3.1): the checks +//! shared by the analyzer and the physical planner. + +use std::collections::HashSet; +use std::sync::Arc; + +use krishiv_common::error::{InvalidPandasUdfPlacement, KrishivError}; + +use crate::{Expr, ExprNode, FunctionId, FunctionKind, PlanNode, PythonUdfKind, ResolvedExpr}; + +/// Whether `e` is a call of a grouped-aggregate Python UDF (202, 217, 252, +/// 254). +#[must_use] +pub fn is_python_udaf(e: &Expr) -> bool { + matches!( + &e.node, + ExprNode::Resolved(ResolvedExpr::Function(f)) if matches!( + &f.id, + FunctionId::PythonUdf { eval_type, .. } if eval_type.kind() == PythonUdfKind::GroupedAgg + ) + ) +} + +/// Whether the Python relation UDF `node` runs, if it runs one, is +/// deterministic. A relation node's UDF is not an expression of it, so a +/// plan's determinism check asks this beside its expressions' (Spark's +/// `QueryPlan.deterministic` sees the node's `functionExpr`, a `PythonUDF`, +/// at v4.2.0). +#[must_use] +pub fn python_relation_is_deterministic(node: &PlanNode) -> bool { + match node { + PlanNode::MapPartitions { udf, .. } + | PlanNode::GroupMap { udf, .. } + | PlanNode::CoGroupMap { udf, .. } + | PlanNode::PythonTableFunction { udf, .. } => udf.deterministic, + _ => true, + } +} + +/// Refuses an aggregate whose expressions call a grouped-aggregate Python UDF +/// beside any other aggregate function (Spark's planner, +/// `SparkStrategies.scala` L670-671 at v4.2.0: `INVALID_PANDAS_UDF_PLACEMENT`, +/// the UDFs' names quoted as `toSQLId` quotes them). Called by the physical +/// planner, where Spark raises it: analysis accepts such an aggregate, so a +/// DataFrame is refused at its first action, not at creation. Subqueries +/// are not looked into. +/// +/// Known difference: the names are the canonical lower-case ones a plan +/// carries, where Spark prints the name as registered. +/// +/// # Errors +/// `INVALID_PANDAS_UDF_PLACEMENT` (0A000). +pub fn check_python_aggregates<'a>( + aggregates: impl IntoIterator>, +) -> Result<(), KrishivError> { + let mut udafs: Vec = Vec::new(); + let mut other = false; + let mut seen = HashSet::new(); + // Depth first, in written order: the names are listed as first seen. + let mut stack: Vec<&Expr> = aggregates.into_iter().map(AsRef::as_ref).collect(); + stack.reverse(); + while let Some(e) = stack.pop() { + if !seen.insert(std::ptr::from_ref(e).addr()) { + continue; + } + if let ExprNode::Resolved(ResolvedExpr::Function(f)) = &e.node { + if is_python_udaf(e) { + let name = crate::function::shown_call_name("", f.id.name()); + let quoted = format!("`{}`", name.replace('`', "``")); + if !udafs.contains("ed) { + udafs.push(quoted); + } + continue; + } + if f.kind == FunctionKind::Aggregate { + other = true; + } + } + let at = stack.len(); + e.for_each_child(&mut |c| stack.push(c)); + stack[at..].reverse(); + } + if other && !udafs.is_empty() { + return Err(InvalidPandasUdfPlacement { + function_list: udafs.join(", "), + } + .into()); + } + Ok(()) +} diff --git a/crates/krishiv-plan/src/tests/inventory.rs b/crates/krishiv-plan/src/tests/inventory.rs index c92066c72..4943e4489 100644 --- a/crates/krishiv-plan/src/tests/inventory.rs +++ b/crates/krishiv-plan/src/tests/inventory.rs @@ -63,6 +63,10 @@ fn kinds_of_node(plan: &LogicalPlan, out: &mut BTreeSet<&'static str>) { PlanNode::Range { .. } => "PlanNode::Range", PlanNode::UnionLoop { .. } => "PlanNode::UnionLoop", PlanNode::UnionLoopRef { .. } => "PlanNode::UnionLoopRef", + PlanNode::MapPartitions { .. } => "PlanNode::MapPartitions", + PlanNode::GroupMap { .. } => "PlanNode::GroupMap", + PlanNode::CoGroupMap { .. } => "PlanNode::CoGroupMap", + PlanNode::PythonTableFunction { .. } => "PlanNode::PythonTableFunction", }); } diff --git a/crates/krishiv-plan/src/tests/mod.rs b/crates/krishiv-plan/src/tests/mod.rs index 8cc8a572a..e50b548fc 100644 --- a/crates/krishiv-plan/src/tests/mod.rs +++ b/crates/krishiv-plan/src/tests/mod.rs @@ -19,6 +19,7 @@ mod nodes; mod pivot; mod prefix_cache; mod props; +mod python_relations; mod rewrite; mod secrets; mod setops; diff --git a/crates/krishiv-plan/src/tests/props.rs b/crates/krishiv-plan/src/tests/props.rs index e9e7e3587..5399cf3ab 100644 --- a/crates/krishiv-plan/src/tests/props.rs +++ b/crates/krishiv-plan/src/tests/props.rs @@ -1111,6 +1111,55 @@ fn m_plan(p: &mut Arc, w: &mut Walk) -> bool { site!(w, *id += 1); false } + PlanNode::MapPartitions { + input, + udf, + output, + barrier, + } => { + site!(w, udf.version += 1); + site!(w, *barrier = !*barrier); + m_plan(input, w) || output.iter_mut().any(|c| m_bound(c, w)) + } + PlanNode::GroupMap { + input, + keys, + udf, + output, + } => { + site!(w, udf.version += 1); + m_plan(input, w) || m_exprs(keys, w) || output.iter_mut().any(|c| m_bound(c, w)) + } + PlanNode::CoGroupMap { + left, + right, + left_keys, + right_keys, + udf, + output, + } => { + site!(w, udf.version += 1); + m_plan(left, w) + || m_plan(right, w) + || m_exprs(left_keys, w) + || m_exprs(right_keys, w) + || output.iter_mut().any(|c| m_bound(c, w)) + } + PlanNode::PythonTableFunction { + args, + table_arg, + udf, + output, + } => { + site!(w, udf.version += 1); + site!(w, *table_arg = None); + m_exprs(args, w) + || output.iter_mut().any(|c| m_bound(c, w)) + || table_arg.as_mut().is_some_and(|t| { + site!(w, t.single_partition = !t.single_partition); + m_plan(&mut t.input, w) || m_exprs(&mut t.partition_by, w) + }) + } PlanNode::UnionLoop { id, anchor, diff --git a/crates/krishiv-plan/src/tests/python_relations.rs b/crates/krishiv-plan/src/tests/python_relations.rs new file mode 100644 index 000000000..2e7e14fec --- /dev/null +++ b/crates/krishiv-plan/src/tests/python_relations.rs @@ -0,0 +1,354 @@ +// SPDX-License-Identifier: Apache-2.0 +//! The Python UDF relation nodes and eval types (design P3-T085 §3.1): +//! codes, kinds, traversal, outputs, encoding, printing and the mixed +//! aggregate check. + +use std::collections::BTreeMap; +use std::sync::Arc; + +use krishiv_common::error::ErrorClass; +use krishiv_types::DataType; + +use crate::{ + AttrId, BoundColumn, Builtin, Expr, ExprNode, FunctionId, FunctionKind, GeneratorKind, + LogicalPlan, Plan, PlanNode, PythonEvalType, PythonUdfKind, PythonUdfRef, ReadSource, + ResolvedExpr, ResolvedFunction, Scan, SortOrder, TableArg, check_python_aggregates, +}; + +fn col(id: u64, name: &str) -> BoundColumn { + BoundColumn { + id: AttrId(id), + name: name.to_owned(), + data_type: DataType::Long, + nullable: true, + } +} + +fn column(id: u64) -> Arc { + Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Column(col( + id, "k", + ))))) +} + +fn scan(first: u64) -> Arc { + Arc::new(LogicalPlan::new(PlanNode::Scan(Scan { + source: ReadSource::Table { + name: vec!["t".to_owned()], + options: BTreeMap::new(), + }, + output: vec![col(first, "k"), col(first + 1, "v")], + version: None, + }))) +} + +fn udf(eval_type: PythonEvalType) -> PythonUdfRef { + PythonUdfRef { + name: "f".to_owned(), + version: 1, + eval_type, + thread_safe: true, + deterministic: true, + } +} + +fn call(id: FunctionId, kind: FunctionKind, args: Vec>) -> Arc { + Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Function( + ResolvedFunction { + id, + kind, + args, + distinct: false, + filter: None, + within_group: Vec::new(), + data_type: DataType::Double, + nullable: true, + deterministic: true, + }, + )))) +} + +fn nodes() -> Vec { + let out = vec![col(10, "x"), col(11, "y")]; + vec![ + LogicalPlan::new(PlanNode::MapPartitions { + input: scan(1), + udf: udf(PythonEvalType::MapPandasIter), + output: out.clone(), + barrier: false, + }), + LogicalPlan::new(PlanNode::GroupMap { + input: scan(1), + keys: vec![column(1)], + udf: udf(PythonEvalType::GroupedMapArrow), + output: out.clone(), + }), + LogicalPlan::new(PlanNode::CoGroupMap { + left: scan(1), + right: scan(3), + left_keys: vec![column(1)], + right_keys: vec![column(3)], + udf: udf(PythonEvalType::CogroupedMapPandas), + output: out.clone(), + }), + LogicalPlan::new(PlanNode::PythonTableFunction { + args: vec![Arc::new(Expr::from(crate::Literal::Integer(1)))], + table_arg: Some(TableArg { + input: scan(1), + partition_by: vec![column(1)], + order_by: vec![SortOrder::asc(column(2))], + single_partition: false, + }), + udf: udf(PythonEvalType::ArrowUdtf), + output: out, + }), + ] +} + +#[test] +fn eval_type_codes_are_sparks_and_round_trip() { + let codes: Vec = PythonEvalType::ALL.iter().map(|t| t.code()).collect(); + assert_eq!( + codes, + [ + 101, 200, 201, 202, 203, 204, 205, 206, 207, 209, 210, 215, 216, 217, 250, 251, 252, + 253, 254, 300, 301, 302 + ] + ); + for t in PythonEvalType::ALL { + assert_eq!(PythonEvalType::from_code(t.code()), Some(t)); + } + // Streaming state (P4) and the plain batched UDF are not plan kinds. + for code in [100, 208, 211, 212, 213, 214, 0, 999] { + assert_eq!(PythonEvalType::from_code(code), None, "{code}"); + } +} + +#[test] +fn eval_types_classify_by_what_the_udf_sees() { + let kind = |c| PythonEvalType::from_code(c).map(PythonEvalType::kind); + for c in [101, 200, 250] { + assert_eq!(kind(c), Some(PythonUdfKind::Scalar)); + assert!(PythonEvalType::from_code(c).is_some_and(PythonEvalType::is_scalar)); + } + for (codes, k) in [ + (&[204, 251][..], PythonUdfKind::ScalarIter), + (&[205, 207], PythonUdfKind::Map), + (&[201, 209, 215, 216], PythonUdfKind::GroupedMap), + (&[206, 210], PythonUdfKind::CogroupedMap), + (&[202, 217, 252, 254], PythonUdfKind::GroupedAgg), + (&[203, 253], PythonUdfKind::WindowAgg), + (&[300, 301, 302], PythonUdfKind::Udtf), + ] { + for &c in codes { + assert_eq!(kind(c), Some(k), "{c}"); + assert!(!PythonEvalType::from_code(c).is_some_and(PythonEvalType::is_scalar)); + } + } +} + +#[test] +fn a_udf_reference_is_its_function_identity() { + let r = udf(PythonEvalType::GroupedAggArrow); + let id = r.function_id(); + assert_eq!(PythonUdfRef::of(&id, true), Some(r.clone())); + assert_eq!( + PythonUdfRef::of(&id, false), + Some(PythonUdfRef { + deterministic: false, + ..r + }) + ); + assert_eq!( + PythonUdfRef::of(&FunctionId::Builtin(Builtin::Add), true), + None + ); +} + +/// A relation node is as deterministic as its UDF; any other node has none. +#[test] +fn a_relation_node_is_as_deterministic_as_its_udf() { + for node in nodes() { + assert!( + crate::python_relation_is_deterministic(&node.node), + "{node}" + ); + let mut node = node.node.clone(); + match &mut node { + PlanNode::MapPartitions { udf, .. } + | PlanNode::GroupMap { udf, .. } + | PlanNode::CoGroupMap { udf, .. } + | PlanNode::PythonTableFunction { udf, .. } => udf.deterministic = false, + other => panic!("not a relation node: {other:?}"), + } + assert!(!crate::python_relation_is_deterministic(&node), "{node:?}"); + } + assert!(crate::python_relation_is_deterministic(&scan(1).node)); +} + +#[test] +fn relation_nodes_walk_their_children_and_keys_and_output_their_declared_columns() { + let expected_children = [1, 1, 2, 1]; + let expected_exprs = [0, 1, 2, 3]; + for ((plan, children), exprs) in nodes().iter().zip(expected_children).zip(expected_exprs) { + let mut n = 0; + plan.for_each_child(&mut |_| n += 1); + assert_eq!(n, children, "{plan}"); + let mut e = 0; + plan.for_each_expr(&mut |_| e += 1); + assert_eq!(e, exprs, "{plan}"); + let out = plan.resolved_output().unwrap_or_default(); + assert_eq!( + out.iter().map(|c| c.id.0).collect::>(), + [10, 11], + "{plan}" + ); + assert!(plan.is_resolved(), "{plan}"); + // A rewrite that changes nothing keeps the node. + let same = Arc::new(plan.clone()) + .transform_up(&mut |p| Ok::<_, std::convert::Infallible>(p)) + .unwrap(); + assert_eq!(*same, *plan); + } +} + +#[test] +fn relation_nodes_encode_apart_by_kind_udf_and_barrier() { + let plans: Vec = nodes() + .into_iter() + .map(|p| Plan::try_new(p).unwrap()) + .collect(); + for (i, a) in plans.iter().enumerate() { + for (j, b) in plans.iter().enumerate() { + assert_eq!(a.structurally_eq(b), i == j); + } + } + let with = |version: u32, barrier: bool| { + let mut r = udf(PythonEvalType::MapArrowIter); + r.version = version; + Plan::try_new(LogicalPlan::new(PlanNode::MapPartitions { + input: scan(1), + udf: r, + output: vec![col(10, "x")], + barrier, + })) + .unwrap() + }; + let base = with(1, false); + assert_ne!(base.structural_hash(), with(2, false).structural_hash()); + assert_ne!(base.structural_hash(), with(1, true).structural_hash()); + assert_eq!(base.structural_hash(), with(1, false).structural_hash()); +} + +#[test] +fn relation_nodes_print_their_udf_and_columns() { + let text: Vec = nodes() + .iter() + .map(|p| p.to_string().lines().next().unwrap_or_default().to_owned()) + .collect(); + assert_eq!( + text[0], + "MapPartitions f#1(205), [x#10, y#11], barrier=false" + ); + assert_eq!(text[1], "GroupMap [k#1], f#1(209), [x#10, y#11]"); + assert_eq!(text[2], "CoGroupMap [k#1], [k#3], f#1(206), [x#10, y#11]"); + assert!( + text[3] + .starts_with("PythonTableFunction f#1(302)(1), [x#10, y#11], TABLE PARTITION BY [k#1]"), + "{}", + text[3] + ); +} + +#[test] +fn a_lateral_python_udtf_carries_its_call_and_keeps_the_input_columns() { + let udtf = call( + udf(PythonEvalType::TableUdf).function_id(), + FunctionKind::Generator, + vec![column(1)], + ); + let plan = LogicalPlan::new(PlanNode::Generate { + input: scan(1), + kind: GeneratorKind::PythonUdtf, + outer: false, + arg: udtf, + output: vec![col(10, "x")], + names: Vec::new(), + qualifier: None, + }); + assert_eq!(GeneratorKind::PythonUdtf.name(), "python_udtf"); + let out = plan.resolved_output().unwrap_or_default(); + assert_eq!(out.iter().map(|c| c.id.0).collect::>(), [1, 2, 10]); + let mut explode = plan.clone(); + if let PlanNode::Generate { kind, .. } = &mut explode.node { + *kind = GeneratorKind::Explode; + } + assert!( + !Plan::try_new(plan) + .unwrap() + .structurally_eq(&Plan::try_new(explode).unwrap()) + ); +} + +#[test] +fn a_python_udaf_beside_a_builtin_aggregate_is_refused() { + let udaf = |name: &str| { + let mut r = udf(PythonEvalType::GroupedAggPandas); + r.name = name.to_owned(); + call(r.function_id(), FunctionKind::Aggregate, vec![column(2)]) + }; + let builtin = call( + FunctionId::Registry { + name: "sum".to_owned(), + id: 0, + }, + FunctionKind::Aggregate, + vec![column(2)], + ); + let scalar_python = call( + udf(PythonEvalType::ScalarPandas).function_id(), + FunctionKind::Scalar, + vec![column(2)], + ); + assert!(check_python_aggregates(&[udaf("f"), udaf("g")]).is_ok()); + assert!(check_python_aggregates(&[Arc::clone(&builtin), scalar_python]).is_ok()); + let e = + check_python_aggregates(&[udaf("mean_udf"), Arc::clone(&builtin), udaf("g")]).unwrap_err(); + assert_eq!(e.error_class(), ErrorClass::InvalidPandasUdfPlacement); + assert_eq!(e.sql_state(), Some("0A000")); + assert!(e.to_string().contains("`mean_udf`, `g`"), "{e}"); + // Nested in an expression too. + let plus = call( + FunctionId::Builtin(Builtin::Add), + FunctionKind::Scalar, + vec![udaf("f"), builtin], + ); + assert!(check_python_aggregates(&[plus]).is_err()); +} + +#[test] +fn rewrites_reach_the_relation_nodes_children_and_expressions() { + let one = Arc::new(Expr::from(crate::Literal::Integer(1))); + for plan in nodes() { + let plan = Arc::new(plan); + let mut children = 0; + let rebuilt = Arc::clone(&plan) + .map_children(|c| { + children += 1; + Ok::<_, std::convert::Infallible>(Arc::new((*c).clone())) + }) + .unwrap(); + assert!(children > 0, "{plan}"); + assert_eq!(*rebuilt, *plan); + let mut exprs = 0; + let rewritten = Arc::clone(&plan) + .map_expressions(|_| { + exprs += 1; + Ok::<_, std::convert::Infallible>(Arc::clone(&one)) + }) + .unwrap(); + let mut after = Vec::new(); + rewritten.for_each_expr(&mut |e| after.push(Arc::clone(e))); + assert_eq!(after.len(), exprs, "{plan}"); + assert!(after.iter().all(|e| *e == one), "{plan}"); + } +} diff --git a/crates/krishiv-plan/src/walk.rs b/crates/krishiv-plan/src/walk.rs index 136565e9d..51cbc1a89 100644 --- a/crates/krishiv-plan/src/walk.rs +++ b/crates/krishiv-plan/src/walk.rs @@ -203,6 +203,10 @@ pub(crate) fn all_resolved(root: Node<'_>) -> bool { | PlanNode::Expand { .. } | PlanNode::Window { .. } | PlanNode::Sample { .. } + | PlanNode::MapPartitions { .. } + | PlanNode::GroupMap { .. } + | PlanNode::CoGroupMap { .. } + | PlanNode::PythonTableFunction { .. } | PlanNode::ReplicateRows { .. } => true, // Lowering builds the generator without its columns. PlanNode::Generate { output, .. } @@ -441,6 +445,33 @@ impl Teardown { self.exprs.extend(params); } PlanNode::SetOp { left, right, .. } => self.plans.extend([left, right]), + PlanNode::MapPartitions { input, .. } => self.plans.push(input), + PlanNode::GroupMap { input, keys, .. } => { + self.plans.push(input); + self.exprs.extend(keys); + } + PlanNode::CoGroupMap { + left, + right, + left_keys, + right_keys, + .. + } => { + self.plans.extend([left, right]); + self.exprs.extend(left_keys); + self.exprs.extend(right_keys); + } + PlanNode::PythonTableFunction { + args, table_arg, .. + } => { + self.exprs.extend(args); + if let Some(t) = table_arg { + self.plans.push(t.input); + self.exprs.extend(t.partition_by); + self.exprs + .extend(t.order_by.into_iter().map(|SortOrder { expr, .. }| expr)); + } + } PlanNode::Union { inputs, .. } => self.plans.extend(inputs), PlanNode::WithCte { input, ctes } => { self.plans.push(input); diff --git a/crates/krishiv-session/src/catalog.rs b/crates/krishiv-session/src/catalog.rs index a83543671..5241abc7d 100644 --- a/crates/krishiv-session/src/catalog.rs +++ b/crates/krishiv-session/src/catalog.rs @@ -961,6 +961,10 @@ fn streaming_leaves(root: &PhysNode) -> (Vec<&krishiv_exec::physical::ScanNode>, | PhysNode::WindowGroupLimit { input, .. } | PhysNode::Aggregate { input, .. } | PhysNode::PythonEval { input, .. } + | PhysNode::PythonMap { input, .. } + | PhysNode::PythonGroupMap { input, .. } + | PhysNode::PythonCoGroupMap { input, .. } + | PhysNode::PythonAggregate { input, .. } | PhysNode::WriteFiles { input, .. } => go(input, out), PhysNode::Exchange(e) => go(&e.input, out), PhysNode::HashJoin(join) => go(&join.probe, out), diff --git a/proto/krishiv/plan/v1/plan.proto b/proto/krishiv/plan/v1/plan.proto index 52114b23c..ad238621e 100644 --- a/proto/krishiv/plan/v1/plan.proto +++ b/proto/krishiv/plan/v1/plan.proto @@ -174,6 +174,14 @@ message Node { PythonEvalNode python_eval = 26; // feature: sample SampleNode sample = 27; + // feature: python_udf_relations + PythonMapNode python_map = 28; + // feature: python_udf_relations + PythonGroupMapNode python_group_map = 29; + // feature: python_udf_relations + PythonCoGroupMapNode python_co_group_map = 30; + // feature: python_udf_relations + PythonAggregateNode python_aggregate = 31; } } @@ -189,6 +197,69 @@ message SampleNode { int64 seed = 4; } +// feature: python_udf_relations +// `mapInPandas` / `mapInArrow`: each partition of `input` sent to a Python +// UDF task. +message PythonMapNode { + uint32 input = 1; + PythonRelation udf = 2; +} + +// feature: python_udf_relations +// A grouped map over `input` clustered and sorted by the key columns. +message PythonGroupMapNode { + uint32 input = 1; + repeated uint64 keys = 2; + PythonRelation udf = 3; +} + +// feature: python_udf_relations +// A cogrouped map over both relations in one stream, sorted by the key +// columns then the side column. +message PythonCoGroupMapNode { + uint32 input = 1; + repeated uint64 keys = 2; + uint64 side = 3; + PythonRelation udf = 4; +} + +// feature: python_udf_relations +// An aggregate of grouped-aggregate Python UDF calls over `input` clustered +// and sorted by the key columns (`keys`: a collated key's comparison key); +// each group shows its first row's `key_values`. +message PythonAggregateNode { + uint32 input = 1; + repeated uint64 layout = 2; + repeated uint64 keys = 3; + repeated PythonCall calls = 4; + repeated uint64 key_values = 5; +} + +// feature: python_udf_relations +// The Python UDF of a relation node: by name and version, the columns of its +// frames per side, and its declared result. +message PythonRelation { + PythonUdfFunction udf = 1; + repeated PythonFrameSide sides = 2; + repeated PythonOutputColumn output = 3; + // Whether the UDF is deterministic (Spark's `udfDeterministic`). + bool deterministic = 4; +} + +// feature: python_udf_relations +message PythonFrameSide { + repeated uint64 columns = 1; + repeated string names = 2; + repeated DataType types = 3; +} + +// feature: python_udf_relations +message PythonOutputColumn { + string name = 1; + DataType data_type = 2; + bool nullable = 3; +} + // feature: python_udf // The rows of `input`, then one column per Python UDF call over every row. message PythonEvalNode { @@ -379,6 +450,8 @@ enum GeneratorKind { GENERATOR_KIND_INLINE = 3; // feature: variant_explode GENERATOR_KIND_VARIANT_EXPLODE = 4; + // feature: python_udf_relations + GENERATOR_KIND_PYTHON_UDTF = 5; } // feature: expand