From 67c5e53e232a3e5fa3da64e27c83159410931d50 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Fri, 13 Nov 2026 11:37:27 +0000 Subject: [PATCH 01/29] =?UTF-8?q?plan:=20Python=20UDF=20relation=20IR=20an?= =?UTF-8?q?d=20eval=20types=20(P3-T014a,=20T085=20=C2=A73.1)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PythonEvalType gains Spark's relation-kind codes (201-217, 251-254, 300-302; pyspark/util.py at 4.2.0) with a kind classifier and is_scalar. The logical IR gains MapPartitions, GroupMap, CoGroupMap, PythonTableFunction (with its TABLE argument) and GeneratorKind::PythonUdtf; each references its UDF by PythonUdfRef, never the callable. Walkers, printer, canonical encoding and cache fingerprint cover them. An aggregate mixing a grouped-aggregate Python UDF with another aggregate is refused with Spark's INVALID_PANDAS_UDF_PLACEMENT (0A000). The plan codec refuses the Python UDTF generator until its feature lands. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-plan-codec/src/encode.rs | 5 + crates/krishiv-plan-codec/src/features.rs | 1 + .../src/analyzer/ops/aggregate.rs | 2 + crates/krishiv-plan/src/analyzer/resolve.rs | 13 + crates/krishiv-plan/src/cache.rs | 66 ++++ crates/krishiv-plan/src/display.rs | 73 +++++ crates/krishiv-plan/src/encode.rs | 56 ++++ crates/krishiv-plan/src/function.rs | 200 +++++++++++- crates/krishiv-plan/src/lib.rs | 8 +- crates/krishiv-plan/src/plan.rs | 149 ++++++++- crates/krishiv-plan/src/python.rs | 70 +++++ crates/krishiv-plan/src/tests/inventory.rs | 4 + crates/krishiv-plan/src/tests/mod.rs | 1 + crates/krishiv-plan/src/tests/props.rs | 49 +++ .../src/tests/python_relations.rs | 294 ++++++++++++++++++ crates/krishiv-plan/src/walk.rs | 31 ++ 16 files changed, 1009 insertions(+), 13 deletions(-) create mode 100644 crates/krishiv-plan/src/python.rs create mode 100644 crates/krishiv-plan/src/tests/python_relations.rs diff --git a/crates/krishiv-plan-codec/src/encode.rs b/crates/krishiv-plan-codec/src/encode.rs index 07f76e547..bd720057f 100644 --- a/crates/krishiv-plan-codec/src/encode.rs +++ b/crates/krishiv-plan-codec/src/encode.rs @@ -351,6 +351,11 @@ impl Encoder { krishiv_plan::GeneratorKind::VariantExplode => { pb::GeneratorKind::VariantExplode } + krishiv_plan::GeneratorKind::PythonUdtf => { + return Err(CodecError::Unsupported( + "a Python UDTF generator".to_owned(), + )); + } } .into(), outer: *outer, diff --git a/crates/krishiv-plan-codec/src/features.rs b/crates/krishiv-plan-codec/src/features.rs index c4a6d5eb1..c4294b598 100644 --- a/crates/krishiv-plan-codec/src/features.rs +++ b/crates/krishiv-plan-codec/src/features.rs @@ -467,6 +467,7 @@ pub fn feature_of(construct: Construct<'_>) -> Option { Construct::ProbeSelect => Some(Feature::FusedProbe), Construct::Generator(kind) => match kind { krishiv_plan::GeneratorKind::VariantExplode => Some(Feature::VariantExplode), + krishiv_plan::GeneratorKind::PythonUdtf => Some(Feature::Generate), krishiv_plan::GeneratorKind::Explode | krishiv_plan::GeneratorKind::PosExplode | krishiv_plan::GeneratorKind::Inline => None, diff --git a/crates/krishiv-plan/src/analyzer/ops/aggregate.rs b/crates/krishiv-plan/src/analyzer/ops/aggregate.rs index 5ab1c5cfa..b8b7a0fa1 100644 --- a/crates/krishiv-plan/src/analyzer/ops/aggregate.rs +++ b/crates/krishiv-plan/src/analyzer/ops/aggregate.rs @@ -344,6 +344,8 @@ impl Resolver<'_> { } } let items = self.select_items(&c, aggregates, env, "Aggregate")?; + crate::check_python_aggregates(items.iter().map(|b| &b.expr)) + .map_err(|e| AnalysisError::new(e, p.origin))?; // The declared `CHAR`/`VARCHAR` types follow the aliases of an explicit `GROUP BY` too. self.carry_declared(&items, aggregates); // Spark's `ExtractGenerator`: the generator's argument takes the diff --git a/crates/krishiv-plan/src/analyzer/resolve.rs b/crates/krishiv-plan/src/analyzer/resolve.rs index 843d2377e..e46e61899 100644 --- a/crates/krishiv-plan/src/analyzer/resolve.rs +++ b/crates/krishiv-plan/src/analyzer/resolve.rs @@ -883,6 +883,10 @@ impl Resolver<'_> { | PlanNode::ReplicateRows { .. } | PlanNode::Expand { .. } | PlanNode::Range { .. } + | PlanNode::MapPartitions { .. } + | PlanNode::GroupMap { .. } + | PlanNode::CoGroupMap { .. } + | PlanNode::PythonTableFunction { .. } | PlanNode::JoinHint { .. } => self.leaf_node(p, env), } } @@ -1062,6 +1066,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/cache.rs b/crates/krishiv-plan/src/cache.rs index e08f9908c..429fa5afc 100644 --- a/crates/krishiv-plan/src/cache.rs +++ b/crates/krishiv-plan/src/cache.rs @@ -573,6 +573,15 @@ impl Encoder { } /// Writes `plan` as the digest of its encoding, computed once per node. + /// 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); + } + } + fn plan(&mut self, plan: &LogicalPlan, bias: u64) { let key = (NodeKind::Plan, std::ptr::from_ref(plan) as usize, bias); let digest = match self.seen.get(&key) { @@ -852,6 +861,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); diff --git a/crates/krishiv-plan/src/display.rs b/crates/krishiv-plan/src/display.rs index be51c22d1..b6c53c001 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); @@ -512,6 +531,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 582c68f8f..01ef8b79c 100644 --- a/crates/krishiv-plan/src/encode.rs +++ b/crates/krishiv-plan/src/encode.rs @@ -113,6 +113,10 @@ pub(crate) fn plan_tag(node: &PlanNode) -> u8 { PlanNode::Generate { .. } => 33, PlanNode::PipeOperator { .. } => 34, PlanNode::Repartition { .. } => 35, + PlanNode::MapPartitions { .. } => 36, + PlanNode::GroupMap { .. } => 37, + PlanNode::CoGroupMap { .. } => 38, + PlanNode::PythonTableFunction { .. } => 39, } } @@ -648,6 +652,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); @@ -1030,6 +1041,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); @@ -1048,6 +1060,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..d2938b83b 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,110 @@ 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, +} + +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`]; `None` for any other + /// function. + #[must_use] + pub fn of(id: &FunctionId) -> 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, + }), + _ => None, + } + } } impl FunctionId { diff --git a/crates/krishiv-plan/src/lib.rs b/crates/krishiv-plan/src/lib.rs index 16de1d556..66799672e 100644 --- a/crates/krishiv-plan/src/lib.rs +++ b/crates/krishiv-plan/src/lib.rs @@ -131,6 +131,7 @@ mod literal; mod origin; mod partial; mod plan; +mod python; mod raises; pub(crate) mod spark_fmt; #[cfg(test)] @@ -148,7 +149,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; @@ -170,8 +173,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, Scan, SetOperation, SqlArgs, - TABLE_FUNCTION_NAME_SEP, TableVersion, UnpivotColumnSet, generator_arg, + TABLE_FUNCTION_NAME_SEP, TableArg, TableVersion, UnpivotColumnSet, generator_arg, without_generated_marker, }; +pub use python::{check_python_aggregates, is_python_udaf}; 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 4a377a245..c4b503208 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. @@ -481,6 +481,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 @@ -510,6 +586,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 { @@ -522,6 +603,7 @@ impl GeneratorKind { Self::PosExplode => "posexplode", Self::Inline => "inline", Self::VariantExplode => "variant_explode", + Self::PythonUdtf => "python_udtf", } } @@ -988,11 +1070,20 @@ impl LogicalPlan { | PlanNode::Pivot { input, .. } | PlanNode::Unpivot { input, .. } | PlanNode::Window { 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)); @@ -1038,11 +1129,20 @@ impl LogicalPlan { | PlanNode::Pivot { input, .. } | PlanNode::Unpivot { input, .. } | PlanNode::Window { 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)); @@ -1069,6 +1169,7 @@ impl LogicalPlan { | PlanNode::JoinHint { .. } | PlanNode::Union { .. } | PlanNode::SetOp { .. } + | PlanNode::MapPartitions { .. } | PlanNode::Distinct { .. } => {} PlanNode::Hint { params, .. } | PlanNode::Repartition { exprs: params, .. } => { params.iter().for_each(f); @@ -1125,6 +1226,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, @@ -1153,6 +1269,7 @@ impl LogicalPlan { | PlanNode::JoinHint { .. } | PlanNode::Union { .. } | PlanNode::SetOp { .. } + | PlanNode::MapPartitions { .. } | PlanNode::Distinct { .. } => {} PlanNode::Hint { params, .. } | PlanNode::Repartition { exprs: params, .. } => { params.iter_mut().for_each(f); @@ -1221,6 +1338,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, @@ -1299,7 +1431,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()), diff --git a/crates/krishiv-plan/src/python.rs b/crates/krishiv-plan/src/python.rs new file mode 100644 index 000000000..4b0025628 --- /dev/null +++ b/crates/krishiv-plan/src/python.rs @@ -0,0 +1,70 @@ +// 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, 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 + ) + ) +} + +/// 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). Subqueries are not +/// looked into. +/// +/// # 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 36d8e88a4..8f3f2a566 100644 --- a/crates/krishiv-plan/src/tests/inventory.rs +++ b/crates/krishiv-plan/src/tests/inventory.rs @@ -62,6 +62,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 9b37360c9..a0fea2ad0 100644 --- a/crates/krishiv-plan/src/tests/mod.rs +++ b/crates/krishiv-plan/src/tests/mod.rs @@ -17,6 +17,7 @@ mod json_tuple; mod nodes; mod pivot; 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 6f9c97799..5135f363f 100644 --- a/crates/krishiv-plan/src/tests/props.rs +++ b/crates/krishiv-plan/src/tests/props.rs @@ -1103,6 +1103,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..eebedf43d --- /dev/null +++ b/crates/krishiv-plan/src/tests/python_relations.rs @@ -0,0 +1,294 @@ +// 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, + } +} + +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), Some(r)); + assert_eq!(PythonUdfRef::of(&FunctionId::Builtin(Builtin::Add)), None); +} + +#[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()); +} diff --git a/crates/krishiv-plan/src/walk.rs b/crates/krishiv-plan/src/walk.rs index a85d50358..f4b6aef07 100644 --- a/crates/krishiv-plan/src/walk.rs +++ b/crates/krishiv-plan/src/walk.rs @@ -167,6 +167,10 @@ pub(crate) fn all_resolved(root: Node<'_>) -> bool { | PlanNode::Repartition { .. } | PlanNode::Expand { .. } | PlanNode::Window { .. } + | PlanNode::MapPartitions { .. } + | PlanNode::GroupMap { .. } + | PlanNode::CoGroupMap { .. } + | PlanNode::PythonTableFunction { .. } | PlanNode::ReplicateRows { .. } => true, // Lowering builds the generator without its columns. PlanNode::Generate { output, .. } @@ -404,6 +408,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); From 49ca0b5d6c38f12d66f8cc2d9d4506c4575bc5c7 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Fri, 13 Nov 2026 11:45:26 +0000 Subject: [PATCH 02/29] =?UTF-8?q?optimizer:=20Python=20UDF=20relations=20a?= =?UTF-8?q?re=20opaque=20(P3-T014a,=20T085=20=C2=A73.1)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit No filter moves through MapPartitions, GroupMap or CoGroupMap (the pushdown's default arm keeps it above), and column pruning reads every column of their inputs: the UDF sees them all (group_ops.py L305 at 4.2.0; ADR 0022 §5). Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-optimizer/src/prune.rs | 12 ++ crates/krishiv-optimizer/tests/it/main.rs | 1 + .../tests/it/python_relations.rs | 132 ++++++++++++++++++ 3 files changed, 145 insertions(+) create mode 100644 crates/krishiv-optimizer/tests/it/python_relations.rs diff --git a/crates/krishiv-optimizer/src/prune.rs b/crates/krishiv-optimizer/src/prune.rs index 708558638..3a15a8b4c 100644 --- a/crates/krishiv-optimizer/src/prune.rs +++ b/crates/krishiv-optimizer/src/prune.rs @@ -154,6 +154,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 09efad1b3..8f7195645 100644 --- a/crates/krishiv-optimizer/tests/it/main.rs +++ b/crates/krishiv-optimizer/tests/it/main.rs @@ -15,6 +15,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..4b921cd96 --- /dev/null +++ b/crates/krishiv-optimizer/tests/it/python_relations.rs @@ -0,0 +1,132 @@ +// 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, + } +} + +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}" + ); + } +} From 457bf7ac1a9e63629d220ef96b176bec9934091c Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Fri, 13 Nov 2026 12:08:21 +0000 Subject: [PATCH 03/29] =?UTF-8?q?exec:=20Python=20UDF=20relation=20seam=20?= =?UTF-8?q?and=20its=20sink=20(P3-T014a,=20T085=20=C2=A73.2)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The relation kinds get a second runtime seam beside PreparedUdf: RelationCallSpec, PreparedRelationUdf::open, a UdfTask fed by Frames (Batch with a side, EndGroup, End) and read through next. A task's send and next may be awaited at once, so a runtime can wait for its results to be taken. UdfRuntime::prepare_relation defaults to KRISHIV_UDF_NOT_AVAILABLE, so the existing runtimes keep compiling. RelationSink runs map, grouped, cogrouped and aggregate shapes in PythonEval's exchange shape: one task per driver opened at its first frame; rows held until a frame is full (maxRecordsPerBatch, and maxBytesPerBatch when set) or the group or side ends; groups cut where the sorted key changes; Blocked while a lane's window of frames is full; results stored in the shared plan's spillable list. Held rows and queued frames are reserved; a refused batch changes nothing; cancel and failure drop the tasks. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-exec/src/lib.rs | 1 + crates/krishiv-exec/src/python_eval.rs | 20 +- crates/krishiv-exec/src/python_relation.rs | 144 +++ .../krishiv-exec/src/python_relation/sink.rs | 955 ++++++++++++++++++ .../krishiv-exec/src/python_relation/tests.rs | 561 ++++++++++ crates/krishiv-exec/src/shared.rs | 8 +- 6 files changed, 1686 insertions(+), 3 deletions(-) create mode 100644 crates/krishiv-exec/src/python_relation.rs create mode 100644 crates/krishiv-exec/src/python_relation/sink.rs create mode 100644 crates/krishiv-exec/src/python_relation/tests.rs diff --git a/crates/krishiv-exec/src/lib.rs b/crates/krishiv-exec/src/lib.rs index 739f10480..a84c06023 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/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..26addf2eb --- /dev/null +++ b/crates/krishiv-exec/src/python_relation.rs @@ -0,0 +1,144 @@ +// 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 partition (one per driver), sends it the partition'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. +//! +//! [`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. + +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; + +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 hold + /// every input column, its keys among them. + pub input_types: Vec>, + /// The grouping keys (0 for a map kind): with [`Self::arity`] what the + /// Python side needs to call `f(pdf)` or `f(key, pdf)`. + pub key_count: usize, + /// The columns of a side's frames (the first side's for a cogroup). + 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`: most bytes per frame + /// (cut between rows, 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). + 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 partition: 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. +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/sink.rs b/crates/krishiv-exec/src/python_relation/sink.rs new file mode 100644 index 000000000..da19bf1b7 --- /dev/null +++ b/crates/krishiv-exec/src/python_relation/sink.rs @@ -0,0 +1,955 @@ +// 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`] 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. +//! +//! **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. +//! Rows are held until a frame is full (`maxRecordsPerBatch`, and +//! `maxBytesPerBatch` when set) or its group or side ends, as Spark cuts a +//! group's Arrow stream; a group is never held whole. +//! +//! **Memory.** Held rows and queued frames are reserved in the driver's +//! consumer (a refused batch returns `NeedMemory` before anything changes, +//! so the retry is exact); a frame's reservation is released once the task +//! took it. Results wait in the lane's queue until the next `sink` call +//! stores them; an aggregate keeps one key row per group until `finalize`. +//! +//! **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. + +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; +use krishiv_memory::{Reservation, SpillScope}; +use krishiv_types::batch::Batch; + +use super::{Frame, PreparedRelationUdf, UdfTask}; +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 group's keys then one column per lane. + Aggregate { + /// The key columns. + keys: 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 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("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, + lanes: Vec, + 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, + /// Most bytes per frame (cut between rows); 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 } => { + let 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::>()?; + Some(Arc::new(Schema::new(fields))) + } + _ => None, + }; + Ok(Arc::new(Self { + shape, + lanes, + 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) + } +} + +/// A lane's queues, shared by the driver and the lane's pump. +#[derive(Default)] +struct LaneIo { + /// Frames not taken by the task yet, each with its reservation. + inbound: VecDeque<(Frame, Reservation)>, + /// The reservation of the frame the task is taking. + sending: Option, + /// Batches the task yielded, not stored yet. + outputs: VecDeque, +} + +impl LaneIo { + fn in_flight(&self) -> usize { + self.inbound.len() + usize::from(self.sending.is_some()) + } +} + +/// 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)] +struct Lane { + io: Arc>, + pump: Arc>, +} + +impl Lane { + fn new() -> Self { + Self { + io: Arc::new(Mutex::new(LaneIo::default())), + pump: Arc::new(Mutex::new(Pump::Idle)), + } + } + + /// Polls the pump; its outcome once it ended (`None` while it runs or + /// before the task opened). + fn poll(pump: &Mutex, cx: &mut Context<'_>) -> Option> { + let mut p = lock(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, + } + } + + fn opened(&self) -> bool { + !matches!(*lock(&self.pump), Pump::Idle) + } + + /// Queues `frame`, opening the task first if it is not. + fn queue(&self, udf: &dyn PreparedRelationUdf, frame: Frame, r: Reservation) -> Result<()> { + { + let mut p = lock(&self.pump); + if matches!(*p, Pump::Idle) { + let task = udf.open()?; + *p = Pump::Running(pump(task, Arc::clone(&self.io))); + } + } + lock(&self.io).inbound.push_back((frame, r)); + Ok(()) + } +} + +/// Sends the queued frames to `task` while it reads the task's results. +fn pump(task: Box, io: Arc>) -> BoxFuture<'static, Result<()>> { + Box::pin(async move { + let task = task; + 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, r)) => { + q.sending = Some(r); + Poll::Ready(frame) + } + None => Poll::Pending, + } + }) + .await; + let end = matches!(frame, Frame::End); + task.send(frame).await?; + lock(&io).sending = None; + if end { + return Ok::<(), KrishivError>(()); + } + } + }; + let receive = async { + while let Some(batch) = task.next().await? { + lock(&io).outputs.push_back(batch); + } + Ok(()) + }; + futures::future::try_join(send, receive).await.map(drop) + }) +} + +/// One driver's state. +struct RelLocal { + store: LocalState, + lanes: Vec, + /// A finished batch the store could not take yet (it asked for memory). + ready: Option, + 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 sent yet. + held: Vec, + held_rows: usize, + held_side: u8, + /// The reservation of the held rows' bytes. + held_bytes: Reservation, + /// An aggregate's group keys, one row per group, in group order. + key_rows: Vec, + key_bytes: Vec, +} + +/// What a driver left at `combine`. +struct Leftover { + lanes: Vec, + ready: Option, + key_rows: Vec, + _key_bytes: Vec, +} + +/// 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. + fn frame( + &self, + lane: &RelationLane, + side: u8, + rows: &RecordBatch, + partition: u32, + ) -> Result { + 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()); + for (e, name) in exprs.iter().zip(names) { + let a = e.evaluate_in(rows, partition)?; + let client = krishiv_types::arrow::client_arrow_type(a.data_type()); + let a = crate::carrier::recarried_array(&a, &client)?; + fields.push(Field::new(name, client, 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 }) + } + + /// Sends the held rows in frames of at most `maxRecordsPerBatch` rows + /// (and `maxBytesPerBatch` bytes); with `all` the rows left over too + /// (the end of a group or side). + fn send_held(&self, l: &mut RelLocal, all: bool) -> Result<()> { + if l.held_rows == 0 { + return Ok(()); + } + let held = std::mem::take(&mut l.held); + let rows = std::mem::take(&mut l.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 mut step = if self.max_rows == 0 { + rows + } else { + self.max_rows + }; + if self.max_bytes > 0 { + let bytes = data.get_array_memory_size().max(1); + let fit = (self.max_bytes.saturating_mul(rows) / bytes).max(1); + step = step.min(fit); + } + let mut offset = 0; + while offset < rows { + let len = step.min(rows - offset); + if len < step && !all { + // Too few for a full frame: held for the next input. + l.held.push(data.slice(offset, len)); + l.held_rows = len; + break; + } + let chunk = data.slice(offset, len); + let share = l.held_bytes.size() * len / (rows - offset); + let mut part = Some(l.held_bytes.split(share)); + for (lane, state) in self.lanes.iter().zip(&l.lanes) { + let frame = self.frame(lane, l.held_side, &chunk, l.partition)?; + // The chunk's bytes go with the first lane's frame. + let r = match part.take() { + Some(r) => r, + None => l.held_bytes.split(0), + }; + state.queue(lane.udf.as_ref(), frame, r)?; + } + offset += len; + } + Ok(()) + } + + /// Ends the open group: its held rows, then `EndGroup` to every lane. + fn end_group(&self, l: &mut RelLocal) -> Result<()> { + self.send_held(l, true)?; + for (lane, state) in self.lanes.iter().zip(&l.lanes) { + let r = l.held_bytes.split(0); + state.queue(lane.udf.as_ref(), Frame::EndGroup, r)?; + } + Ok(()) + } + + /// Feeds `batch` (reserved in `held_bytes`) to the lanes. + fn feed( + &self, + l: &mut RelLocal, + batch: &RecordBatch, + runs: Vec, + keys: Vec, + ) -> Result<()> { + let grouped = self.encoder.is_some(); + let mut keys = keys.into_iter(); + for run in runs { + if run.new_group { + if l.group_open { + self.end_group(l)?; + } + l.group_open = true; + if let Some(k) = keys.next() { + l.key_rows.push(k); + } + } else if grouped && !l.group_open { + return Err(internal("Python UDF rows before their group".to_owned())); + } + if run.side != l.held_side { + self.send_held(l, true)?; + l.held_side = run.side; + } + l.held.push(batch.slice(run.start, run.len)); + l.held_rows += run.len; + self.send_held(l, false)?; + } + Ok(()) + } + + /// The copied key row of each run starting a group (an aggregate's). + fn key_rows(&self, batch: &RecordBatch, runs: &[Run]) -> Result> { + let (RelationShape::Aggregate { keys }, Some(schema)) = (&self.shape, &self.key_schema) + else { + return Ok(Vec::new()); + }; + let projected = batch + .project(keys) + .map_err(|e| internal(format!("aggregate keys: {e}")))?; + let projected = RecordBatch::try_new_with_options( + Arc::clone(schema), + projected.columns().to_vec(), + &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())), + ) + .map_err(|e| internal(format!("aggregate keys: {e}")))?; + runs.iter() + .filter(|r| r.new_group) + .map(|r| { + arrow_select::concat::concat_batches(schema, [&projected.slice(r.start, 1)]) + .map_err(|e| internal(format!("aggregate keys: {e}"))) + }) + .collect() + } + + /// 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(internal(format!( + "a Python UDF returned {} 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(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| internal(format!("a Python UDF result does not fit its schema: {e}")))?; + Ok(Batch::new(out)) + } + + /// An aggregate driver's rows: each group's keys, then each lane's value + /// of it. + fn aggregate_rows(&self, key_rows: &[RecordBatch], lanes: &[Lane]) -> Result> { + let Some(schema) = &self.key_schema else { + return Ok(None); + }; + if key_rows.is_empty() { + return Ok(None); + } + let groups = key_rows.iter().map(RecordBatch::num_rows).sum::(); + let keys = arrow_select::concat::concat_batches(schema, key_rows) + .map_err(|e| internal(format!("aggregate keys: {e}")))?; + let mut columns: Vec = keys.columns().to_vec(); + let width = columns.len(); + for (i, (lane, state)) in self.lanes.iter().zip(lanes).enumerate() { + let outputs: Vec = lock(&state.io).outputs.drain(..).collect(); + let values: Vec<&ArrayRef> = outputs + .iter() + .map(|b| { + if b.num_columns() == 1 { + Ok(b.column(0)) + } else { + Err(internal(format!( + "a Python aggregate UDF returned {} columns", + b.num_columns() + ))) + } + }) + .collect::>()?; + let refs: Vec<&dyn Array> = values.iter().map(|a| AsRef::as_ref(*a)).collect(); + let client = lane + .client + .first() + .ok_or_else(|| internal("an aggregate lane without a type".to_owned()))?; + let column = if refs.is_empty() { + arrow_array::new_empty_array(client) + } else { + arrow_select::concat::concat(&refs) + .map_err(|e| internal(format!("aggregate results: {e}")))? + }; + if column.len() != groups { + return Err(internal(format!( + "a Python aggregate UDF returned {} rows for {groups} groups", + column.len() + ))); + } + let field = self + .schema + .fields() + .get(width + i) + .ok_or_else(|| internal("an aggregate result out of its schema".to_owned()))?; + columns.push(recarried(&column, client, field.data_type())?); + } + let options = RecordBatchOptions::new().with_row_count(Some(groups)); + let out = RecordBatch::try_new_with_options(Arc::clone(&self.schema), columns, &options) + .map_err(|e| internal(format!("a Python aggregate result: {e}")))?; + Ok(Some(Batch::new(out))) + } + + /// Stores `batch` in the list, or keeps it for the retry after the + /// memory the store asked for. + fn store(&self, l: &mut RelLocal, batch: Batch, g: &RelGlobal) -> Result> { + match self.store.sink(batch.clone(), &mut l.store, &g.store)? { + SinkStatus::NeedMemory { bytes } => { + l.ready = Some(batch); + Ok(Some(SinkStatus::NeedMemory { bytes })) + } + _ => Ok(None), + } + } + + /// Moves each lane's frames and results without waiting; the first + /// failure of a lane. + fn progress(l: &RelLocal) -> Result<()> { + let mut cx = Context::from_waker(noop_waker_ref()); + for lane in &l.lanes { + if let Some(Err(e)) = Lane::poll(&lane.pump, &mut cx) { + return Err(e); + } + } + Ok(()) + } +} + +/// `c` checked to be `client`, moved to the canonical `canonical`. +fn recarried(c: &ArrayRef, client: &ArrowType, canonical: &ArrowType) -> Result { + if c.data_type() != client { + return Err(internal(format!( + "a Python UDF returned {} where {client} was declared", + c.data_type() + ))); + } + crate::carrier::recarried_array(c, canonical) +} + +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 { + store: self.store.local_state(ctx, &g.store)?, + lanes: self.lanes.iter().map(|_| Lane::new()).collect(), + ready: None, + partition: ctx.partition(), + current: None, + group_open: false, + held: Vec::new(), + held_rows: 0, + held_side: 0, + held_bytes: consumer.reservation(), + key_rows: Vec::new(), + key_bytes: Vec::new(), + })) + } + + fn sink( + &self, + input: Batch, + local: &mut LocalState, + global: &GlobalState, + ) -> Result { + let g = global_ref::(global, "PythonRelation")?; + let l = local_mut::(local, "PythonRelation")?; + // Store what finished first: a `NeedMemory` here leaves the input + // unsent, so the retry with the same batch is exact. + if let Some(ready) = l.ready.take() + && let Some(status) = self.store(l, ready, g)? + { + return Ok(status); + } + Self::progress(l)?; + if !matches!(self.shape, RelationShape::Aggregate { .. }) { + loop { + let next = l.lanes.iter().zip(&self.lanes).find_map(|(state, lane)| { + lock(&state.io).outputs.pop_front().map(|b| (lane, b)) + }); + let Some((lane, batch)) = next else { break }; + let batch = self.output(lane, &batch)?; + if let Some(status) = self.store(l, batch, g)? { + return Ok(status); + } + } + } + if input.num_rows() > 0 { + let input = crate::operators::owned_rows(input)?; + let data = input.data(); + let (runs, last) = self.runs(data, l.current.as_ref())?; + let keys = self.key_rows(data, &runs)?; + let key_bytes: usize = keys.iter().map(RecordBatch::get_array_memory_size).sum(); + let bytes = data.get_array_memory_size(); + if l.held_bytes.try_grow(bytes + key_bytes).is_err() { + return Ok(SinkStatus::NeedMemory { + bytes: bytes + key_bytes, + }); + } + let part = l.held_bytes.split(key_bytes); + l.key_bytes.push(part); + self.feed(l, data, runs, keys)?; + l.current = last; + Self::progress(l)?; + } + let full = l + .lanes + .iter() + .any(|lane| lock(&lane.io).in_flight() >= self.window); + if full { + let lanes: Vec = l.lanes.clone(); + let window = self.window; + return Ok(SinkStatus::Blocked(Box::pin(poll_fn(move |cx| { + let mut room = true; + for lane in &lanes { + // A lane that ended (or failed) frees the driver: the + // next `sink` call returns its failure. + if Lane::poll(&lane.pump, cx).is_some() { + return Poll::Ready(()); + } + room &= lock(&lane.io).in_flight() < window; + } + if room { Poll::Ready(()) } else { Poll::Pending } + })))); + } + 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 mut l = *l; + // The input ended: the open group ends, then the partition. + if l.group_open { + self.end_group(&mut l)?; + } else { + self.send_held(&mut l, true)?; + } + for (lane, state) in self.lanes.iter().zip(&l.lanes) { + if state.opened() { + let r = l.held_bytes.split(0); + state.queue(lane.udf.as_ref(), Frame::End, r)?; + } + } + RelationSink::progress(&l)?; + let RelLocal { + store, + lanes, + ready, + key_rows, + key_bytes, + .. + } = l; + lock(&g.leftover).push(Leftover { + lanes, + ready, + key_rows, + _key_bytes: key_bytes, + }); + 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 leftover = std::mem::take(&mut *lock(&g.leftover)); + // Every task at once: the first failure (or a cancel) ends the + // wait, and the error drops the other tasks, which cancels them. + 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())); + } + let mut all = true; + for lane in leftover.iter().flat_map(|d| &d.lanes) { + match Lane::poll(&lane.pump, cx) { + Some(Err(e)) => return Poll::Ready(Err(e)), + Some(Ok(())) => {} + None => all &= !lane.opened(), + } + } + if all { + Poll::Ready(Ok(())) + } else { + Poll::Pending + } + }) + .await?; + for driver in leftover { + ctx.query().check_cancelled()?; + let mut batches: Vec = driver.ready.into_iter().collect(); + if matches!(self.shape, RelationShape::Aggregate { .. }) { + batches.extend(self.aggregate_rows(&driver.key_rows, &driver.lanes)?); + } else { + for (lane, state) in self.lanes.iter().zip(&driver.lanes) { + let outputs: Vec = lock(&state.io).outputs.drain(..).collect(); + for b in &outputs { + batches.push(self.output(lane, b)?); + } + } + } + if batches.is_empty() { + continue; + } + 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)?; + self.store.add_part(&g.store, batches, reservation)?; + } + 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 + }) + } +} + +/// 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..00a516801 --- /dev/null +++ b/crates/krishiv-exec/src/python_relation/tests.rs @@ -0,0 +1,561 @@ +// 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::{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, + 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, +} + +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, +} + +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(), + })) + } +} + +#[derive(Default)] +struct TaskState { + out: VecDeque, + rows: i64, + ended: bool, +} + +struct Task { + shared: Arc, + yields: Yield, + state: std::sync::Mutex, + ready: AtomicWaker, +} + +impl Drop for Task { + fn drop(&mut self) { + self.shared.dropped.fetch_add(1, Ordering::SeqCst); + } +} + +fn count(n: i64) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![Field::new("n", ArrowType::Int64, true)])); + RecordBatch::try_new( + schema, + vec![Arc::new(Int64Array::from(vec![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.shared.gate.register(cx.waker()); + if 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(); + match frame { + Frame::Batch { batch, .. } => { + s.rows += i64::try_from(batch.num_rows()).unwrap(); + if matches!(self.yields, Yield::Twice) { + s.out.push_back(batch.clone()); + s.out.push_back(batch); + } + } + Frame::EndGroup => { + if matches!(self.yields, Yield::Count) { + let n = std::mem::take(&mut s.rows); + s.out.push_back(count(n)); + } + } + Frame::End => s.ended = true, + } + drop(s); + self.ready.wake(); + Ok(()) + }) + } + + fn next(&self) -> BoxFuture<'_, crate::error::Result>> { + 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 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 { + 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 => ( + 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 => ( + vec![ArrowType::Int32, ArrowType::Int64], + vec![ + Field::new("k", ArrowType::Int32, true), + Field::new("v", ArrowType::Int64, true), + ], + ), + }; + let lane = RelationLane { + 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, + None, + ) + .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 { + loop { + match r.sink.sink(b.clone(), &mut l, &g)? { + SinkStatus::Blocked(wait) => wait.await, + SinkStatus::NeedMemory { bytes } => panic!("no memory for {bytes}"), + _ => break, + } + } + } + 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_only_that_side_then_end_group() { + 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", "G", "B0:1", "B1:2", "G", "B1:1", "G", "E"] + ); + done(r); +} + +#[tokio::test] +async fn an_aggregate_lane_gets_one_group_per_key() { + let r = rig( + RelationShape::Aggregate { keys: 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); +} + +#[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]], + 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}"); +} diff --git a/crates/krishiv-exec/src/shared.rs b/crates/krishiv-exec/src/shared.rs index 3f3f42481..baf0ccc24 100644 --- a/crates/krishiv-exec/src/shared.rs +++ b/crates/krishiv-exec/src/shared.rs @@ -623,8 +623,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(|| { From 6d2289d4846b87ef623ca0cbeafb537b9514bd20 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Fri, 13 Nov 2026 14:01:35 +0000 Subject: [PATCH 04/29] =?UTF-8?q?exec,=20plan-codec:=20Python=20relation?= =?UTF-8?q?=20nodes,=20planning=20and=20codec=20feature=20(P3-T014a,=20T08?= =?UTF-8?q?5=20=C2=A73.1)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Physical nodes PythonMap, PythonGroupMap, PythonCoGroupMap and PythonAggregate (pure data: the UDF by name and version, frame columns, declared result) run on RelationSink. The planner gives a grouped kind a hash exchange on its keys (one partition without keys) and the spillable sort; a cogroup's two sides go through one union stream sorted by key then side, so a key's left rows reach the UDF before its right rows. An aggregate whose aggregates are Python UDAFs plans onto PythonAggregate; mixing them with built-ins is INVALID_PANDAS_UDF_PLACEMENT. Barrier mode, Python window aggregates and UDTFs are KRISHIV_UNSUPPORTED_FEATURE (PR-26c); a scalar PythonEval refuses a relation eval type. A worker without a runtime refuses with KRISHIV_UDF_NOT_AVAILABLE. Plan codec feature python_udf_relations (semantic): the four nodes, GENERATOR_KIND_PYTHON_UDTF and every eval type past the scalar ones. Proto messages, encode and decode, inventory cases, goldens, and a reader without the feature refuses each plan naming it. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-exec/CLAUDE.md | 1 + crates/krishiv-exec/src/physical/data.rs | 1 + .../krishiv-exec/src/physical/instantiate.rs | 5 + .../physical/instantiate/python_relation.rs | 280 ++++++++ crates/krishiv-exec/src/physical/mod.rs | 95 +++ crates/krishiv-exec/src/planner/build.rs | 33 + .../krishiv-exec/src/planner/build_project.rs | 28 + crates/krishiv-exec/src/planner/mod.rs | 1 + .../src/planner/python_relation.rs | 520 +++++++++++++++ crates/krishiv-exec/src/planner/python_udf.rs | 10 + crates/krishiv-exec/tests/planner/main.rs | 1 + .../tests/planner/python_relations.rs | 630 ++++++++++++++++++ crates/krishiv-plan-codec/src/decode.rs | 70 ++ crates/krishiv-plan-codec/src/encode.rs | 90 ++- crates/krishiv-plan-codec/src/features.rs | 36 +- .../src/proto/krishiv.plan.v1.pb.rs | 99 ++- .../tests/common/inventory.rs | 167 +++++ crates/krishiv-plan-codec/tests/features.rs | 52 ++ .../python_udf_relations/python_aggregate.bin | Bin 0 -> 902 bytes .../python_cogroup_map.bin | Bin 0 -> 938 bytes .../python_eval_types.bin | Bin 0 -> 2000 bytes .../python_udf_relations/python_group_map.bin | Bin 0 -> 914 bytes .../python_udf_relations/python_map.bin | Bin 0 -> 910 bytes .../python_udtf_generate.bin | Bin 0 -> 852 bytes crates/krishiv-plan-codec/tests/roundtrip.rs | 14 + crates/krishiv-session/src/catalog.rs | 4 + proto/krishiv/plan/v1/plan.proto | 69 ++ 27 files changed, 2197 insertions(+), 9 deletions(-) create mode 100644 crates/krishiv-exec/src/physical/instantiate/python_relation.rs create mode 100644 crates/krishiv-exec/src/planner/python_relation.rs create mode 100644 crates/krishiv-exec/tests/planner/python_relations.rs create mode 100644 crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_aggregate.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_cogroup_map.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_eval_types.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_group_map.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_map.bin create mode 100644 crates/krishiv-plan-codec/tests/golden/python_udf_relations/python_udtf_generate.bin diff --git a/crates/krishiv-exec/CLAUDE.md b/crates/krishiv-exec/CLAUDE.md index 10e674160..1ab56e2d8 100644 --- a/crates/krishiv-exec/CLAUDE.md +++ b/crates/krishiv-exec/CLAUDE.md @@ -126,6 +126,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, `Blocked` while a lane's frame window is full, results in the shared plan's list); `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 (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/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 cbb7f9136..f978f0cb3 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; @@ -923,6 +924,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..41c81bafa --- /dev/null +++ b/crates/krishiv-exec/src/physical/instantiate/python_relation.rs @@ -0,0 +1,280 @@ +// 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. +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, .. } => { + (input, RelationShape::Aggregate { keys: keys.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, + keys, + calls, + .. + } => self.aggregate_lanes(&input.schema, layout, keys, 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()) + .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(), + key_count, + arity: plan.sides.first().map_or(0, |s| s.columns.len()), + 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 { + 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) 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 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], + 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 { + udf: prepared, + names: vec![(0..args.len()).map(|j| format!("_{j}")).collect()], + 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 ccc10544c..36473dc7c 100644 --- a/crates/krishiv-exec/src/physical/mod.rs +++ b/crates/krishiv-exec/src/physical/mod.rs @@ -421,6 +421,51 @@ 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 keys, then one value per call. + PythonAggregate { + /// The input. + input: Box, + /// The attribute of each input column. + layout: Vec, + /// The key columns of `input`. + keys: Vec, + /// The calls, in output column order after the keys. + calls: Vec, + }, /// Each row of `input` `count` times (`INTERSECT ALL`, `EXCEPT ALL`). ReplicateRows { /// The input. @@ -1023,6 +1068,10 @@ plan_data!(enum PhysNode { Limit { input, offset, limit }, ReplicateRows { input, layout, count, subqueries }, PythonEval { input, layout, calls }, + PythonMap { input, udf }, + PythonGroupMap { input, keys, udf }, + PythonCoGroupMap { input, keys, side, udf }, + PythonAggregate { input, layout, keys, calls }, Generate { input, layout, kind, outer, arg, subqueries }, Expand { input, layout, projections, nullable, subqueries }, Sort { input, keys }, @@ -1043,6 +1092,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, @@ -1127,6 +1179,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()], @@ -1173,6 +1229,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()], @@ -1197,6 +1257,41 @@ 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. +#[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 40f78f056..17241f200 100644 --- a/crates/krishiv-exec/src/planner/build.rs +++ b/crates/krishiv-exec/src/planner/build.rs @@ -562,6 +562,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) } @@ -579,6 +584,11 @@ impl<'a> Planner<'a> { partition_by, order_by, } => { + if exprs.iter().any(|e| super::python_udf::has_python_udf(e)) { + return Err(super::python_relation::later_kind( + "a Python UDF in a window", + )); + } let input = self.plan(input)?; self.window(input, exprs, partition_by, order_by, None) } @@ -599,6 +609,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 99e5042f2..23842a090 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..9c17f1539 --- /dev/null +++ b/crates/krishiv-exec/src/planner/python_relation.rs @@ -0,0 +1,520 @@ +// 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::>(), + )) +} + +/// The frame side reading the first `columns.len()` columns of a relation. +fn side_of(columns: &[usize], cols: &[BoundColumn]) -> PythonFrameSide { + PythonFrameSide { + columns: columns.to_vec(), + names: cols.iter().map(|c| c.name.clone()).collect(), + types: cols.iter().map(|c| c.data_type.clone()).collect(), + } +} + +/// 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`. + 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, key_columns) = self.with_columns(planned, keys)?; + let columns = cols + .iter() + .map(|c| planned.column(c.id)) + .collect::>>()?; + let sorted = self.clustered_and_sorted(planned, &key_columns, &[])?; + let keys = key_columns.clone(); + let plan = PythonRelationPlan { + udf: udf.clone(), + sides: vec![side_of(&columns, &cols)], + output: outputs(output), + }; + let explain = format!( + "[{}], {}#{}", + key_names(&sorted, &keys), + udf.name, + udf.version + ); + Ok(relation( + sorted, + |input| PhysNode::PythonGroupMap { + input, + keys, + udf: plan, + }, + output, + ("PythonGroupMap", explain), + )) + } + + /// `cogroup(...).applyInPandas` / `applyInArrow`. + 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, lk) = self.with_columns(l, left_keys)?; + let r = self.plan(right)?; + let (r, rk) = self.with_columns(r, right_keys)?; + let k = lk.len(); + // One layout for both branches: keys, side, left, right. + let width = k + 1 + left_cols.len() + right_cols.len(); + let layout: Vec = (0..width).map(|_| self.fresh_id()).collect(); + let nullable = vec![true; width]; + let mut branches = Vec::with_capacity(2); + for (side, (planned, keys)) in [(l, &lk), (r, &rk)].into_iter().enumerate() { + let mut items = Vec::with_capacity(width); + 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()))), + }); + } + } + } + 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 sorted = self.clustered_and_sorted(union, &key_columns, &[k])?; + let plan = PythonRelationPlan { + udf: udf.clone(), + sides: vec![ + side_of( + &(k + 1..k + 1 + left_cols.len()).collect::>(), + &left_cols, + ), + side_of( + &(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: key_columns, + 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); + } + let planned = self.plan(input)?; + let (planned, key_columns) = self.with_columns(planned, keys)?; + let prepared = self.prepare(&[])?; + for call in &calls { + let ExprNode::Resolved(ResolvedExpr::Function(f)) = &call.node else { + return Err(internal("a Python UDAF call is not a function")); + }; + if f.args.iter().any(|a| super::build::has_subquery(a)) { + return Err(unsupported("a subquery in the arguments of a Python UDF")); + } + let args: Vec<&Expr> = f.args.iter().map(AsRef::as_ref).collect(); + compile_all( + &args, + &planned.layout, + &planned.schema, + &self.values(&prepared), + )?; + } + let sorted = self.clustered_and_sorted(planned, &key_columns, &[])?; + 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) in calls.iter().enumerate() { + let ExprNode::Resolved(ResolvedExpr::Function(f)) = &call.node else { + return Err(internal("a Python UDAF call is not a function")); + }; + 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(call), + output: column.id, + }); + columns.push((Arc::clone(call), column)); + } + let detail = calls + .iter() + .zip(&columns) + .map(|(c, (_, col))| format!("{c} AS {}", col.name)) + .collect::>() + .join(", "); + let input_layout = sorted.layout.clone(); + let keys_at = key_columns.clone(); + let schema = Arc::new(Schema::new(fields)); + let node = sorted.wrap( + |input| PhysNode::PythonAggregate { + input, + layout: input_layout, + keys: keys_at, + 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::>>()?; + let mut out = self.project( + Planned { + placement: Placement::Unknown, + ordering: Vec::new(), + sites: Vec::new(), + ..node + }, + &exprs, + )?; + out.sites.clear(); + Ok(out) + } + + /// `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))) +} + +/// 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..8a8a72a53 100644 --- a/crates/krishiv-exec/src/planner/python_udf.rs +++ b/crates/krishiv-exec/src/planner/python_udf.rs @@ -179,6 +179,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/tests/planner/main.rs b/crates/krishiv-exec/tests/planner/main.rs index 4d11e5030..668715dbf 100644 --- a/crates/krishiv-exec/tests/planner/main.rs +++ b/crates/krishiv-exec/tests/planner/main.rs @@ -29,6 +29,7 @@ mod orientation; mod partial_aggregate; mod physical; mod plan_stats; +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..f63ac4590 --- /dev/null +++ b/crates/krishiv-exec/tests/planner/python_relations.rs @@ -0,0 +1,630 @@ +// 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, + /// 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>>, +} + +type Log = Arc>>; + +impl UdfRuntime for Runtime { + fn prepare(&self, spec: &UdfCallSpec) -> Result, KrishivError> { + Err(krishiv_common::error::InternalError { + message: format!("no scalar {}", spec.name), + } + .into()) + } + + fn prepare_relation( + &self, + spec: &RelationCallSpec, + ) -> Result, KrishivError> { + let mode = match spec.name.as_str() { + "echo" => Mode::Echo, + "twice" => Mode::Twice, + _ => Mode::Count, + }; + Ok(Arc::new(Prepared { + mode, + output: Arc::clone(&spec.output_arrow), + log: Arc::clone(&self.log), + per_task: Arc::clone(&self.per_task), + })) + } +} + +/// The fake's prepared site. +struct Prepared { + mode: Mode, + output: arrow_schema::SchemaRef, + log: Arc>>, + per_task: Arc>>, +} + +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), + 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>>, + state: std::sync::Mutex, + ready: AtomicWaker, +} + +impl Task { + 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 { + let mut s = self.state.lock().unwrap(); + match frame { + Frame::Batch { 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 => { + let back = back.unwrap(); + s.out.push_back(back.clone()); + s.out.push_back(back); + } + 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, + } +} + +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:?}"); +} + +#[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); + assert_eq!( + log.iter().filter(|e| e.starts_with("B1")).count(), + 2, + "{log:?}" + ); + // Within each task, a key's left rows come before its right rows. + for own in rt.per_task.lock().unwrap().iter() { + let own = own.lock().unwrap(); + for w in own.windows(2) { + assert!( + !(w[0].starts_with("B1") && w[1].starts_with("B0")), + "{own:?}" + ); + } + } +} + +/// A grouped-aggregate Python UDF call over `arg`. +fn udaf(arg: &BoundColumn) -> Arc { + Arc::new(Expr::new(ExprNode::Resolved(ResolvedExpr::Function( + ResolvedFunction { + id: udf("count", 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:?}"); +} + +#[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, + }); + let err = run_with(&env, node, Some(Arc::new(Runtime::default()))).unwrap_err(); + assert_eq!( + err.error_class(), + ErrorClass::InvalidPandasUdfPlacement, + "{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_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}" + ); +} + +#[test] +fn a_skewed_group_runs_through_the_spilling_sort_in_bounded_memory() { + use krishiv_exec::planner::{MemoryTable, MemoryTables}; + use krishiv_memory::{MemoryConfig, SpillConfig, SpillManager, TreePool}; + use krishiv_types::batch::Batch; + + const ROWS: i32 = 300_000; + let schema = Arc::new(Schema::new(vec![ + Field::new("k", ArrowType::Int32, false), + Field::new("v", ArrowType::Int64, false), + ])); + // 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 batches: Vec = (0..ROWS) + .step_by(8192) + .map(|start| { + let end = (start + 8192).min(ROWS); + let k: Int32Array = (start..end) + .map(|i| if i < 270_000 { 0 } else { 1 + i / 100 }) + .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 node = LogicalPlan::new(PlanNode::GroupMap { + input: Arc::new(scan), + keys: vec![col(&k)], + udf: udf("count", PythonEvalType::GroupedMapArrow), + output: counted(), + }); + let plan = Plan::try_new(node).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 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)); + let files: u64 = query + .metrics() + .operators + .values() + .map(|o| o.metrics.spill_files) + .sum(); + assert!(files > 0, "the sort did not spill"); + drop(query); + assert_eq!(krishiv_memory::MemoryPool::reserved(&*root), 0); +} diff --git a/crates/krishiv-plan-codec/src/decode.rs b/crates/krishiv-plan-codec/src/decode.rs index 931282e73..7657dee35 100644 --- a/crates/krishiv-plan-codec/src/decode.rs +++ b/crates/krishiv-plan-codec/src/decode.rs @@ -552,6 +552,36 @@ 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)?, + 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), @@ -560,6 +590,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)?, @@ -1683,3 +1714,42 @@ 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, + }; + 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 bd720057f..a805ce48f 100644 --- a/crates/krishiv-plan-codec/src/encode.rs +++ b/crates/krishiv-plan-codec/src/encode.rs @@ -334,6 +334,49 @@ 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, + calls, + } => Kind::PythonAggregate(pb::PythonAggregateNode { + input: self.node_index(input)?, + layout: layout(l), + keys: indexes(keys), + calls: calls + .iter() + .map(|c| { + Ok(pb::PythonCall { + call: self.expr(&c.call)?, + output: c.output.0, + }) + }) + .collect::>()?, + }), PhysNode::Generate { input, layout: l, @@ -351,11 +394,7 @@ impl Encoder { krishiv_plan::GeneratorKind::VariantExplode => { pb::GeneratorKind::VariantExplode } - krishiv_plan::GeneratorKind::PythonUdtf => { - return Err(CodecError::Unsupported( - "a Python UDTF generator".to_owned(), - )); - } + krishiv_plan::GeneratorKind::PythonUdtf => pb::GeneratorKind::PythonUdtf, } .into(), outer: *outer, @@ -748,6 +787,47 @@ 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, + }), + 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 c4294b598..9e8dd7020 100644 --- a/crates/krishiv-plan-codec/src/features.rs +++ b/crates/krishiv-plan-codec/src/features.rs @@ -152,6 +152,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)."), @@ -237,6 +238,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 @@ -398,6 +401,10 @@ pub fn feature_of(construct: Construct<'_>) -> Option { PhysNode::SortMergeJoin(_) => Some(Feature::SortMergeJoin), PhysNode::WindowGroupLimit { .. } => Some(Feature::WindowGroupLimit), PhysNode::PythonEval { .. } => Some(Feature::PythonUdf), + 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). @@ -465,9 +472,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::Generate), + krishiv_plan::GeneratorKind::PythonUdtf => Some(Feature::PythonUdfRelations), krishiv_plan::GeneratorKind::Explode | krishiv_plan::GeneratorKind::PosExplode | krishiv_plan::GeneratorKind::Inline => None, @@ -688,7 +697,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); @@ -784,6 +799,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 } @@ -807,7 +829,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( @@ -834,7 +858,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(); @@ -861,6 +888,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 aaa288058..98a7543a6 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" + 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" )] pub kind: ::core::option::Option, } @@ -217,8 +217,101 @@ pub mod node { /// feature: python_udf #[prost(message, tag = "26")] PythonEval(super::PythonEvalNode), + /// feature: python_udf_relations + #[prost(message, tag = "27")] + PythonMap(super::PythonMapNode), + /// feature: python_udf_relations + #[prost(message, tag = "28")] + PythonGroupMap(super::PythonGroupMapNode), + /// feature: python_udf_relations + #[prost(message, tag = "29")] + PythonCoGroupMap(super::PythonCoGroupMapNode), + /// feature: python_udf_relations + #[prost(message, tag = "30")] + PythonAggregate(super::PythonAggregateNode), } } +/// 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. +#[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, +} +/// 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, +} +/// 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)] @@ -1656,6 +1749,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. @@ -1669,6 +1764,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. @@ -1679,6 +1775,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 05c5a73dc..3713894f7 100644 --- a/crates/krishiv-plan-codec/tests/common/inventory.rs +++ b/crates/krishiv-plan-codec/tests/common/inventory.rs @@ -1717,6 +1717,167 @@ 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, + } +} + +/// 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], + 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 @@ -1899,6 +2060,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 2e3870699..d673e968e 100644 --- a/crates/krishiv-plan-codec/tests/features.rs +++ b/crates/krishiv-plan-codec/tests/features.rs @@ -1955,3 +1955,55 @@ 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}" + ); + } +} 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 0000000000000000000000000000000000000000..74740a6121ff11e92145cb7d786bf1306743aa9b GIT binary patch literal 902 zcmbtSy^hmB5Z+n;ZRT>sw+Ub@gsfD^5=|;}1)Q!zhz3IyA?22{lUV1oBYUGG9Szda zQSbn~0FOX{#7po7S5Q$Sg;^&YI)OyNuIBU3%ziVo;}MSUhUXmLaN)^k_!M`Wg{n^` z$4aZBnxU6wQ?)EKf|fqor&XFsfo`Mg{8-B&I=VW^N_k7(M0-(JsqCGpI-6I^CcD?V z!W{t;RLF%BLD&&oLPJU?dLCN{Dn^6`oxCiwdL)no_w4}ob74v24v)HwF+rO#dy9+# zX#*0U&}Y|Jk9_8qPZ*FyBy@}4CKRko-kp{k8S{qT#?Yi3cEyFa*V@^|!|=v;=u4_1 z1U9sV>vsr+j}gT0_}@w~iNBDXlN^v3{l$8{HmS{h%7#(lDQpbQ+$ z?bF8(?$A@EDprk7cCXA(62<{`@}{hGxh#qRhQRi@;0N#s5?k`SgQOqtAv_=tWxI4; zkNaq;GBsukF3l!5&E!j(NjHKditmhj=q^^1IxA=te?SBQAA&;!(WdwZ;p>**{UVHU YvC{LZoGjDnc<}pd5-6QlWpe~S0Zlc1ZU6uP literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..ed5c33e9c5217a515cf3497a18fce6ca1418a937 GIT binary patch literal 938 zcmbtSPiqrF6rVS{*-WOX*-k62i-%=T66z%hTq$Lzc=sAyg#fF2Jt!NuPJMaC;TzrG&5P9 z3`f$+qMYz%rp9t!SPoW%+&?WdB^VEkt@4o-DR0>FL}_tdT;u+%Dl^eJla-p5bE9t6 zrg(#a02x!(NFZ)8D!`!9N)Edo2ifHu3~%IGtE$I11seJh^mFD3XO55Cgb;?C5dTsV z0;ugELX;u7LON)3r$UqwB0=COdIKdRZM1IGzKNq(?sW`{%0Zi5G`DJ=T=f0H_BlcJ)z4Icq ziScttci+)}mBS_f4RMax zM>zfIa=CQryYrYaPwKA|0hS`b?G~-26RxDL9Piitnl6T0u8i!m1xS zxzVN7*dm-2AbrZ{KHP(WFX+vk!FG3x!(FuCZGTKxyW8B8I&-uKzv-j4R>rDq=!U}_ V?aUT-TI%6E8@tJqHZQd~hF@B?kPrX> literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..23b0e1bca17854bcae164f6e4eea805fcc77d79a GIT binary patch literal 2000 zcmbu9%Wl&^6hJ+WlQ`3~nNCvbnm$Z;6+;1Kg}Q*MD}>k(SVc&!RXd5PYX^Hmn;jdZ zBC%k>2k-%uKVX5xf;XZY;19as3wVV)&QNG!F%ll#6VKz!WUj~7vm|G3s`!qo=jl~^ z1rH}3Cth7zahMZCEu43~rqc}=MvM}ytwpX!b(~9B?5{9d#aR}udLf;l<5=s&kxNVK zPVBX#ZsN_QUFMLEC~{0y&sxYV=qg1CR0`HLJA;v9V~i4<^~2DMYdThuN;Ct_`+A1* zZgaV)D2fh6Q8YnCK~Sqe5Fk-TloIsuQ3Pl}SP(EuodoJi5qgJG*W%3sJ)MI=WwEI5 zc29;_xlWD6Jh~^ z*RS{ceZFHp9yW=s^p_(CyU#)4mzt7;XOU;%!Co)V=inF5oG!A?Ej}-IT!=yDmA^7S zSX(jpli-(@mVuW*{ePdB19Q|0Si9^s8AZY1)@f3-QNuV}_gigdB2+75tsX^;!s%s> zAHacA)N`3d>R4?*WY)TqSbi&vV$Z!~=!o>6W{lf+`=6g|JZtLA%ZSCZDmV z8Y4K@=`F=x0B?Y2A-+Nu#^=%%_Kmmf=|<^~@z2pf$jb^Zgiw$bUJ7AIR(K_ZVOim| z5K6Mb8zGF!3U7r_mKEL!VN6zdFNASfVOt0jvcd-;9F!G43SmlC_#}iGS>cNi4#^5% zg>YC__$Gv-vch*E9G4Y-2;qdR@JJoJaVKSkEg{Uw3Xg?w=9kdGogQmP;ZoOaHY&gG N`8UI1eiSCl=qDXpAV2^B literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..6bdc79080b17a1b88fd8bc524233846ffd5da90e GIT binary patch literal 914 zcmbtS!EO^V5cN2ljWbQj)(v7=aaipMMR3Up1&LNY0YXTGo+_l&be&|?tvA|rR3Zd| z15&wh-~;#pK7j)gAHkVZ5Bx(IZwgWcH*CFG&y2_3_dKE*_|X~1cU-jOQ+$G()m)Xw z<0EBNUQDr->WNz97QxB@y^|ubRQh(E ziZ%pDP!Si+7^1e|5-M`q@qWL7p!$eVVKXyEmji(uc;JOF%tS-FbadEZj0u{Ic^8&3 zAg@OV$Oh~R>r&3`3dn|pm_WPu4YGoDD7#VD#xA@rug|dL>~+L>Ypag4^PT7|{D9XI zShU10`t#+u)H(gpcZfkskLI?uc%3_hn+2bDazmRhDX* z+)8w=PqeZ7>!?%DOeuU+7^^Z$-~k|9vJ!XAIitJ<<~^Azv@mM8b;!Vsn&DXEp*DY@;_S3M;8)Nz|OWjPW~-J*nbZ zW^8<_syLgPqSWcbP(ZNy+6dpgU;X^@`IDctlu`MzvZKvQ{W^;JA-1#16xPrag^&On za=~}uK8!rcZ}vvp{Vjxhl%d=Hq^|b2v7t=r!VUNxKIj--S)Gnub5sYt`O?k`GhU<< M*Ja9Pg{cnV7pwD*I{*Lx literal 0 HcmV?d00001 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 0000000000000000000000000000000000000000..c9d37f7efb6d584620ce51ce1fed4fe5e90156b2 GIT binary patch literal 910 zcmbtS!EO^V5Vbd(jWbQj)+u6HaaipMMR3Up1&LNY0YXTGo+_l&be&|?tv6abDp3XE zf>drC_yB%@PvC&WojYI91OE_aHU+7I8@Ar8XU1didmhmYgZSKocb;g;=kN?R?OavI z<0Iu%QBI+i>4{nt4!Dyc_$OtirGUV?DnD{Eg{CWywUIaF4e;kxnaS>{s`RW}Sbeun z#T$Z4u3}F#6CSsPC%HwZo$U7;z*QfBTWIFS=xQLq<6h`Te3*-dq;!1PVT=i!jQN+A zF^=9gB19XqYpjbo+7+UW5D5ZzdAHCC*1_yXU7JvNU0$Ez(An#Vi`G^hXBRv1n;-m4 zE};v-d^xUkfq(OzCtOPK81az&B~d{HaSOPSw0H!*Q&THeqIB9sbB}1|i<&RhET66{tl1g04OjUF))ntgZ)a);**LSmsq^IH6b6scA z|KBH4c*dX&{u4cMk{9cxnn?a?0MH!e)0vAoAN0T?GEib!2aJq@zJP zItm^D(Zd^1An`C)@D5?tIUPEnV^%ZXnb~i?o%IMW@S_Wk@3`p57x*0anw6^0W~a)i zvRYs#%jarc8U!N)bYE9lE(Q9Hsf$x1$LN{rEZ6d`yo2tlt};0|S9QLu)=mDPO+`Hc z5>&*67emw)TtY)mH-4Eo2r5B@2E9V-yq*Z;zymjg!$LUHrlXTSV@%Lu%)PRV0eK_R zfNa2Su>s}mseo)q8k5o~evhnReaiM)Y;40jcw57ebKDn~o&7e>E+0p|&mbh-BJ>XY z7?KbCU;LKs2k8ar5vlcGZMR#?XU`J_#_iV{pSXQmR(NZzu{~J+k2i50Tf>UIx47!r zqHJA@$w+s7v%mV@akJ3AgB7=3J6uEl|Gt-^8j~)%&+@sEQ0}&BF5%4t(MyZP(nJ80 zA-ZW*8A-EaR8i9JX8fUZELFTLw29AE6Bi3z)p_A&h|?a>0+_38b#%2cvYD>?1s)46$9OZYM)@l+J9}f$Q)e z#pC4ubb$V9GpqBG=I<+yA#}&!5x}<;Q?bZtK$Xd-xZ0RyrDyAGJ{|quTck2YrJED@ E1)iUHs{jB1 literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/roundtrip.rs b/crates/krishiv-plan-codec/tests/roundtrip.rs index b1dfbc008..d3443aff5 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, @@ -75,6 +79,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, @@ -609,6 +617,12 @@ fn readers_that_predate_features_still_refuse_what_they_cannot_read() { | "exchange_source" | "rust_udf" | "python_udf" + | "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-session/src/catalog.rs b/crates/krishiv-session/src/catalog.rs index d68b8b0d3..8b0a44c9b 100644 --- a/crates/krishiv-session/src/catalog.rs +++ b/crates/krishiv-session/src/catalog.rs @@ -929,6 +929,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 7f4623841..921cc9e72 100644 --- a/proto/krishiv/plan/v1/plan.proto +++ b/proto/krishiv/plan/v1/plan.proto @@ -172,9 +172,76 @@ message Node { ExchangeSourceNode exchange_source = 25; // feature: python_udf PythonEvalNode python_eval = 26; + // feature: python_udf_relations + PythonMapNode python_map = 27; + // feature: python_udf_relations + PythonGroupMapNode python_group_map = 28; + // feature: python_udf_relations + PythonCoGroupMapNode python_co_group_map = 29; + // feature: python_udf_relations + PythonAggregateNode python_aggregate = 30; } } +// 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. +message PythonAggregateNode { + uint32 input = 1; + repeated uint64 layout = 2; + repeated uint64 keys = 3; + repeated PythonCall calls = 4; +} + +// 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; +} + +// 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 { @@ -365,6 +432,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 From 5aa431ebeb6f662a274903be4eef31801f434af6 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Fri, 13 Nov 2026 14:56:13 +0000 Subject: [PATCH 05/29] plan, exec: more tests of the Python relation nodes (P3-T014a) The structural cache key of each relation node, rewrites of their children and expressions, frames cut by maxBytesPerBatch, and a task result of another width failing the phase. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../krishiv-exec/src/python_relation/tests.rs | 57 +++++++++++++++++ crates/krishiv-plan/src/cache.rs | 62 +++++++++++++++++++ .../src/tests/python_relations.rs | 28 +++++++++ 3 files changed, 147 insertions(+) diff --git a/crates/krishiv-exec/src/python_relation/tests.rs b/crates/krishiv-exec/src/python_relation/tests.rs index 00a516801..086cd342d 100644 --- a/crates/krishiv-exec/src/python_relation/tests.rs +++ b/crates/krishiv-exec/src/python_relation/tests.rs @@ -559,3 +559,60 @@ fn a_runtime_without_relation_kinds_refuses_them_as_not_available() { 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); +} + +#[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 { + 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!( + err.to_string().contains("columns where 2 were declared"), + "{err}" + ); + done(r); +} diff --git a/crates/krishiv-plan/src/cache.rs b/crates/krishiv-plan/src/cache.rs index 429fa5afc..936f57370 100644 --- a/crates/krishiv-plan/src/cache.rs +++ b/crates/krishiv-plan/src/cache.rs @@ -1720,6 +1720,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 { @@ -1745,6 +1797,16 @@ 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, + } + } + #[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/tests/python_relations.rs b/crates/krishiv-plan/src/tests/python_relations.rs index eebedf43d..fe93f5ac7 100644 --- a/crates/krishiv-plan/src/tests/python_relations.rs +++ b/crates/krishiv-plan/src/tests/python_relations.rs @@ -292,3 +292,31 @@ fn a_python_udaf_beside_a_builtin_aggregate_is_refused() { ); 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}"); + } +} From 5f88b74c0f540d0fc77b67ee48bd35ce5913220c Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Fri, 13 Nov 2026 17:21:56 +0000 Subject: [PATCH 06/29] exec: refuse DISTINCT, FILTER and WITHIN GROUP on a Python UDAF call (P3-T014a review W1) The Python aggregate feeds every row of a group to the UDF, so these clauses would be ignored silently. Refused at planning with KRISHIV_UNSUPPORTED_FEATURE until Spark's class is verified. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../src/planner/python_relation.rs | 23 +++++++++++ .../tests/planner/python_relations.rs | 38 +++++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/crates/krishiv-exec/src/planner/python_relation.rs b/crates/krishiv-exec/src/planner/python_relation.rs index 9c17f1539..02d6fe4a6 100644 --- a/crates/krishiv-exec/src/planner/python_relation.rs +++ b/crates/krishiv-exec/src/planner/python_relation.rs @@ -319,6 +319,7 @@ impl Planner<'_> { let ExprNode::Resolved(ResolvedExpr::Function(f)) = &call.node else { return Err(internal("a Python UDAF call is not a function")); }; + 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")); } @@ -478,6 +479,28 @@ 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 attribute `e` reads if it is a plain column. fn column_of(e: &Expr) -> Option { match &e.node { diff --git a/crates/krishiv-exec/tests/planner/python_relations.rs b/crates/krishiv-exec/tests/planner/python_relations.rs index f63ac4590..5cf4a84c5 100644 --- a/crates/krishiv-exec/tests/planner/python_relations.rs +++ b/crates/krishiv-exec/tests/planner/python_relations.rs @@ -436,6 +436,44 @@ fn a_python_aggregate_beside_a_builtin_is_refused() { ); } +/// 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(); From 180338008c4dc8debfe45d7d1c6b7f5d1400144a Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Fri, 13 Nov 2026 17:30:46 +0000 Subject: [PATCH 07/29] exec: frames lead with the grouping keys; a one-sided cogroup key gets an empty frame (P3-T014a review M1, M2) Spark's layout (PandasGroupUtils.resolveArgOffsets): a grouped or cogrouped kind's frames lead with its key columns, computed keys included, so a runtime can build `key` for f(key, pdf). RelationCallSpec carries each side's column names; arity counts the data columns only. A cogroup key without rows on a side sends one zero-row Batch for that side (left before the right rows, right before EndGroup), so the Python side builds Spark's empty DataFrame from the frame's schema. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../physical/instantiate/python_relation.rs | 10 +- crates/krishiv-exec/src/physical/mod.rs | 4 +- .../src/planner/python_relation.rs | 52 ++++++-- crates/krishiv-exec/src/python_relation.rs | 41 ++++-- .../krishiv-exec/src/python_relation/sink.rs | 33 ++++- .../krishiv-exec/src/python_relation/tests.rs | 7 +- .../tests/planner/python_relations.rs | 119 ++++++++++++++++-- 7 files changed, 226 insertions(+), 40 deletions(-) diff --git a/crates/krishiv-exec/src/physical/instantiate/python_relation.rs b/crates/krishiv-exec/src/physical/instantiate/python_relation.rs index 41c81bafa..b4927e570 100644 --- a/crates/krishiv-exec/src/physical/instantiate/python_relation.rs +++ b/crates/krishiv-exec/src/physical/instantiate/python_relation.rs @@ -182,8 +182,12 @@ impl Instantiation<'_> { 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()), + 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)), @@ -249,6 +253,7 @@ impl Instantiation<'_> { 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, @@ -256,6 +261,7 @@ impl Instantiation<'_> { 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()], @@ -266,7 +272,7 @@ impl Instantiation<'_> { fields.push(Field::new(&name, canonical, true)); lanes.push(RelationLane { udf: prepared, - names: vec![(0..args.len()).map(|j| format!("_{j}")).collect()], + names: vec![arg_names], sides: vec![ args.into_iter() .map(|e| Arc::new(Compiled(e)) as Arc) diff --git a/crates/krishiv-exec/src/physical/mod.rs b/crates/krishiv-exec/src/physical/mod.rs index 36473dc7c..ccf2a5df6 100644 --- a/crates/krishiv-exec/src/physical/mod.rs +++ b/crates/krishiv-exec/src/physical/mod.rs @@ -1270,7 +1270,9 @@ pub struct PythonRelationPlan { pub output: Vec, } -/// The columns of one side's frames. +/// 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. diff --git a/crates/krishiv-exec/src/planner/python_relation.rs b/crates/krishiv-exec/src/planner/python_relation.rs index 02d6fe4a6..eca95a866 100644 --- a/crates/krishiv-exec/src/planner/python_relation.rs +++ b/crates/krishiv-exec/src/planner/python_relation.rs @@ -80,13 +80,36 @@ fn output_schema(output: &[BoundColumn]) -> SchemaRef { )) } -/// The frame side reading the first `columns.len()` columns of a relation. -fn side_of(columns: &[usize], cols: &[BoundColumn]) -> PythonFrameSide { - PythonFrameSide { - columns: columns.to_vec(), - names: cols.iter().map(|c| c.name.clone()).collect(), - types: cols.iter().map(|c| c.data_type.clone()).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. @@ -118,7 +141,7 @@ impl Planner<'_> { .collect::>>()?; let plan = PythonRelationPlan { udf: udf.clone(), - sides: vec![side_of(&columns, &cols)], + sides: vec![side_of(&[], &[], &columns, &cols)?], output: outputs(output), }; let explain = format!("{}#{}", udf.name, udf.version); @@ -147,15 +170,14 @@ impl Planner<'_> { .map(|c| planned.column(c.id)) .collect::>>()?; let sorted = self.clustered_and_sorted(planned, &key_columns, &[])?; - let keys = key_columns.clone(); let plan = PythonRelationPlan { udf: udf.clone(), - sides: vec![side_of(&columns, &cols)], + sides: vec![side_of(keys, &key_columns, &columns, &cols)?], output: outputs(output), }; let explain = format!( "[{}], {}#{}", - key_names(&sorted, &keys), + key_names(&sorted, &key_columns), udf.name, udf.version ); @@ -163,7 +185,7 @@ impl Planner<'_> { sorted, |input| PhysNode::PythonGroupMap { input, - keys, + keys: key_columns, udf: plan, }, output, @@ -274,13 +296,17 @@ impl Planner<'_> { 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), }; diff --git a/crates/krishiv-exec/src/python_relation.rs b/crates/krishiv-exec/src/python_relation.rs index 26addf2eb..4b5412258 100644 --- a/crates/krishiv-exec/src/python_relation.rs +++ b/crates/krishiv-exec/src/python_relation.rs @@ -7,8 +7,9 @@ //! (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 partition (one per driver), sends it the partition's rows -//! as [`Frame`]s and reads whatever batches it yields. Output rows are not +//! [`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` @@ -21,6 +22,19 @@ //! (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 @@ -55,13 +69,21 @@ pub struct RelationCallSpec { /// 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 hold - /// every input column, its keys among them. + /// 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 grouping keys (0 for a map kind): with [`Self::arity`] what the - /// Python side needs to call `f(pdf)` or `f(key, pdf)`. + /// 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. + 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 columns of a side's frames (the first side's for a cogroup). + /// 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, @@ -92,7 +114,8 @@ pub struct RunnerConf { /// One message to a [`UdfTask`]. #[derive(Debug, Clone)] pub enum Frame { - /// Rows of the current group (of the partition, for a map kind). + /// 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, @@ -107,7 +130,7 @@ pub enum Frame { /// One prepared relation call site. pub trait PreparedRelationUdf: Send + Sync { - /// One task per partition: the frames in, the result batches out. + /// One task per driver: the frames in, the result batches out. /// /// # Errors /// What the runtime raises starting it. diff --git a/crates/krishiv-exec/src/python_relation/sink.rs b/crates/krishiv-exec/src/python_relation/sink.rs index da19bf1b7..a4d329151 100644 --- a/crates/krishiv-exec/src/python_relation/sink.rs +++ b/crates/krishiv-exec/src/python_relation/sink.rs @@ -19,7 +19,8 @@ //! 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. +//! 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 held until a frame is full (`maxRecordsPerBatch`, and //! `maxBytesPerBatch` when set) or its group or side ends, as Spark cuts a //! group's Arrow stream; a group is never held whole. @@ -359,6 +360,8 @@ struct RelLocal { 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 held rows' bytes. held_bytes: Reservation, /// An aggregate's group keys, one row per group, in group order. @@ -521,9 +524,28 @@ impl RelationSink { Ok(()) } - /// Ends the open group: its held rows, then `EndGroup` to every lane. + /// Sends every lane a zero-row frame of `side`: a cogroup key without + /// rows on that side (Spark gives the UDF an empty DataFrame for it). + fn empty_side(&self, l: &mut RelLocal, side: u8) -> Result<()> { + let empty = RecordBatch::new_empty(Arc::clone(&self.input)); + for (lane, state) in self.lanes.iter().zip(&l.lanes) { + let frame = self.frame(lane, side, &empty, l.partition)?; + state.queue(lane.udf.as_ref(), frame, l.held_bytes.split(0))?; + } + Ok(()) + } + + fn cogrouped(&self) -> bool { + matches!(self.shape, RelationShape::CoGrouped { .. }) + } + + /// Ends the open group: its held rows (a cogroup's empty right side), + /// then `EndGroup` to every lane. fn end_group(&self, l: &mut RelLocal) -> Result<()> { self.send_held(l, true)?; + if self.cogrouped() && !l.sides_seen[1] { + self.empty_side(l, 1)?; + } for (lane, state) in self.lanes.iter().zip(&l.lanes) { let r = l.held_bytes.split(0); state.queue(lane.udf.as_ref(), Frame::EndGroup, r)?; @@ -547,6 +569,7 @@ impl RelationSink { self.end_group(l)?; } l.group_open = true; + l.sides_seen = [false; 2]; if let Some(k) = keys.next() { l.key_rows.push(k); } @@ -557,6 +580,11 @@ impl RelationSink { self.send_held(l, true)?; l.held_side = run.side; } + if self.cogrouped() && run.side == 1 && !l.sides_seen[0] { + // A key on the right only: its empty left side first. + self.empty_side(l, 0)?; + } + l.sides_seen[usize::from(run.side)] = true; l.held.push(batch.slice(run.start, run.len)); l.held_rows += run.len; self.send_held(l, false)?; @@ -735,6 +763,7 @@ impl Sink for RelationSink { held: Vec::new(), held_rows: 0, held_side: 0, + sides_seen: [false; 2], held_bytes: consumer.reservation(), key_rows: Vec::new(), key_bytes: Vec::new(), diff --git a/crates/krishiv-exec/src/python_relation/tests.rs b/crates/krishiv-exec/src/python_relation/tests.rs index 086cd342d..9498b006e 100644 --- a/crates/krishiv-exec/src/python_relation/tests.rs +++ b/crates/krishiv-exec/src/python_relation/tests.rs @@ -391,7 +391,7 @@ async fn a_grouped_task_without_keys_sees_one_group() { } #[tokio::test] -async fn a_cogroup_key_on_one_side_gets_only_that_side_then_end_group() { +async fn a_cogroup_key_on_one_side_gets_a_zero_row_frame_for_the_other() { let r = rig( RelationShape::CoGrouped { keys: vec![0], @@ -414,7 +414,9 @@ async fn a_cogroup_key_on_one_side_gets_only_that_side_then_end_group() { .unwrap(); assert_eq!( r.shared.log(), - ["B0:1", "G", "B0:1", "B1:2", "G", "B1:1", "G", "E"] + [ + "B0:1", "B1:0", "G", "B0:1", "B1:2", "G", "B0:0", "B1:1", "G", "E" + ] ); done(r); } @@ -546,6 +548,7 @@ fn a_runtime_without_relation_kinds_refuses_them_as_not_available() { 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()], diff --git a/crates/krishiv-exec/tests/planner/python_relations.rs b/crates/krishiv-exec/tests/planner/python_relations.rs index 5cf4a84c5..6a2baa780 100644 --- a/crates/krishiv-exec/tests/planner/python_relations.rs +++ b/crates/krishiv-exec/tests/planner/python_relations.rs @@ -41,6 +41,9 @@ struct Runtime { 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, } type Log = Arc>>; @@ -62,11 +65,24 @@ impl UdfRuntime for Runtime { "twice" => Mode::Twice, _ => 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), + names: spec.input_names.clone(), + key_count: spec.key_count, })) } } @@ -77,6 +93,9 @@ struct Prepared { output: arrow_schema::SchemaRef, log: Arc>>, per_task: Arc>>, + keys: Log, + names: Vec>, + key_count: usize, } impl PreparedRelationUdf for Prepared { @@ -88,6 +107,9 @@ impl PreparedRelationUdf for Prepared { mode: self.mode, output: Arc::clone(&self.output), log: Arc::clone(&self.log), + keys: Arc::clone(&self.keys), + names: self.names.clone(), + key_count: self.key_count, state: std::sync::Mutex::new(State::default()), ready: AtomicWaker::new(), })) @@ -107,11 +129,47 @@ struct Task { mode: Mode, output: arrow_schema::SchemaRef, log: Arc>>, + keys: Log, + 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) + }; + for r in 1..batch.num_rows() { + assert_eq!(key(r), first, "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); @@ -132,6 +190,7 @@ impl UdfTask for Task { 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 = @@ -350,21 +409,59 @@ fn a_cogroup_gives_each_key_both_sides_and_a_one_sided_key_its_side() { assert_eq!(rows.len(), 9, "{rows:?}"); assert_eq!(rows.iter().filter(|r| *r == "2").count(), 3, "{rows:?}"); let log = log_of(&rt); - assert_eq!( - log.iter().filter(|e| e.starts_with("B1")).count(), - 2, - "{log:?}" - ); - // Within each task, a key's left rows come before its right rows. + 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 w in own.windows(2) { - assert!( - !(w[0].starts_with("B1") && w[1].starts_with("B0")), - "{own:?}" - ); + 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"]); } /// A grouped-aggregate Python UDF call over `arg`. From 62dea5cc3096562de69a74df6a5ea25400d92bec Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Fri, 13 Nov 2026 18:03:59 +0000 Subject: [PATCH 08/29] exec: Python relation memory: reserved results, per-frame window, streamed aggregate (P3-T014a review Md1-Md3, m3) - Each task result is reserved as the pump reads it; a refused one stops the task's reads until the driver reserves it (sink returns NeedMemory, finalize waits). finalize stores results as they arrive, writing a part to a spill file when the pool has no room. - The frame window is checked per frame: rows cut but not framed wait in a reserved backlog. A chunk's reservation is shared by the lanes' frames and released by the last taker; client-form copies are reserved. - A grouped aggregate keeps one key batch per input batch and joins the oldest pending groups with the values as every lane yields them, so the rows reach the spillable store; a task yielding another number of values than groups fails. - Result mismatches raise KRISHIV_PYTHON_UDF_RESULT_MISMATCH (39000), not INTERNAL_ERROR. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../krishiv-common/krishiv-error-classes.json | 6 + crates/krishiv-exec/CLAUDE.md | 2 +- .../physical/instantiate/python_relation.rs | 2 + crates/krishiv-exec/src/python_relation.rs | 2 + .../krishiv-exec/src/python_relation/lane.rs | 224 +++++ .../krishiv-exec/src/python_relation/pair.rs | 235 +++++ .../krishiv-exec/src/python_relation/sink.rs | 908 +++++++++--------- .../krishiv-exec/src/python_relation/tests.rs | 118 ++- crates/krishiv-exec/src/shared.rs | 45 + .../tests/planner/python_relations.rs | 102 +- 10 files changed, 1172 insertions(+), 472 deletions(-) create mode 100644 crates/krishiv-exec/src/python_relation/lane.rs create mode 100644 crates/krishiv-exec/src/python_relation/pair.rs diff --git a/crates/krishiv-common/krishiv-error-classes.json b/crates/krishiv-common/krishiv-error-classes.json index 395752cc9..54061bb0e 100644 --- a/crates/krishiv-common/krishiv-error-classes.json +++ b/crates/krishiv-common/krishiv-error-classes.json @@ -105,6 +105,12 @@ "sqlState" : "58030", "audience" : "operator" }, + "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_QUERY_ABORTED_OUT_OF_MEMORY" : { "message" : [ "The query was aborted to free memory for other work: it held bytes, the most of any query, and spilling did not free enough memory in time. Increase \"krishiv.memory.limit\" or run fewer concurrent queries." diff --git a/crates/krishiv-exec/CLAUDE.md b/crates/krishiv-exec/CLAUDE.md index 1ab56e2d8..77e53d27a 100644 --- a/crates/krishiv-exec/CLAUDE.md +++ b/crates/krishiv-exec/CLAUDE.md @@ -126,7 +126,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, `Blocked` while a lane's frame window is full, results in the shared plan's list); `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/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, 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 (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/physical/instantiate/python_relation.rs b/crates/krishiv-exec/src/physical/instantiate/python_relation.rs index b4927e570..b42ccd38f 100644 --- a/crates/krishiv-exec/src/physical/instantiate/python_relation.rs +++ b/crates/krishiv-exec/src/physical/instantiate/python_relation.rs @@ -195,6 +195,7 @@ impl Instantiation<'_> { }; 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(), @@ -271,6 +272,7 @@ impl Instantiation<'_> { 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![ diff --git a/crates/krishiv-exec/src/python_relation.rs b/crates/krishiv-exec/src/python_relation.rs index 4b5412258..59a3a4a0d 100644 --- a/crates/krishiv-exec/src/python_relation.rs +++ b/crates/krishiv-exec/src/python_relation.rs @@ -51,6 +51,8 @@ use krishiv_types::DataType; use crate::error::Result; +mod lane; +mod pair; pub(crate) mod sink; pub use sink::{ 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..3d4e36cb8 --- /dev/null +++ b/crates/krishiv-exec/src/python_relation/lane.rs @@ -0,0 +1,224 @@ +// 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 reserved it +//! (`sink` returns `NeedMemory`, `finalize` awaits the memory): back +//! pressure into the task, never unaccounted results. + +use std::collections::VecDeque; +use std::future::poll_fn; +use std::sync::Arc; +use std::task::{Context, Poll}; + +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, +} + +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 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>, +} + +impl Lane { + /// A lane whose results reserve in `memory`'s consumer. + pub(super) fn new(memory: Reservation) -> Self { + Self { + io: Arc::new(Mutex::new(LaneIo { + inbound: VecDeque::new(), + sending: None, + outputs: VecDeque::new(), + refused: None, + memory, + })), + pump: Arc::new(Mutex::new(Pump::Idle)), + } + } + + /// 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.io))); + } + } + 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 + } + + /// Reserves the refused result, if any, waiting for the memory, and + /// moves it to the outputs. + pub(super) async fn take_refused(&self) -> Result<()> { + let (bytes, mut r) = { + let mut q = lock(&self.io); + let Some(batch) = &q.refused else { + return Ok(()); + }; + (batch.get_array_memory_size(), q.memory.split(0)) + }; + r.grow(bytes).await.map_err(KrishivError::from)?; + let mut q = lock(&self.io); + if let Some(batch) = q.refused.take() { + q.outputs.push_back((batch, r)); + } + Ok(()) + } +} + +/// Sends the queued frames to `task` while it reads the task's results. +fn pump(task: Box, io: Arc>) -> BoxFuture<'static, Result<()>> { + Box::pin(async move { + let task = task; + 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?; + lock(&io).sending = None; + if end { + return Ok::<(), KrishivError>(()); + } + } + }; + let receive = async { + loop { + // Never woken either: whoever reserves the refused batch + // polls the pump next. + poll_fn(|_| { + if lock(&io).refused.is_some() { + Poll::Pending + } else { + Poll::Ready(()) + } + }) + .await; + let Some(batch) = task.next().await? else { + return Ok(()); + }; + let bytes = batch.get_array_memory_size(); + let mut q = lock(&io); + 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); + } + } + }; + futures::future::try_join(send, receive).await.map(drop) + }) +} 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..9dab87333 --- /dev/null +++ b/crates/krishiv-exec/src/python_relation/pair.rs @@ -0,0 +1,235 @@ +// 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.outputs.pop_front() { + memory.push(r); + } + } else { + *front = front.slice(len, front.num_rows() - 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 index a4d329151..ec6f1e334 100644 --- a/crates/krishiv-exec/src/python_relation/sink.rs +++ b/crates/krishiv-exec/src/python_relation/sink.rs @@ -7,12 +7,13 @@ //! 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`] 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. +//! `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 @@ -21,19 +22,27 @@ //! 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 held until a frame is full (`maxRecordsPerBatch`, and -//! `maxBytesPerBatch` when set) or its group or side ends, as Spark cuts a -//! group's Arrow stream; a group is never held whole. +//! Rows are cut into frames of `maxRecordsPerBatch` rows (and +//! `maxBytesPerBatch` bytes when set) or at the end of their group or side, +//! as Spark cuts a group's Arrow stream; a group is never held whole. //! -//! **Memory.** Held rows and queued frames are reserved in the driver's -//! consumer (a refused batch returns `NeedMemory` before anything changes, -//! so the retry is exact); a frame's reservation is released once the task -//! took it. Results wait in the lane's queue until the next `sink` call -//! stores them; an aggregate keeps one key row per group until `finalize`. +//! **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. +//! +//! **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 the copies made for the client forms), 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; `finalize` waits for the memory. //! //! **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. +//! which cancels it, and every reservation. use std::collections::VecDeque; use std::fmt; @@ -46,11 +55,13 @@ 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; +use krishiv_common::error::{KrishivError, KrishivPythonUdfResultMismatch}; use krishiv_memory::{Reservation, SpillScope}; use krishiv_types::batch::Batch; -use super::{Frame, PreparedRelationUdf, UdfTask}; +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}; @@ -101,6 +112,8 @@ impl RelationShape { /// 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. @@ -115,6 +128,7 @@ pub struct RelationLane { 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) @@ -236,127 +250,31 @@ impl RelationSink { } } -/// A lane's queues, shared by the driver and the lane's pump. -#[derive(Default)] -struct LaneIo { - /// Frames not taken by the task yet, each with its reservation. - inbound: VecDeque<(Frame, Reservation)>, - /// The reservation of the frame the task is taking. - sending: Option, - /// Batches the task yielded, not stored yet. - outputs: VecDeque, -} - -impl LaneIo { - fn in_flight(&self) -> usize { - self.inbound.len() + usize::from(self.sending.is_some()) - } -} - -/// 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)] -struct Lane { - io: Arc>, - pump: Arc>, -} - -impl Lane { - fn new() -> Self { - Self { - io: Arc::new(Mutex::new(LaneIo::default())), - pump: Arc::new(Mutex::new(Pump::Idle)), - } - } - - /// Polls the pump; its outcome once it ended (`None` while it runs or - /// before the task opened). - fn poll(pump: &Mutex, cx: &mut Context<'_>) -> Option> { - let mut p = lock(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, - } - } - - fn opened(&self) -> bool { - !matches!(*lock(&self.pump), Pump::Idle) - } - - /// Queues `frame`, opening the task first if it is not. - fn queue(&self, udf: &dyn PreparedRelationUdf, frame: Frame, r: Reservation) -> Result<()> { - { - let mut p = lock(&self.pump); - if matches!(*p, Pump::Idle) { - let task = udf.open()?; - *p = Pump::Running(pump(task, Arc::clone(&self.io))); - } - } - lock(&self.io).inbound.push_back((frame, r)); - Ok(()) - } -} - -/// Sends the queued frames to `task` while it reads the task's results. -fn pump(task: Box, io: Arc>) -> BoxFuture<'static, Result<()>> { - Box::pin(async move { - let task = task; - 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, r)) => { - q.sending = Some(r); - Poll::Ready(frame) - } - None => Poll::Pending, - } - }) - .await; - let end = matches!(frame, Frame::End); - task.send(frame).await?; - lock(&io).sending = None; - if end { - return Ok::<(), KrishivError>(()); - } - } - }; - let receive = async { - while let Some(batch) = task.next().await? { - lock(&io).outputs.push_back(batch); - } - Ok(()) - }; - futures::future::try_join(send, receive).await.map(drop) - }) +/// 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. -struct RelLocal { - store: LocalState, +/// 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). - ready: Option, + /// 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 sent yet. + /// Rows of the open group (and side) not cut yet. held: Vec, held_rows: usize, held_side: u8, @@ -364,23 +282,35 @@ struct RelLocal { sides_seen: [bool; 2], /// The reservation of the held rows' bytes. held_bytes: Reservation, - /// An aggregate's group keys, one row per group, in group order. - key_rows: Vec, - key_bytes: Vec, + /// Cut, waiting for room in the lanes' windows. + backlog: VecDeque, + /// Memory reserved for the next frame's copies (`finalize` waits for + /// it; `sink` asks the driver for it). + spare: Reservation, + /// An aggregate's groups waiting for their values. + pending: Pending, } -/// What a driver left at `combine`. -struct Leftover { - lanes: Vec, - ready: Option, - key_rows: Vec, - _key_bytes: Vec, +/// What [`RelationSink::flush`] stopped at. +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 global state: the store's, and what drivers left. struct RelGlobal { store: GlobalState, - leftover: Mutex>, + leftover: Mutex>, } /// A run of input rows within one group and side. @@ -445,22 +375,29 @@ impl RelationSink { } /// The frame of `lane` over `rows` (one side): its columns evaluated, - /// in their client forms. + /// in their client forms; and the bytes of the columns copied to reach + /// them. fn frame( - &self, lane: &RelationLane, side: u8, rows: &RecordBatch, partition: u32, - ) -> Result { + ) -> 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; for (e, name) in exprs.iter().zip(names) { let a = e.evaluate_in(rows, partition)?; let client = krishiv_types::arrow::client_arrow_type(a.data_type()); - let a = crate::carrier::recarried_array(&a, &client)?; + let a = if a.data_type() == &client { + a + } else { + let moved = crate::carrier::recarried_array(&a, &client)?; + copied += moved.get_array_memory_size(); + moved + }; fields.push(Field::new(name, client, e.nullable())); columns.push(a); } @@ -468,18 +405,18 @@ impl RelationSink { 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 }) + Ok((Frame::Batch { side, batch }, copied)) } - /// Sends the held rows in frames of at most `maxRecordsPerBatch` rows + /// 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 send_held(&self, l: &mut RelLocal, all: bool) -> Result<()> { - if l.held_rows == 0 { + fn cut_held(&self, d: &mut Driver, all: bool) -> Result<()> { + if d.held_rows == 0 { return Ok(()); } - let held = std::mem::take(&mut l.held); - let rows = std::mem::take(&mut l.held_rows); + 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() @@ -503,36 +440,29 @@ impl RelationSink { let len = step.min(rows - offset); if len < step && !all { // Too few for a full frame: held for the next input. - l.held.push(data.slice(offset, len)); - l.held_rows = len; + d.held.push(data.slice(offset, len)); + d.held_rows = len; break; } - let chunk = data.slice(offset, len); - let share = l.held_bytes.size() * len / (rows - offset); - let mut part = Some(l.held_bytes.split(share)); - for (lane, state) in self.lanes.iter().zip(&l.lanes) { - let frame = self.frame(lane, l.held_side, &chunk, l.partition)?; - // The chunk's bytes go with the first lane's frame. - let r = match part.take() { - Some(r) => r, - None => l.held_bytes.split(0), - }; - state.queue(lane.udf.as_ref(), frame, r)?; - } + let share = d.held_bytes.size() * len / (rows - offset); + d.backlog.push_back(Cut::Rows { + side: d.held_side, + rows: data.slice(offset, len), + memory: d.held_bytes.split(share), + }); offset += len; } Ok(()) } - /// Sends every lane a zero-row frame of `side`: a cogroup key without - /// rows on that side (Spark gives the UDF an empty DataFrame for it). - fn empty_side(&self, l: &mut RelLocal, side: u8) -> Result<()> { - let empty = RecordBatch::new_empty(Arc::clone(&self.input)); - for (lane, state) in self.lanes.iter().zip(&l.lanes) { - let frame = self.frame(lane, side, &empty, l.partition)?; - state.queue(lane.udf.as_ref(), frame, l.held_bytes.split(0))?; - } - Ok(()) + /// A zero-row frame of `side`: a cogroup key without rows on that side + /// (Spark gives the UDF an empty DataFrame for it). + fn empty_side(&self, d: &mut Driver, side: u8) { + d.backlog.push_back(Cut::Rows { + side, + rows: RecordBatch::new_empty(Arc::clone(&self.input)), + memory: d.held_bytes.split(0), + }); } fn cogrouped(&self) -> bool { @@ -540,80 +470,101 @@ impl RelationSink { } /// Ends the open group: its held rows (a cogroup's empty right side), - /// then `EndGroup` to every lane. - fn end_group(&self, l: &mut RelLocal) -> Result<()> { - self.send_held(l, true)?; - if self.cogrouped() && !l.sides_seen[1] { - self.empty_side(l, 1)?; - } - for (lane, state) in self.lanes.iter().zip(&l.lanes) { - let r = l.held_bytes.split(0); - state.queue(lane.udf.as_ref(), Frame::EndGroup, r)?; + /// 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); } + d.backlog.push_back(Cut::EndGroup); Ok(()) } - /// Feeds `batch` (reserved in `held_bytes`) to the lanes. - fn feed( - &self, - l: &mut RelLocal, - batch: &RecordBatch, - runs: Vec, - keys: Vec, - ) -> Result<()> { + /// 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(); - let mut keys = keys.into_iter(); for run in runs { if run.new_group { - if l.group_open { - self.end_group(l)?; - } - l.group_open = true; - l.sides_seen = [false; 2]; - if let Some(k) = keys.next() { - l.key_rows.push(k); + if d.group_open { + self.end_group(d)?; } - } else if grouped && !l.group_open { + 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 != l.held_side { - self.send_held(l, true)?; - l.held_side = run.side; + if run.side != d.held_side { + self.cut_held(d, true)?; + d.held_side = run.side; } - if self.cogrouped() && run.side == 1 && !l.sides_seen[0] { + if self.cogrouped() && run.side == 1 && !d.sides_seen[0] { // A key on the right only: its empty left side first. - self.empty_side(l, 0)?; + self.empty_side(d, 0); } - l.sides_seen[usize::from(run.side)] = true; - l.held.push(batch.slice(run.start, run.len)); - l.held_rows += run.len; - self.send_held(l, false)?; + 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(()) } - /// The copied key row of each run starting a group (an aggregate's). - fn key_rows(&self, batch: &RecordBatch, runs: &[Run]) -> Result> { - let (RelationShape::Aggregate { keys }, Some(schema)) = (&self.shape, &self.key_schema) - else { - return Ok(Vec::new()); - }; - let projected = batch - .project(keys) - .map_err(|e| internal(format!("aggregate keys: {e}")))?; - let projected = RecordBatch::try_new_with_options( - Arc::clone(schema), - projected.columns().to_vec(), - &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())), - ) - .map_err(|e| internal(format!("aggregate keys: {e}")))?; - runs.iter() - .filter(|r| r.new_group) - .map(|r| { - arrow_select::concat::concat_batches(schema, [&projected.slice(r.start, 1)]) - .map_err(|e| internal(format!("aggregate keys: {e}"))) - }) - .collect() + /// Frames the backlog into the lanes while every lane's window has + /// room. Nothing changes for a frame whose copies the pool refuses. + fn flush(&self, d: &mut Driver) -> Result { + loop { + if d.lanes + .iter() + .any(|l| lock(&l.io).in_flight() >= self.window) + { + return Ok(Flushed::Window); + } + let Some(cut) = d.backlog.front() else { + return Ok(Flushed::All); + }; + match cut { + Cut::Rows { side, rows, .. } => { + let mut frames = Vec::with_capacity(self.lanes.len()); + for lane in &self.lanes { + frames.push(Self::frame(lane, *side, rows, d.partition)?); + } + let copied: usize = frames.iter().map(|(_, c)| c).sum(); + let more = copied.saturating_sub(d.spare.size()); + if more > 0 && d.spare.try_grow(more).is_err() { + return Ok(Flushed::Memory(more)); + } + let Some(Cut::Rows { memory, .. }) = d.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 + self.lanes.iter().zip(&d.lanes).zip(frames) + { + let memory = FrameMem { + _shared: Some(Arc::clone(&shared)), + _own: d.spare.split(copied), + }; + state.queue(lane.udf.as_ref(), frame, memory)?; + } + } + Cut::EndGroup | Cut::End => { + let end = matches!(cut, Cut::End); + d.backlog.pop_front(); + for (lane, state) in self.lanes.iter().zip(&d.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: d.spare.split(0), + }; + state.queue(lane.udf.as_ref(), frame, memory)?; + } + } + } + } } /// A task's batch as an output batch: its columns checked against the @@ -622,11 +573,14 @@ impl RelationSink { /// the engine's own guard. fn output(&self, lane: &RelationLane, batch: &RecordBatch) -> Result { if batch.num_columns() != lane.client.len() { - return Err(internal(format!( - "a Python UDF returned {} columns where {} were declared", - 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 @@ -635,79 +589,49 @@ impl RelationSink { .zip(&lane.client) .zip(self.schema.fields()) { - columns.push(recarried(c, client, field.data_type())?); + 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| internal(format!("a Python UDF result does not fit its schema: {e}")))?; + .map_err(|e| mismatch(lane, format!("{e}")))?; Ok(Batch::new(out)) } - /// An aggregate driver's rows: each group's keys, then each lane's value - /// of it. - fn aggregate_rows(&self, key_rows: &[RecordBatch], lanes: &[Lane]) -> Result> { - let Some(schema) = &self.key_schema else { - return Ok(None); - }; - if key_rows.is_empty() { - return Ok(None); - } - let groups = key_rows.iter().map(RecordBatch::num_rows).sum::(); - let keys = arrow_select::concat::concat_batches(schema, key_rows) - .map_err(|e| internal(format!("aggregate keys: {e}")))?; - let mut columns: Vec = keys.columns().to_vec(); - let width = columns.len(); - for (i, (lane, state)) in self.lanes.iter().zip(lanes).enumerate() { - let outputs: Vec = lock(&state.io).outputs.drain(..).collect(); - let values: Vec<&ArrayRef> = outputs - .iter() - .map(|b| { - if b.num_columns() == 1 { - Ok(b.column(0)) - } else { - Err(internal(format!( - "a Python aggregate UDF returned {} columns", - b.num_columns() - ))) - } - }) - .collect::>()?; - let refs: Vec<&dyn Array> = values.iter().map(|a| AsRef::as_ref(*a)).collect(); - let client = lane - .client - .first() - .ok_or_else(|| internal("an aggregate lane without a type".to_owned()))?; - let column = if refs.is_empty() { - arrow_array::new_empty_array(client) - } else { - arrow_select::concat::concat(&refs) - .map_err(|e| internal(format!("aggregate results: {e}")))? - }; - if column.len() != groups { - return Err(internal(format!( - "a Python aggregate UDF returned {} rows for {groups} groups", - column.len() - ))); + /// 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).outputs.pop_front(); + if let Some((batch, memory)) = next { + return Ok(Some((self.output(lane, &batch)?, vec![memory]))); } - let field = self - .schema - .fields() - .get(width + i) - .ok_or_else(|| internal("an aggregate result out of its schema".to_owned()))?; - columns.push(recarried(&column, client, field.data_type())?); - } - let options = RecordBatchOptions::new().with_row_count(Some(groups)); - let out = RecordBatch::try_new_with_options(Arc::clone(&self.schema), columns, &options) - .map_err(|e| internal(format!("a Python aggregate result: {e}")))?; - Ok(Some(Batch::new(out))) + } + Ok(None) } /// Stores `batch` in the list, or keeps it for the retry after the /// memory the store asked for. - fn store(&self, l: &mut RelLocal, batch: Batch, g: &RelGlobal) -> Result> { + fn store( + &self, + l: &mut RelLocal, + batch: Batch, + memory: Vec, + g: &RelGlobal, + ) -> Result> { match self.store.sink(batch.clone(), &mut l.store, &g.store)? { SinkStatus::NeedMemory { bytes } => { - l.ready = Some(batch); + l.driver.ready = Some((batch, memory)); Ok(Some(SinkStatus::NeedMemory { bytes })) } _ => Ok(None), @@ -716,28 +640,121 @@ impl RelationSink { /// Moves each lane's frames and results without waiting; the first /// failure of a lane. - fn progress(l: &RelLocal) -> Result<()> { + fn progress(d: &Driver) -> Result<()> { let mut cx = Context::from_waker(noop_waker_ref()); - for lane in &l.lanes { - if let Some(Err(e)) = Lane::poll(&lane.pump, &mut cx) { + 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 { keys }, 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, keys, 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); + } + 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`. -fn recarried(c: &ArrayRef, client: &ArrowType, canonical: &ArrowType) -> Result { +pub(super) fn recarried( + lane: &RelationLane, + c: &ArrayRef, + client: &ArrowType, + canonical: &ArrowType, +) -> Result { if c.data_type() != client { - return Err(internal(format!( - "a Python UDF returned {} where {client} was declared", - c.data_type() - ))); + return Err(mismatch( + lane, + format!("{} where {client} was declared", c.data_type()), + )); } crate::carrier::recarried_array(c, canonical) } +/// The future of a `Blocked` driver: ready once every lane's window has +/// room, a lane ended (its failure is the next call's), or a result waits +/// for memory (the next call asks for it). +fn blocked(lanes: Vec, window: usize) -> BoxFuture<'static, ()> { + Box::pin(poll_fn(move |cx| { + let mut room = true; + for lane in &lanes { + if lane.poll(cx).is_some() || lock(&lane.io).refused.is_some() { + return Poll::Ready(()); + } + room &= lock(&lane.io).in_flight() < window; + } + if room { Poll::Ready(()) } else { Poll::Pending } + })) +} + impl Sink for RelationSink { fn name(&self) -> &'static str { "PythonRelation" @@ -755,18 +772,25 @@ impl Sink for RelationSink { let consumer = ctx.consumer().consumer(); Ok(Box::new(RelLocal { store: self.store.local_state(ctx, &g.store)?, - lanes: self.lanes.iter().map(|_| Lane::new()).collect(), - 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(), - key_rows: Vec::new(), - key_bytes: Vec::new(), + driver: Driver { + lanes: self + .lanes + .iter() + .map(|_| Lane::new(consumer.reservation())) + .collect(), + 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(), + backlog: VecDeque::new(), + spare: consumer.reservation(), + pending: Pending::default(), + }, })) } @@ -778,63 +802,28 @@ impl Sink for RelationSink { ) -> Result { let g = global_ref::(global, "PythonRelation")?; let l = local_mut::(local, "PythonRelation")?; - // Store what finished first: a `NeedMemory` here leaves the input - // unsent, so the retry with the same batch is exact. - if let Some(ready) = l.ready.take() - && let Some(status) = self.store(l, ready, g)? - { + // 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); } - Self::progress(l)?; - if !matches!(self.shape, RelationShape::Aggregate { .. }) { - loop { - let next = l.lanes.iter().zip(&self.lanes).find_map(|(state, lane)| { - lock(&state.io).outputs.pop_front().map(|b| (lane, b)) - }); - let Some((lane, batch)) = next else { break }; - let batch = self.output(lane, &batch)?; - if let Some(status) = self.store(l, batch, g)? { - return Ok(status); - } - } + if let Flushed::Memory(bytes) = self.flush(&mut l.driver)? { + return Ok(SinkStatus::NeedMemory { bytes }); } if input.num_rows() > 0 { - let input = crate::operators::owned_rows(input)?; - let data = input.data(); - let (runs, last) = self.runs(data, l.current.as_ref())?; - let keys = self.key_rows(data, &runs)?; - let key_bytes: usize = keys.iter().map(RecordBatch::get_array_memory_size).sum(); - let bytes = data.get_array_memory_size(); - if l.held_bytes.try_grow(bytes + key_bytes).is_err() { - return Ok(SinkStatus::NeedMemory { - bytes: bytes + key_bytes, - }); + if let Some(status) = self.take_input(&mut l.driver, input)? { + return Ok(status); } - let part = l.held_bytes.split(key_bytes); - l.key_bytes.push(part); - self.feed(l, data, runs, keys)?; - l.current = last; - Self::progress(l)?; - } - let full = l - .lanes + // A refusal here waits for the next call, which asks for it. + self.flush(&mut l.driver)?; + Self::progress(&l.driver)?; + } + let d = &l.driver; + if d.lanes .iter() - .any(|lane| lock(&lane.io).in_flight() >= self.window); - if full { - let lanes: Vec = l.lanes.clone(); - let window = self.window; - return Ok(SinkStatus::Blocked(Box::pin(poll_fn(move |cx| { - let mut room = true; - for lane in &lanes { - // A lane that ended (or failed) frees the driver: the - // next `sink` call returns its failure. - if Lane::poll(&lane.pump, cx).is_some() { - return Poll::Ready(()); - } - room &= lock(&lane.io).in_flight() < window; - } - if room { Poll::Ready(()) } else { Poll::Pending } - })))); + .any(|lane| lock(&lane.io).in_flight() >= self.window) + { + return Ok(SinkStatus::Blocked(blocked(d.lanes.clone(), self.window))); } Ok(SinkStatus::NeedMoreInput) } @@ -844,34 +833,18 @@ impl Sink for RelationSink { let l = local .downcast::() .map_err(|_| internal("wrong local state for PythonRelation".to_owned()))?; - let mut l = *l; - // The input ended: the open group ends, then the partition. - if l.group_open { - self.end_group(&mut l)?; + 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.send_held(&mut l, true)?; - } - for (lane, state) in self.lanes.iter().zip(&l.lanes) { - if state.opened() { - let r = l.held_bytes.split(0); - state.queue(lane.udf.as_ref(), Frame::End, r)?; - } + self.cut_held(&mut driver, true)?; } - RelationSink::progress(&l)?; - let RelLocal { - store, - lanes, - ready, - key_rows, - key_bytes, - .. - } = l; - lock(&g.leftover).push(Leftover { - lanes, - ready, - key_rows, - _key_bytes: key_bytes, - }); + driver.backlog.push_back(Cut::End); + self.flush(&mut driver)?; + Self::progress(&driver)?; + lock(&g.leftover).push(driver); self.store.combine(store, &g.store) } @@ -882,53 +855,27 @@ impl Sink for RelationSink { ) -> BoxFuture<'a, Result<()>> { Box::pin(async move { let g = global_ref::(global, "PythonRelation")?; - let leftover = std::mem::take(&mut *lock(&g.leftover)); - // Every task at once: the first failure (or a cancel) ends the - // wait, and the error drops the other tasks, which cancels them. - 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())); - } - let mut all = true; - for lane in leftover.iter().flat_map(|d| &d.lanes) { - match Lane::poll(&lane.pump, cx) { - Some(Err(e)) => return Poll::Ready(Err(e)), - Some(Ok(())) => {} - None => all &= !lane.opened(), - } - } - if all { - Poll::Ready(Ok(())) - } else { - Poll::Pending - } - }) - .await?; - for driver in leftover { + 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 batches: Vec = driver.ready.into_iter().collect(); - if matches!(self.shape, RelationShape::Aggregate { .. }) { - batches.extend(self.aggregate_rows(&driver.key_rows, &driver.lanes)?); - } else { - for (lane, state) in self.lanes.iter().zip(&driver.lanes) { - let outputs: Vec = lock(&state.io).outputs.drain(..).collect(); - for b in &outputs { - batches.push(self.output(lane, b)?); - } - } + let mut done = true; + for d in &mut drivers { + done &= self.finish_step(d, g, ctx).await?; } - if batches.is_empty() { - continue; + if done { + break; + } + self.wait(&drivers, ctx).await?; + } + if matches!(self.shape, RelationShape::Aggregate { .. }) { + for d in &drivers { + pair::check_complete(&d.pending, &self.lanes, &d.lanes)?; } - 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)?; - self.store.add_part(&g.store, batches, reservation)?; } + drop(drivers); self.store.finalize(&g.store, ctx).await }) } @@ -946,6 +893,97 @@ impl Sink for RelationSink { } } +impl RelationSink { + /// One step of a finished driver in `finalize`: the backlog framed as + /// far as the windows allow (waiting for the copies' memory), the + /// results stored (written to a spill file when the pool has no room), + /// then refused results reserved. Whether the driver + /// is done: no backlog, every task ended, every result stored. + async fn finish_step(&self, d: &mut Driver, g: &RelGlobal, ctx: &TaskContext) -> Result { + while let Flushed::Memory(bytes) = self.flush(d)? { + d.spare.grow(bytes).await.map_err(KrishivError::from)?; + } + 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); + } + 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 mut reservation = ctx.consumer().reservation(); + if reservation.try_grow(bytes).is_ok() { + self.store.add_part(&g.store, batches, reservation)?; + } else { + drop(reservation); + self.store.add_spilled_part(&g.store, batches, ctx).await?; + } + } + drop(held); + // The results' memory is free again: a refused result next. + for lane in &d.lanes { + lane.take_refused().await?; + } + Ok(d.backlog.is_empty() + && d.ready.is_none() + && d.lanes + .iter() + .all(|l| l.finished() && !lock(&l.io).has_results())) + } + + /// Waits until a driver can move: a lane yielded, 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 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())); + } + let mut ready = false; + let mut all_done = true; + for d in drivers { + let mut room = true; + 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); + ready |= q.has_results(); + room &= q.in_flight() < self.window; + } + ready |= room && !d.backlog.is_empty(); + all_done &= d.backlog.is_empty() && d.ready.is_none(); + } + if ready || all_done { + Poll::Ready(Ok(())) + } else { + Poll::Pending + } + }) + .await + } +} + +/// 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(); + Ok((lanes, d.backlog.len())) +} + /// 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 { diff --git a/crates/krishiv-exec/src/python_relation/tests.rs b/crates/krishiv-exec/src/python_relation/tests.rs index 9498b006e..2339a992a 100644 --- a/crates/krishiv-exec/src/python_relation/tests.rs +++ b/crates/krishiv-exec/src/python_relation/tests.rs @@ -45,6 +45,8 @@ struct Shared { 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, } impl Shared { @@ -67,6 +69,10 @@ enum Yield { 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), } struct Fake { @@ -146,15 +152,24 @@ impl UdfTask for Task { match frame { Frame::Batch { batch, .. } => { s.rows += i64::try_from(batch.num_rows()).unwrap(); - if matches!(self.yields, Yield::Twice) { + let copies = match self.yields { + Yield::Twice => 2, + Yield::Burst(n) => n, + Yield::Nothing | Yield::Count | Yield::CountTwice => 0, + }; + for _ in 0..copies { s.out.push_back(batch.clone()); - s.out.push_back(batch); } } Frame::EndGroup => { - if matches!(self.yields, Yield::Count) { - let n = std::mem::take(&mut s.rows); - s.out.push_back(count(n)); + 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)); + } + _ => {} } } Frame::End => s.ended = true, @@ -170,7 +185,10 @@ impl UdfTask for Task { self.ready.register(cx.waker()); let mut s = self.state.lock().unwrap(); match s.out.pop_front() { - Some(b) => Poll::Ready(Ok(Some(b))), + Some(b) => { + self.shared.taken.fetch_add(1, Ordering::SeqCst); + Poll::Ready(Ok(Some(b))) + } None if s.ended => Poll::Ready(Ok(None)), None => Poll::Pending, } @@ -228,7 +246,7 @@ fn rig(shape: RelationShape, yields: Yield, limits: RelationLimits, limit: usize let names = || vec!["k".to_owned(), "v".to_owned()]; let sides = if cogroup { 2 } else { 1 }; let (client, output) = match yields { - Yield::Count => ( + Yield::Count | Yield::CountTwice => ( vec![ArrowType::Int64], match shape { RelationShape::Aggregate { .. } => vec![ @@ -238,7 +256,7 @@ fn rig(shape: RelationShape, yields: Yield, limits: RelationLimits, limit: usize _ => vec![Field::new("n", ArrowType::Int64, true)], }, ), - Yield::Twice | Yield::Nothing => ( + Yield::Twice | Yield::Nothing | Yield::Burst(_) => ( vec![ArrowType::Int32, ArrowType::Int64], vec![ Field::new("k", ArrowType::Int32, true), @@ -247,6 +265,7 @@ fn rig(shape: RelationShape, yields: Yield, limits: RelationLimits, limit: usize ), }; let lane = RelationLane { + name: "f".to_owned(), udf: Arc::new(Fake { shared: Arc::clone(&shared), yields, @@ -583,6 +602,7 @@ 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, @@ -613,9 +633,91 @@ async fn a_result_of_another_width_fails_the_phase() { ) .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] }, + 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); + } +} diff --git a/crates/krishiv-exec/src/shared.rs b/crates/krishiv-exec/src/shared.rs index baf0ccc24..5bc6e7526 100644 --- a/crates/krishiv-exec/src/shared.rs +++ b/crates/krishiv-exec/src/shared.rs @@ -381,6 +381,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 { diff --git a/crates/krishiv-exec/tests/planner/python_relations.rs b/crates/krishiv-exec/tests/planner/python_relations.rs index 6a2baa780..0c72f32b2 100644 --- a/crates/krishiv-exec/tests/planner/python_relations.rs +++ b/crates/krishiv-exec/tests/planner/python_relations.rs @@ -31,6 +31,8 @@ enum Mode { Echo, /// Each batch back twice. Twice, + /// Each batch back eight times. + Burst, /// One `BIGINT` per group (per partition without groups): its rows. Count, } @@ -63,6 +65,7 @@ impl UdfRuntime for Runtime { let mode = match spec.name.as_str() { "echo" => Mode::Echo, "twice" => Mode::Twice, + "burst" => Mode::Burst, _ => Mode::Count, }; assert_eq!(spec.input_names.len(), spec.input_types.len(), "{spec:?}"); @@ -197,10 +200,12 @@ impl UdfTask for Task { 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::Twice | Mode::Burst => { let back = back.unwrap(); - s.out.push_back(back.clone()); - s.out.push_back(back); + let n = if self.mode == Mode::Twice { 2 } else { 8 }; + for _ in 0..n { + s.out.push_back(back.clone()); + } } Mode::Count => {} } @@ -660,26 +665,27 @@ fn a_worker_without_a_runtime_refuses_the_plan() { ); } -#[test] -fn a_skewed_group_runs_through_the_spilling_sort_in_bounded_memory() { +/// 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; - const ROWS: i32 = 300_000; let schema = Arc::new(Schema::new(vec![ Field::new("k", ArrowType::Int32, false), Field::new("v", ArrowType::Int64, false), ])); - // 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 batches: Vec = (0..ROWS) + let batches: Vec = (0..rows) .step_by(8192) .map(|start| { - let end = (start + 8192).min(ROWS); - let k: Int32Array = (start..end) - .map(|i| if i < 270_000 { 0 } else { 1 + i / 100 }) - .collect(); + 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(), @@ -704,13 +710,7 @@ fn a_skewed_group_runs_through_the_spilling_sort_in_bounded_memory() { output: vec![k.clone(), v], version: None, })); - let node = LogicalPlan::new(PlanNode::GroupMap { - input: Arc::new(scan), - keys: vec![col(&k)], - udf: udf("count", PythonEvalType::GroupedMapArrow), - output: counted(), - }); - let plan = Plan::try_new(node).unwrap(); + 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()), @@ -738,6 +738,34 @@ fn a_skewed_group_runs_through_the_spilling_sort_in_bounded_memory() { }); 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| { @@ -753,13 +781,31 @@ fn a_skewed_group_runs_through_the_spilling_sort_in_bounded_memory() { assert_eq!(counts.len(), 301); assert_eq!(counts.last(), Some(&270_000)); assert_eq!(counts.iter().sum::(), i64::from(ROWS)); - let files: u64 = query - .metrics() - .operators - .values() - .map(|o| o.metrics.spill_files) - .sum(); assert!(files > 0, "the sort did not spill"); - drop(query); - assert_eq!(krishiv_memory::MemoryPool::reserved(&*root), 0); +} + +/// 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"); } From 0e4113184390020ffda6c19110507fa569b5ca2b Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Fri, 13 Nov 2026 18:22:07 +0000 Subject: [PATCH 09/29] plan, exec: Python relation review minors (P3-T014a review m1, m2, m4-m6, m8-m10) - Test helper: Blocked took the batch, so the next batch follows (m1); a driver blocked on every batch sends each row once. - A slow fake task blocks the real driver and loses nothing (m2). - A weighted (change-stream) batch is refused (m4). - Only a Python aggregate UDF over a window gets the PR-26c refusal; a scalar one keeps the window planner's own error (m5). - INVALID_PANDAS_UDF_PLACEMENT is raised at physical planning only, as Spark does; analysis accepts the aggregate (m6). - Seam docs: one task per driver; UdfTask takes &self and must not call back into the operator (m8, m10); cache.rs doc comment back on plan (m9). Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-exec/src/planner/build.rs | 7 +- crates/krishiv-exec/src/planner/python_udf.rs | 28 +++++- crates/krishiv-exec/src/python_relation.rs | 4 + .../krishiv-exec/src/python_relation/tests.rs | 45 ++++++++-- .../tests/planner/python_relations.rs | 87 ++++++++++++++++++- .../src/analyzer/ops/aggregate.rs | 2 - crates/krishiv-plan/src/cache.rs | 2 +- crates/krishiv-plan/src/python.rs | 9 +- 8 files changed, 165 insertions(+), 19 deletions(-) diff --git a/crates/krishiv-exec/src/planner/build.rs b/crates/krishiv-exec/src/planner/build.rs index 17241f200..7cf285c58 100644 --- a/crates/krishiv-exec/src/planner/build.rs +++ b/crates/krishiv-exec/src/planner/build.rs @@ -584,9 +584,12 @@ impl<'a> Planner<'a> { partition_by, order_by, } => { - if exprs.iter().any(|e| super::python_udf::has_python_udf(e)) { + if exprs + .iter() + .any(|e| super::python_udf::has_python_aggregate_udf(e)) + { return Err(super::python_relation::later_kind( - "a Python UDF in a window", + "a Python aggregate UDF over a window", )); } let input = self.plan(input)?; diff --git a/crates/krishiv-exec/src/planner/python_udf.rs b/crates/krishiv-exec/src/planner/python_udf.rs index 8a8a72a53..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 } diff --git a/crates/krishiv-exec/src/python_relation.rs b/crates/krishiv-exec/src/python_relation.rs index 59a3a4a0d..00d7f91c6 100644 --- a/crates/krishiv-exec/src/python_relation.rs +++ b/crates/krishiv-exec/src/python_relation.rs @@ -144,6 +144,10 @@ pub trait PreparedRelationUdf: Send + Sync { /// 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. /// diff --git a/crates/krishiv-exec/src/python_relation/tests.rs b/crates/krishiv-exec/src/python_relation/tests.rs index 2339a992a..bbecae69e 100644 --- a/crates/krishiv-exec/src/python_relation/tests.rs +++ b/crates/krishiv-exec/src/python_relation/tests.rs @@ -331,12 +331,12 @@ 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 { - loop { - match r.sink.sink(b.clone(), &mut l, &g)? { - SinkStatus::Blocked(wait) => wait.await, - SinkStatus::NeedMemory { bytes } => panic!("no memory for {bytes}"), - _ => break, - } + // `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)?; @@ -721,3 +721,36 @@ async fn an_aggregate_task_yielding_another_number_of_values_than_groups_fails() 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); +} diff --git a/crates/krishiv-exec/tests/planner/python_relations.rs b/crates/krishiv-exec/tests/planner/python_relations.rs index 0c72f32b2..3fa388370 100644 --- a/crates/krishiv-exec/tests/planner/python_relations.rs +++ b/crates/krishiv-exec/tests/planner/python_relations.rs @@ -46,6 +46,8 @@ struct Runtime { /// 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>>; @@ -62,7 +64,8 @@ impl UdfRuntime for Runtime { &self, spec: &RelationCallSpec, ) -> Result, KrishivError> { - let mode = match spec.name.as_str() { + 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, @@ -84,6 +87,7 @@ impl UdfRuntime for Runtime { 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, })) @@ -97,6 +101,7 @@ struct Prepared { log: Arc>>, per_task: Arc>>, keys: Log, + waits: Option>, names: Vec>, key_count: usize, } @@ -111,6 +116,7 @@ impl PreparedRelationUdf for Prepared { 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()), @@ -133,6 +139,7 @@ struct Task { output: arrow_schema::SchemaRef, log: Arc>>, keys: Log, + waits: Option>, names: Vec>, key_count: usize, state: std::sync::Mutex, @@ -190,6 +197,21 @@ impl Task { 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 } => { @@ -380,6 +402,31 @@ fn a_grouped_map_sees_each_key_once_with_all_its_rows() { 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(); @@ -530,12 +577,21 @@ fn a_python_aggregate_beside_a_builtin_is_refused() { grouping: grouping.clone(), aggregates, }); - let err = run_with(&env, node, Some(Arc::new(Runtime::default()))).unwrap_err(); + // 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 @@ -647,6 +703,33 @@ fn python_udtfs_and_window_aggregates_are_refused_until_pr_26c() { } } +#[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(); diff --git a/crates/krishiv-plan/src/analyzer/ops/aggregate.rs b/crates/krishiv-plan/src/analyzer/ops/aggregate.rs index b8b7a0fa1..5ab1c5cfa 100644 --- a/crates/krishiv-plan/src/analyzer/ops/aggregate.rs +++ b/crates/krishiv-plan/src/analyzer/ops/aggregate.rs @@ -344,8 +344,6 @@ impl Resolver<'_> { } } let items = self.select_items(&c, aggregates, env, "Aggregate")?; - crate::check_python_aggregates(items.iter().map(|b| &b.expr)) - .map_err(|e| AnalysisError::new(e, p.origin))?; // The declared `CHAR`/`VARCHAR` types follow the aliases of an explicit `GROUP BY` too. self.carry_declared(&items, aggregates); // Spark's `ExtractGenerator`: the generator's argument takes the diff --git a/crates/krishiv-plan/src/cache.rs b/crates/krishiv-plan/src/cache.rs index 936f57370..319c889ed 100644 --- a/crates/krishiv-plan/src/cache.rs +++ b/crates/krishiv-plan/src/cache.rs @@ -572,7 +572,6 @@ impl Encoder { } } - /// Writes `plan` as the digest of its encoding, computed once per node. /// 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); @@ -582,6 +581,7 @@ impl Encoder { } } + /// Writes `plan` as the digest of its encoding, computed once per node. fn plan(&mut self, plan: &LogicalPlan, bias: u64) { let key = (NodeKind::Plan, std::ptr::from_ref(plan) as usize, bias); let digest = match self.seen.get(&key) { diff --git a/crates/krishiv-plan/src/python.rs b/crates/krishiv-plan/src/python.rs index 4b0025628..13a340a6d 100644 --- a/crates/krishiv-plan/src/python.rs +++ b/crates/krishiv-plan/src/python.rs @@ -25,8 +25,13 @@ pub fn is_python_udaf(e: &Expr) -> bool { /// 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). Subqueries are not -/// looked into. +/// 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). From 76362376a2bd19958298dc2d091b6314a424775b Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Fri, 13 Nov 2026 20:16:03 +0000 Subject: [PATCH 10/29] exec: a one-sided cogroup key over two batches gets one empty frame (P3-T014a re-review N1) A right-only cogroup key whose rows spanned two input batches got a second zero-row left frame after its first right frames. The empty side now counts as seen when its frame is queued. Regression test for a left-only and a right-only key split over two batches. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../krishiv-exec/src/python_relation/sink.rs | 4 ++- .../krishiv-exec/src/python_relation/tests.rs | 25 +++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/crates/krishiv-exec/src/python_relation/sink.rs b/crates/krishiv-exec/src/python_relation/sink.rs index ec6f1e334..37db90a78 100644 --- a/crates/krishiv-exec/src/python_relation/sink.rs +++ b/crates/krishiv-exec/src/python_relation/sink.rs @@ -456,8 +456,10 @@ impl RelationSink { } /// A zero-row frame of `side`: a cogroup key without rows on that side - /// (Spark gives the UDF an empty DataFrame for it). + /// (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; d.backlog.push_back(Cut::Rows { side, rows: RecordBatch::new_empty(Arc::clone(&self.input)), diff --git a/crates/krishiv-exec/src/python_relation/tests.rs b/crates/krishiv-exec/src/python_relation/tests.rs index bbecae69e..95652079c 100644 --- a/crates/krishiv-exec/src/python_relation/tests.rs +++ b/crates/krishiv-exec/src/python_relation/tests.rs @@ -440,6 +440,31 @@ async fn a_cogroup_key_on_one_side_gets_a_zero_row_frame_for_the_other() { 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( From 798906c827105955be4083435d08be8a8eb6c171 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Fri, 13 Nov 2026 20:35:59 +0000 Subject: [PATCH 11/29] exec: Python relation finalize: wait for the slowest aggregate lane, spill refused results, hand stored bytes over (P3-T014a re-review N2-N5) - N2: `wait` counts an aggregate driver ready only when every lane has values (what `pair::joined` needs), so `finalize` no longer loops inside one poll while one lane is ahead. - N3: in `finalize` a refused map or grouped result is stored with the others (spill file when the pool is full) instead of waiting for memory; an aggregate's refused value is reserved with `try_grow` after each step and otherwise waits for its memory while the other lanes keep moving. A finished aggregate with values left over is `check_complete`'s error rather than another loop. - N4: a stored result's reservation passes to the store's reservation through the consumer's grant credit (`lend`), never counted twice. - N5: tests for two aggregate lanes at different speeds (unit, bounded by a poll count; planner, values paired per key) and for a result refused in `finalize` under a small memory limit. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../krishiv-exec/src/python_relation/lane.rs | 41 +++--- .../krishiv-exec/src/python_relation/sink.rs | 111 +++++++++++++--- .../krishiv-exec/src/python_relation/tests.rs | 122 +++++++++++++++++- .../tests/planner/python_relations.rs | 31 ++++- 4 files changed, 267 insertions(+), 38 deletions(-) diff --git a/crates/krishiv-exec/src/python_relation/lane.rs b/crates/krishiv-exec/src/python_relation/lane.rs index 3d4e36cb8..5449c28eb 100644 --- a/crates/krishiv-exec/src/python_relation/lane.rs +++ b/crates/krishiv-exec/src/python_relation/lane.rs @@ -9,9 +9,11 @@ //! 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 reserved it -//! (`sink` returns `NeedMemory`, `finalize` awaits the memory): back -//! pressure into the task, never unaccounted results. +//! 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. use std::collections::VecDeque; use std::future::poll_fn; @@ -149,22 +151,29 @@ impl Lane { None } - /// Reserves the refused result, if any, waiting for the memory, and - /// moves it to the outputs. - pub(super) async fn take_refused(&self) -> Result<()> { + /// 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 { + lock(&self.io).refused.take() + } + + /// 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 Some(batch) = &q.refused else { - return Ok(()); - }; - (batch.get_array_memory_size(), q.memory.split(0)) + let bytes = q.refused.as_ref()?.get_array_memory_size(); + (bytes, q.memory.split(0)) }; - r.grow(bytes).await.map_err(KrishivError::from)?; - let mut q = lock(&self.io); - if let Some(batch) = q.refused.take() { - q.outputs.push_back((batch, r)); - } - Ok(()) + 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(()) + })) } } diff --git a/crates/krishiv-exec/src/python_relation/sink.rs b/crates/krishiv-exec/src/python_relation/sink.rs index 37db90a78..e3535bd61 100644 --- a/crates/krishiv-exec/src/python_relation/sink.rs +++ b/crates/krishiv-exec/src/python_relation/sink.rs @@ -38,7 +38,11 @@ //! 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; `finalize` waits for the memory. +//! 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. 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, @@ -56,7 +60,7 @@ 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::{Reservation, SpillScope}; +use krishiv_memory::{ConsumerHandle, Reservation, SpillScope}; use krishiv_types::batch::Batch; use super::lane::{FrameMem, Lane}; @@ -305,6 +309,8 @@ enum Flushed { 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. @@ -622,8 +628,9 @@ impl RelationSink { Ok(None) } - /// Stores `batch` in the list, or keeps it for the retry after the - /// memory the store asked for. + /// 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, @@ -631,10 +638,16 @@ impl RelationSink { 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 } => { - l.driver.ready = Some((batch, memory)); - Ok(Some(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), } @@ -741,6 +754,25 @@ pub(super) fn recarried( 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 +} + /// The future of a `Blocked` driver: ready once every lane's window has /// room, a lane ended (its failure is the next call's), or a result waits /// for memory (the next call asks for it). @@ -773,6 +805,7 @@ impl Sink for RelationSink { 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: self @@ -835,7 +868,9 @@ impl Sink for RelationSink { let l = local .downcast::() .map_err(|_| internal("wrong local state for PythonRelation".to_owned()))?; - let RelLocal { store, mut driver } = *l; + 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 { @@ -898,9 +933,11 @@ impl Sink for RelationSink { impl RelationSink { /// One step of a finished driver in `finalize`: the backlog framed as /// far as the windows allow (waiting for the copies' memory), the - /// results stored (written to a spill file when the pool has no room), - /// then refused results reserved. Whether the driver - /// is done: no backlog, every task ended, every result stored. + /// results stored, a refused map or grouped result with them (written + /// to a spill file when the pool has no room), then an aggregate's + /// refused values reserved if the pool has room now. 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 { while let Flushed::Memory(bytes) = self.flush(d)? { d.spare.grow(bytes).await.map_err(KrishivError::from)?; @@ -915,6 +952,15 @@ impl RelationSink { 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() @@ -922,48 +968,73 @@ impl RelationSink { .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(); if reservation.try_grow(bytes).is_ok() { self.store.add_part(&g.store, batches, reservation)?; } else { - drop(reservation); + drop((reservation, handle.reservation_from_credit(lent))); self.store.add_spilled_part(&g.store, batches, ctx).await?; } } - drop(held); - // The results' memory is free again: a refused result next. - for lane in &d.lanes { - lane.take_refused().await?; + 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(d.backlog.is_empty() && d.ready.is_none() && d.lanes .iter() - .all(|l| l.finished() && !lock(&l.io).has_results())) + .all(|l| l.finished() && (aggregate || !lock(&l.io).has_results()))) } - /// Waits until a driver can move: a lane yielded, failed, ended or - /// freed room in its window for a backlog; or the query was cancelled. + /// 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); - ready |= q.has_results(); + if aggregate { + joinable &= !q.outputs.is_empty(); + } else { + ready |= q.has_results(); + } room &= q.in_flight() < self.window; } - ready |= room && !d.backlog.is_empty(); + ready |= joinable || (room && !d.backlog.is_empty()); all_done &= d.backlog.is_empty() && d.ready.is_none(); } if ready || all_done { diff --git a/crates/krishiv-exec/src/python_relation/tests.rs b/crates/krishiv-exec/src/python_relation/tests.rs index 95652079c..24ae5ab07 100644 --- a/crates/krishiv-exec/src/python_relation/tests.rs +++ b/crates/krishiv-exec/src/python_relation/tests.rs @@ -47,8 +47,13 @@ struct Shared { 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, } +/// 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); @@ -182,6 +187,12 @@ impl UdfTask for Task { 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() { @@ -239,6 +250,17 @@ struct Rig { } fn rig(shape: RelationShape, yields: Yield, limits: RelationLimits, limit: usize) -> Rig { + rig_in(shape, yields, limits, limit, None) +} + +/// 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); @@ -280,7 +302,7 @@ fn rig(shape: RelationShape, yields: Yield, limits: RelationLimits, limit: usize vec![lane], Arc::new(Schema::new(output)), limits, - None, + scope, ) .unwrap(); let root = pool(limit); @@ -779,3 +801,101 @@ async fn a_weighted_batch_is_refused() { 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] }, + 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] }, + 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); +} + +#[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); +} diff --git a/crates/krishiv-exec/tests/planner/python_relations.rs b/crates/krishiv-exec/tests/planner/python_relations.rs index 3fa388370..8baf5a597 100644 --- a/crates/krishiv-exec/tests/planner/python_relations.rs +++ b/crates/krishiv-exec/tests/planner/python_relations.rs @@ -518,9 +518,13 @@ fn a_computed_grouping_key_leads_every_frame() { /// 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("count", PythonEvalType::GroupedAggArrow).function_id(), + id: udf(name, PythonEvalType::GroupedAggArrow).function_id(), kind: FunctionKind::Aggregate, args: vec![col(arg)], distinct: false, @@ -556,6 +560,31 @@ fn a_python_aggregate_yields_one_row_per_group() { assert!(rows.contains(&"5|1".to_owned()), "{rows:?}"); } +/// 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(); From 329835d24a06679abd4fa75d8bfe9008b36f5f51 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Fri, 13 Nov 2026 20:39:15 +0000 Subject: [PATCH 12/29] exec: frame column names need not be unique (P3-T014a re-review N6) A column key is also a data column, so `groupBy("a")` frames are `a, a, b`: the `RelationCallSpec::input_names` doc says a runtime splits by key_count, never by name. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-exec/src/python_relation.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/crates/krishiv-exec/src/python_relation.rs b/crates/krishiv-exec/src/python_relation.rs index 00d7f91c6..0c9b4d9c6 100644 --- a/crates/krishiv-exec/src/python_relation.rs +++ b/crates/krishiv-exec/src/python_relation.rs @@ -77,7 +77,9 @@ pub struct RelationCallSpec { 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. + /// 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`] From 5e74d8dde8a17d6368bdcd8e4b1306a9048422b4 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 01:25:52 +0000 Subject: [PATCH 13/29] exec: an aggregate lane holding only empty batches is not joinable (P3-T014a re-review R1) wait() counted a lane with any queued result as joinable, but joined() needs rows; a lane holding only 0-row batches made finalize spin. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-exec/src/python_relation/sink.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/crates/krishiv-exec/src/python_relation/sink.rs b/crates/krishiv-exec/src/python_relation/sink.rs index e3535bd61..adb97c3a5 100644 --- a/crates/krishiv-exec/src/python_relation/sink.rs +++ b/crates/krishiv-exec/src/python_relation/sink.rs @@ -1028,7 +1028,9 @@ impl RelationSink { all_done &= lane.finished(); let q = lock(&lane.io); if aggregate { - joinable &= !q.outputs.is_empty(); + // 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(); } From 1a573cab2cbffe6ef0267eb238f5a1067000fa8a Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 02:22:56 +0000 Subject: [PATCH 14/29] changelog: fragment for the Python UDF relation IR (#746) Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- changelog.d/746.added.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changelog.d/746.added.md 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. From 964e9c08cde6adf80e44077081f07139a8092eee Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 04:15:37 +0000 Subject: [PATCH 15/29] exec: resolve the prepare_relation doc links of python_relation (#746 CI docs) The module doc linked UdfRuntime::prepare_relation, which is not in scope there; rustdoc -D warnings failed the docs job. A reference link to crate::python_eval::UdfRuntime::prepare_relation resolves both uses. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-exec/src/python_relation.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/crates/krishiv-exec/src/python_relation.rs b/crates/krishiv-exec/src/python_relation.rs index 0c9b4d9c6..87499b3f7 100644 --- a/crates/krishiv-exec/src/python_relation.rs +++ b/crates/krishiv-exec/src/python_relation.rs @@ -39,6 +39,8 @@ //! `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; From 9a079735d0eba7e6361def655e962318d04f9050 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 05:26:28 +0000 Subject: [PATCH 16/29] exec: scalar Python calls below and above a Python aggregate (#746 review) A Python aggregate's keys and call arguments went to the native compiler with their scalar Python calls in them, and so did a scalar call over a call's result: GROUP BY f(v), udaf(f(v)) and f(udaf(v)) failed at planning. As the ordinary aggregate path does, and as Spark does (ExtractGroupingPythonUDFFromAggregate, ExtractPythonUDFs and ExtractPythonUDFFromAggregate at v4.2.0), the calls in keys and arguments are extracted into PythonEval nodes below the node, which runs each call over their columns, and the calls over its results into nodes above it. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../src/planner/python_relation.rs | 66 ++++++++--- .../tests/planner/python_relations.rs | 108 ++++++++++++++++++ 2 files changed, 159 insertions(+), 15 deletions(-) diff --git a/crates/krishiv-exec/src/planner/python_relation.rs b/crates/krishiv-exec/src/planner/python_relation.rs index eca95a866..305c69c21 100644 --- a/crates/krishiv-exec/src/planner/python_relation.rs +++ b/crates/krishiv-exec/src/planner/python_relation.rs @@ -338,24 +338,51 @@ impl Planner<'_> { for a in aggregates { udaf_calls(a, &mut calls, &mut seen); } - let planned = self.plan(input)?; - let (planned, key_columns) = self.with_columns(planned, keys)?; - let prepared = self.prepare(&[])?; for call in &calls { - let ExprNode::Resolved(ResolvedExpr::Function(f)) = &call.node else { - return Err(internal("a Python UDAF call is not a function")); - }; + 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")); } - let args: Vec<&Expr> = f.args.iter().map(AsRef::as_ref).collect(); + } + // 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)?; + 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( - &args, + &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, &key_columns, &[])?; let mut layout: Vec = key_columns @@ -368,10 +395,8 @@ impl Planner<'_> { .collect(); let mut plans = Vec::with_capacity(calls.len()); let mut columns: Vec<(Arc, BoundColumn)> = Vec::with_capacity(calls.len()); - for (i, call) in calls.iter().enumerate() { - let ExprNode::Resolved(ResolvedExpr::Function(f)) = &call.node else { - return Err(internal("a Python UDAF call is not a function")); - }; + 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}"), @@ -381,12 +406,12 @@ impl Planner<'_> { fields.push(Field::new(&column.name, f.data_type.to_arrow(), true)); layout.push(column.id); plans.push(PythonCallPlan { - call: Arc::clone(call), + call: Arc::clone(ran), output: column.id, }); columns.push((Arc::clone(call), column)); } - let detail = calls + let detail = run .iter() .zip(&columns) .map(|(c, (_, col))| format!("{c} AS {}", col.name)) @@ -428,7 +453,9 @@ impl Planner<'_> { .iter() .map(|e| replaced(e, &columns, &key_bound)) .collect::>>()?; - let mut out = self.project( + // 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(), @@ -437,6 +464,7 @@ impl Planner<'_> { }, &exprs, )?; + let mut out = self.project(node, &exprs)?; out.sites.clear(); Ok(out) } @@ -527,6 +555,14 @@ fn refuse_call_syntax(f: &krishiv_plan::ResolvedFunction) -> Result<()> { ))) } +/// 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 { diff --git a/crates/krishiv-exec/tests/planner/python_relations.rs b/crates/krishiv-exec/tests/planner/python_relations.rs index 8baf5a597..7394c8a03 100644 --- a/crates/krishiv-exec/tests/planner/python_relations.rs +++ b/crates/krishiv-exec/tests/planner/python_relations.rs @@ -54,6 +54,9 @@ 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), } @@ -94,6 +97,18 @@ impl UdfRuntime for Runtime { } } +/// 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, @@ -560,6 +575,99 @@ fn a_python_aggregate_yields_one_row_per_group() { assert!(rows.contains(&"5|1".to_owned()), "{rows:?}"); } +/// 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:?}"); +} + /// Two Python aggregates, one slow: each key's values pair up. #[test] fn two_python_aggregates_pair_their_values_per_key() { From dd7477ee5f388321a3d481a2faa841ec96819b3a Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 05:29:54 +0000 Subject: [PATCH 17/29] exec, optimizer: group and cogroup Python relations by a collated key's collation (#746 review) GroupMap and CoGroupMap passed collated string keys as they are to the exchange, the sort and the group boundaries, which compare bytes: under UTF8_LCASE, 'A' and 'a' were two groups (two one-sided cogroup calls). Spark groups and cogroups by the collation (GroupedIterator and CoGroupedIterator orderings at v4.2.0). The planner now appends each collated key's comparison key, the one the collation rewrite gives aggregates, sorts and exchanges (krishiv_optimizer::collation_key, now public), and clusters, sorts and cuts groups on it; the frames keep each row's own key values. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../src/planner/python_relation.rs | 80 ++++++++++++++++--- .../tests/planner/python_relations.rs | 68 +++++++++++++++- .../krishiv-optimizer/src/collation_keys.rs | 8 +- crates/krishiv-optimizer/src/lib.rs | 1 + 4 files changed, 146 insertions(+), 11 deletions(-) diff --git a/crates/krishiv-exec/src/planner/python_relation.rs b/crates/krishiv-exec/src/planner/python_relation.rs index 305c69c21..b5846e3af 100644 --- a/crates/krishiv-exec/src/planner/python_relation.rs +++ b/crates/krishiv-exec/src/planner/python_relation.rs @@ -165,11 +165,12 @@ impl Planner<'_> { let cols = columns_of(input)?; let planned = self.plan(input)?; let (planned, key_columns) = self.with_columns(planned, keys)?; + let (planned, compared) = self.compared(planned, keys, &key_columns)?; let columns = cols .iter() .map(|c| planned.column(c.id)) .collect::>>()?; - let sorted = self.clustered_and_sorted(planned, &key_columns, &[])?; + let sorted = self.clustered_and_sorted(planned, &compared, &[])?; let plan = PythonRelationPlan { udf: udf.clone(), sides: vec![side_of(keys, &key_columns, &columns, &cols)?], @@ -185,7 +186,7 @@ impl Planner<'_> { sorted, |input| PhysNode::PythonGroupMap { input, - keys: key_columns, + keys: compared, udf: plan, }, output, @@ -210,16 +211,23 @@ impl Planner<'_> { let (left_cols, right_cols) = (columns_of(left)?, columns_of(right)?); let l = self.plan(left)?; let (l, lk) = self.with_columns(l, left_keys)?; + let (l, lc) = self.compared(l, left_keys, &lk)?; let r = self.plan(right)?; let (r, rk) = self.with_columns(r, right_keys)?; + let (r, rc) = self.compared(r, right_keys, &rk)?; let k = lk.len(); - // One layout for both branches: keys, side, left, right. + // 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 layout: Vec = (0..width).map(|_| self.fresh_id()).collect(); - let nullable = vec![true; width]; + 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)) in [(l, &lk), (r, &rk)].into_iter().enumerate() { - let mut items = Vec::with_capacity(width); + 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, @@ -249,6 +257,14 @@ impl Planner<'_> { } } } + 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, @@ -291,7 +307,12 @@ impl Planner<'_> { placement: Placement::Unknown, }; let key_columns: Vec = (0..k).collect(); - let sorted = self.clustered_and_sorted(union, &key_columns, &[k])?; + 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![ @@ -315,7 +336,7 @@ impl Planner<'_> { sorted, |input| PhysNode::PythonCoGroupMap { input, - keys: key_columns, + keys: compared, side: k, udf: plan, }, @@ -469,6 +490,47 @@ impl Planner<'_> { 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( diff --git a/crates/krishiv-exec/tests/planner/python_relations.rs b/crates/krishiv-exec/tests/planner/python_relations.rs index 7394c8a03..f9324ecc7 100644 --- a/crates/krishiv-exec/tests/planner/python_relations.rs +++ b/crates/krishiv-exec/tests/planner/python_relations.rs @@ -186,8 +186,14 @@ impl Task { } 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), first, "one key per group"); + assert_eq!( + key(r).to_lowercase(), + first.to_lowercase(), + "one key per group" + ); } self.keys .lock() @@ -531,6 +537,66 @@ fn a_computed_grouping_key_leads_every_frame() { 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:?}"); +} + /// A grouped-aggregate Python UDF call over `arg`. fn udaf(arg: &BoundColumn) -> Arc { udaf_named("count", arg) diff --git a/crates/krishiv-optimizer/src/collation_keys.rs b/crates/krishiv-optimizer/src/collation_keys.rs index 562d40bbd..1d2cbc651 100644 --- a/crates/krishiv-optimizer/src/collation_keys.rs +++ b/crates/krishiv-optimizer/src/collation_keys.rs @@ -91,7 +91,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 { diff --git a/crates/krishiv-optimizer/src/lib.rs b/crates/krishiv-optimizer/src/lib.rs index 8a361fbf3..af0708a12 100644 --- a/crates/krishiv-optimizer/src/lib.rs +++ b/crates/krishiv-optimizer/src/lib.rs @@ -87,6 +87,7 @@ use krishiv_plan::{ }; use krishiv_udf_sdk::UdfLibrary; +pub use collation_keys::key_if_collated as collation_key; pub use group_limit::{ DEFAULT_WINDOW_GROUP_LIMIT_THRESHOLD, RankLikeFunction, WindowGroupLimit, infer_window_group_limit, From 9d777f3f7e7e3c9e3c1197ac9684ac992a8d5dd4 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 05:30:41 +0000 Subject: [PATCH 18/29] exec: Python relation finalize stores results before waiting for frame memory (#746 review) In finalize a frame whose copies the pool refused was awaited first, while the lanes' results held reservations nothing could reclaim (the lane queues are not spillable, and the store's parts are not until it finalizes): the wait could only time out. A driver's results are now stored before it waits, written out while a frame waits for memory, and the frame tries again; the wait for what is still refused comes after every driver stored its results. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../krishiv-exec/src/python_relation/sink.rs | 41 +++++-- .../krishiv-exec/src/python_relation/tests.rs | 116 ++++++++++++++++++ crates/krishiv-exec/src/shared.rs | 6 + 3 files changed, 155 insertions(+), 8 deletions(-) diff --git a/crates/krishiv-exec/src/python_relation/sink.rs b/crates/krishiv-exec/src/python_relation/sink.rs index adb97c3a5..b3499f61e 100644 --- a/crates/krishiv-exec/src/python_relation/sink.rs +++ b/crates/krishiv-exec/src/python_relation/sink.rs @@ -291,11 +291,14 @@ struct Driver { /// 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, } /// What [`RelationSink::flush`] stopped at. +#[derive(Clone, Copy)] enum Flushed { /// The backlog is empty. All, @@ -824,6 +827,7 @@ impl Sink for RelationSink { held_bytes: consumer.reservation(), backlog: VecDeque::new(), spare: consumer.reservation(), + wants: 0, pending: Pending::default(), }, })) @@ -905,6 +909,14 @@ impl Sink for RelationSink { 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 { .. }) { @@ -932,16 +944,17 @@ impl Sink for RelationSink { impl RelationSink { /// One step of a finished driver in `finalize`: the backlog framed as - /// far as the windows allow (waiting for the copies' memory), the - /// results stored, a refused map or grouped result with them (written - /// to a spill file when the pool has no room), then an aggregate's - /// refused values reserved if the pool has room now. Whether the driver + /// 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 { - while let Flushed::Memory(bytes) = self.flush(d)? { - d.spare.grow(bytes).await.map_err(KrishivError::from)?; - } + d.wants = wanted(self.flush(d)?); let mut batches = Vec::new(); let mut held = Vec::new(); if let Some((batch, memory)) = d.ready.take() { @@ -971,13 +984,17 @@ impl RelationSink { let handle = ctx.consumer().consumer(); let lent = lend(std::mem::take(&mut held), handle, bytes); let mut reservation = ctx.consumer().reservation(); - if reservation.try_grow(bytes).is_ok() { + 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. @@ -1049,6 +1066,14 @@ impl RelationSink { } } +/// 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)] diff --git a/crates/krishiv-exec/src/python_relation/tests.rs b/crates/krishiv-exec/src/python_relation/tests.rs index 24ae5ab07..83a4d4030 100644 --- a/crates/krishiv-exec/src/python_relation/tests.rs +++ b/crates/krishiv-exec/src/python_relation/tests.rs @@ -899,3 +899,119 @@ async fn a_result_refused_in_finalize_spills_instead_of_waiting() { 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); +} diff --git a/crates/krishiv-exec/src/shared.rs b/crates/krishiv-exec/src/shared.rs index 5bc6e7526..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`. From 35f2ef8be569feb48eff6c2acf203164db520dcc Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 05:33:04 +0000 Subject: [PATCH 19/29] exec: a Python relation's cuts share the input reservation by all its rows (#746 review) take_input reserves a whole batch, but cut_held split that reservation over the held rows only: the last cut of a batch's first group took the whole batch's bytes, so its later groups' slices, which keep the batch's buffers, held none once the first group's frames were taken. The driver now counts the rows the reservation covers (taken, not cut yet) and each cut takes its share of all of them. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../krishiv-exec/src/python_relation/sink.rs | 12 +++++- .../krishiv-exec/src/python_relation/tests.rs | 42 +++++++++++++++++++ 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/crates/krishiv-exec/src/python_relation/sink.rs b/crates/krishiv-exec/src/python_relation/sink.rs index b3499f61e..87a83aded 100644 --- a/crates/krishiv-exec/src/python_relation/sink.rs +++ b/crates/krishiv-exec/src/python_relation/sink.rs @@ -284,8 +284,13 @@ struct Driver { held_side: u8, /// The cogroup sides the open group has rows on. sides_seen: [bool; 2], - /// The reservation of the held rows' bytes. + /// 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. backlog: VecDeque, /// Memory reserved for the next frame's copies (`finalize` waits for @@ -453,7 +458,8 @@ impl RelationSink { d.held_rows = len; break; } - let share = d.held_bytes.size() * len / (rows - offset); + let share = d.held_bytes.size() * len / d.held_bytes_rows.max(len); + d.held_bytes_rows = d.held_bytes_rows.saturating_sub(len); d.backlog.push_back(Cut::Rows { side: d.held_side, rows: data.slice(offset, len), @@ -726,6 +732,7 @@ impl RelationSink { 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) @@ -825,6 +832,7 @@ impl Sink for RelationSink { held_side: 0, sides_seen: [false; 2], held_bytes: consumer.reservation(), + held_bytes_rows: 0, backlog: VecDeque::new(), spare: consumer.reservation(), wants: 0, diff --git a/crates/krishiv-exec/src/python_relation/tests.rs b/crates/krishiv-exec/src/python_relation/tests.rs index 83a4d4030..f42671730 100644 --- a/crates/krishiv-exec/src/python_relation/tests.rs +++ b/crates/krishiv-exec/src/python_relation/tests.rs @@ -1015,3 +1015,45 @@ async fn finalize_drains_results_before_waiting_for_frame_copies() { 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)); + r.shared.open_gate(); + tokio::time::timeout(Duration::from_secs(5), wait) + .await + .unwrap(); + // 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!( + r.root.reserved() >= bytes / 2, + "{} of {bytes} bytes reserved", + r.root.reserved() + ); + 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); +} From b0ce59723faa1e7cc5f507e77ddb2d7aee6eb248 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 05:33:32 +0000 Subject: [PATCH 20/29] exec: pin that a keyless Python aggregate over no rows yields no row (#746 review) Spark's ArrowAggregatePythonExec (L156-157 at v4.2.0) returns an empty partition as it is, without calling the UDF, so a keyless grouped-agg Python UDF over empty input yields zero rows, unlike a built-in global aggregate. The test pins that the relation sink does the same. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../tests/planner/python_relations.rs | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/crates/krishiv-exec/tests/planner/python_relations.rs b/crates/krishiv-exec/tests/planner/python_relations.rs index f9324ecc7..0af570a2a 100644 --- a/crates/krishiv-exec/tests/planner/python_relations.rs +++ b/crates/krishiv-exec/tests/planner/python_relations.rs @@ -734,6 +734,25 @@ fn a_scalar_python_call_over_a_python_aggregate_runs_above_it() { 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() { From ccc6be89500223aec9eaa7cf5acc0ead06b517f1 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 07:05:14 +0000 Subject: [PATCH 21/29] exec: scalar Python calls in grouped and cogrouped map keys run below (#746 review) GroupMap and CoGroupMap passed their keys to the native compiler with their scalar Python calls in them: groupBy(plus1(b)).applyInPandas(...) and a cogroup keyed by such a call failed at planning. Spark computes the keys in a Project under FlatMapGroupsInPandas / FlatMapCoGroupsInPandas, which ExtractPythonUDFs at v4.2.0 turns into an ArrowEvalPython. The planner now extracts the calls into a PythonEval below the node (each cogroup side below its own branch), as the Python aggregate path does, and computes and compares the keys over their columns; the frames keep the keys' own names and types. Review thread PRRT_kwDOVMT3-s63aD-C. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../src/planner/python_relation.rs | 24 ++++++--- .../tests/planner/python_relations.rs | 53 +++++++++++++++++++ 2 files changed, 69 insertions(+), 8 deletions(-) diff --git a/crates/krishiv-exec/src/planner/python_relation.rs b/crates/krishiv-exec/src/planner/python_relation.rs index b5846e3af..01ed3627f 100644 --- a/crates/krishiv-exec/src/planner/python_relation.rs +++ b/crates/krishiv-exec/src/planner/python_relation.rs @@ -153,7 +153,11 @@ impl Planner<'_> { )) } - /// `groupBy(keys).applyInPandas` / `applyInArrow`. + /// `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, @@ -164,8 +168,9 @@ impl Planner<'_> { check_kind(udf, PythonUdfKind::GroupedMap, "GroupMap")?; let cols = columns_of(input)?; let planned = self.plan(input)?; - let (planned, key_columns) = self.with_columns(planned, keys)?; - let (planned, compared) = self.compared(planned, keys, &key_columns)?; + 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)) @@ -194,7 +199,8 @@ impl Planner<'_> { )) } - /// `cogroup(...).applyInPandas` / `applyInArrow`. + /// `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, @@ -210,11 +216,13 @@ impl Planner<'_> { } let (left_cols, right_cols) = (columns_of(left)?, columns_of(right)?); let l = self.plan(left)?; - let (l, lk) = self.with_columns(l, left_keys)?; - let (l, lc) = self.compared(l, left_keys, &lk)?; + 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, rk) = self.with_columns(r, right_keys)?; - let (r, rc) = self.compared(r, right_keys, &rk)?; + 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. diff --git a/crates/krishiv-exec/tests/planner/python_relations.rs b/crates/krishiv-exec/tests/planner/python_relations.rs index 0af570a2a..2616f9a39 100644 --- a/crates/krishiv-exec/tests/planner/python_relations.rs +++ b/crates/krishiv-exec/tests/planner/python_relations.rs @@ -597,6 +597,59 @@ fn a_cogroup_matches_collated_keys_by_their_collation() { 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) From 1118b63a40105d777d96688f5a8bc72051637a33 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 07:17:44 +0000 Subject: [PATCH 22/29] exec, optimizer, plan-codec: a Python aggregate groups a collated key by its collation (#746 review) SELECT k, udaf(v) ... GROUP BY k with k under UTF8_LCASE was refused with INVALID_PANDAS_UDF_PLACEMENT: the collation rewrite grouped the aggregate by __collation_key(k) and showed k as any_value(k), a native aggregate beside the Python UDAF. Spark's ArrowAggregatePythonExec (v4.2.0) sorts by the grouping keys and cuts groups with a GroupedIterator comparing them by the collation; each group's key is its first row's. The rewrite now leaves an aggregate of grouped-aggregate Python UDFs alone, and python_aggregate keys it as GroupMap and CoGroupMap do: it appends a collated key's comparison key, clusters, sorts and cuts groups on it, and shows the first row's own key. PythonAggregate gains key_values (the key columns shown; keys are the compared columns), and PythonAggregateNode the field key_values = 5 (python_udf_relations, new in this PR). Review thread PRRT_kwDOVMT3-s63aD-J. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../physical/instantiate/python_relation.rs | 19 +++++-- crates/krishiv-exec/src/physical/mod.rs | 9 ++-- .../src/planner/python_relation.rs | 11 ++-- .../krishiv-exec/src/python_relation/sink.rs | 27 ++++++---- .../krishiv-exec/src/python_relation/tests.rs | 53 +++++++++++++++++-- .../tests/planner/python_relations.rs | 29 ++++++++++ .../krishiv-optimizer/src/collation_keys.rs | 24 +++++++-- crates/krishiv-plan-codec/src/decode.rs | 1 + crates/krishiv-plan-codec/src/encode.rs | 2 + .../src/proto/krishiv.plan.v1.pb.rs | 5 +- .../tests/common/inventory.rs | 1 + proto/krishiv/plan/v1/plan.proto | 4 +- 12 files changed, 155 insertions(+), 30 deletions(-) diff --git a/crates/krishiv-exec/src/physical/instantiate/python_relation.rs b/crates/krishiv-exec/src/physical/instantiate/python_relation.rs index b42ccd38f..1382012f8 100644 --- a/crates/krishiv-exec/src/physical/instantiate/python_relation.rs +++ b/crates/krishiv-exec/src/physical/instantiate/python_relation.rs @@ -42,9 +42,18 @@ impl Instantiation<'_> { side: *side, }, ), - PhysNode::PythonAggregate { input, keys, .. } => { - (input, RelationShape::Aggregate { keys: keys.clone() }) - } + 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)?; @@ -63,10 +72,10 @@ impl Instantiation<'_> { } PhysNode::PythonAggregate { layout, - keys, + key_values, calls, .. - } => self.aggregate_lanes(&input.schema, layout, keys, calls, &conf)?, + } => 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( diff --git a/crates/krishiv-exec/src/physical/mod.rs b/crates/krishiv-exec/src/physical/mod.rs index ccf2a5df6..8102e03e6 100644 --- a/crates/krishiv-exec/src/physical/mod.rs +++ b/crates/krishiv-exec/src/physical/mod.rs @@ -455,14 +455,17 @@ pub enum PhysNode { }, /// An aggregate whose calls are all grouped-aggregate Python UDFs, over /// `input` clustered by the key columns and sorted by them: per group, - /// the keys, then one value per call. + /// 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 key columns of `input`. + /// 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, }, @@ -1071,7 +1074,7 @@ plan_data!(enum PhysNode { PythonMap { input, udf }, PythonGroupMap { input, keys, udf }, PythonCoGroupMap { input, keys, side, udf }, - PythonAggregate { input, layout, keys, calls }, + PythonAggregate { input, layout, keys, key_values, calls }, Generate { input, layout, kind, outer, arg, subqueries }, Expand { input, layout, projections, nullable, subqueries }, Sort { input, keys }, diff --git a/crates/krishiv-exec/src/planner/python_relation.rs b/crates/krishiv-exec/src/planner/python_relation.rs index 01ed3627f..05fd9962f 100644 --- a/crates/krishiv-exec/src/planner/python_relation.rs +++ b/crates/krishiv-exec/src/planner/python_relation.rs @@ -387,6 +387,10 @@ impl Planner<'_> { .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()); @@ -413,7 +417,7 @@ impl Planner<'_> { .with_origin(call.origin), )); } - let sorted = self.clustered_and_sorted(planned, &key_columns, &[])?; + let sorted = self.clustered_and_sorted(planned, &compared, &[])?; let mut layout: Vec = key_columns .iter() .map(|&c| sorted.column_attr(c)) @@ -447,13 +451,14 @@ impl Planner<'_> { .collect::>() .join(", "); let input_layout = sorted.layout.clone(); - let keys_at = key_columns.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: keys_at, + keys: compared, + key_values, calls: plans, }, schema, diff --git a/crates/krishiv-exec/src/python_relation/sink.rs b/crates/krishiv-exec/src/python_relation/sink.rs index 87a83aded..0ac77e040 100644 --- a/crates/krishiv-exec/src/python_relation/sink.rs +++ b/crates/krishiv-exec/src/python_relation/sink.rs @@ -94,10 +94,14 @@ pub enum RelationShape { side: usize, }, /// As [`Self::Grouped`]; each lane's task yields one value per group, - /// and the output is the group's keys then one column per lane. + /// and the output is the values of `key_values` in the group's first + /// row, then one column per lane. Aggregate { - /// The key columns. + /// 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, }, } @@ -105,9 +109,9 @@ impl RelationShape { fn keys(&self) -> Option<&[usize]> { match self { Self::Map => None, - Self::Grouped { keys } | Self::CoGrouped { keys, .. } | Self::Aggregate { keys } => { - Some(keys) - } + Self::Grouped { keys } + | Self::CoGrouped { keys, .. } + | Self::Aggregate { keys, .. } => Some(keys), } } } @@ -218,8 +222,13 @@ impl RelationSink { None => None, }; let key_schema = match &shape { - RelationShape::Aggregate { keys } => { - let fields: Vec = keys + 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 @@ -712,14 +721,14 @@ impl RelationSink { let data = input.data(); let (runs, last) = self.runs(data, d.current.as_ref())?; let keys = match (&self.shape, &self.key_schema) { - (RelationShape::Aggregate { keys }, Some(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, keys, schema, &starts)? + pair::key_rows(data, key_values, schema, &starts)? } _ => None, }; diff --git a/crates/krishiv-exec/src/python_relation/tests.rs b/crates/krishiv-exec/src/python_relation/tests.rs index f42671730..42cc8cf79 100644 --- a/crates/krishiv-exec/src/python_relation/tests.rs +++ b/crates/krishiv-exec/src/python_relation/tests.rs @@ -490,7 +490,10 @@ async fn a_one_sided_cogroup_key_over_two_batches_gets_one_empty_frame_first() { #[tokio::test] async fn an_aggregate_lane_gets_one_group_per_key() { let r = rig( - RelationShape::Aggregate { keys: vec![0] }, + RelationShape::Aggregate { + keys: vec![0], + key_values: vec![0], + }, Yield::Count, limits(0, 16), 64 << 20, @@ -500,6 +503,39 @@ async fn an_aggregate_lane_gets_one_group_per_key() { 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); @@ -752,7 +788,10 @@ 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] }, + RelationShape::Aggregate { + keys: vec![0], + key_values: vec![0], + }, yields, limits(0, 16), 64 << 20, @@ -805,7 +844,10 @@ async fn a_weighted_batch_is_refused() { #[tokio::test] async fn an_aggregate_waits_for_its_slowest_lane_without_spinning() { let mut r = rig( - RelationShape::Aggregate { keys: vec![0] }, + RelationShape::Aggregate { + keys: vec![0], + key_values: vec![0], + }, Yield::Count, limits(0, 16), 64 << 20, @@ -827,7 +869,10 @@ async fn an_aggregate_waits_for_its_slowest_lane_without_spinning() { }; r.sink = RelationSink::try_new( input(false), - RelationShape::Aggregate { keys: vec![0] }, + 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), diff --git a/crates/krishiv-exec/tests/planner/python_relations.rs b/crates/krishiv-exec/tests/planner/python_relations.rs index 2616f9a39..f3ff94058 100644 --- a/crates/krishiv-exec/tests/planner/python_relations.rs +++ b/crates/krishiv-exec/tests/planner/python_relations.rs @@ -694,6 +694,35 @@ fn a_python_aggregate_yields_one_row_per_group() { 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( diff --git a/crates/krishiv-optimizer/src/collation_keys.rs b/crates/krishiv-optimizer/src/collation_keys.rs index 1d2cbc651..7cf2b0d24 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; @@ -245,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( @@ -442,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-plan-codec/src/decode.rs b/crates/krishiv-plan-codec/src/decode.rs index 6a33591f2..b6027e1bc 100644 --- a/crates/krishiv-plan-codec/src/decode.rs +++ b/crates/krishiv-plan-codec/src/decode.rs @@ -571,6 +571,7 @@ impl Decoder { input: self.boxed(p.input)?, layout: layout(&p.layout), keys: usizes(&p.keys)?, + key_values: usizes(&p.key_values)?, calls: p .calls .iter() diff --git a/crates/krishiv-plan-codec/src/encode.rs b/crates/krishiv-plan-codec/src/encode.rs index b21b5149f..7cbf9f315 100644 --- a/crates/krishiv-plan-codec/src/encode.rs +++ b/crates/krishiv-plan-codec/src/encode.rs @@ -362,11 +362,13 @@ impl Encoder { 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| { 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 f21fc83d6..d09ee581f 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 @@ -268,7 +268,8 @@ pub struct PythonCoGroupMapNode { } /// feature: python_udf_relations /// An aggregate of grouped-aggregate Python UDF calls over `input` clustered -/// and sorted by the key columns. +/// 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")] @@ -279,6 +280,8 @@ pub struct PythonAggregateNode { 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 diff --git a/crates/krishiv-plan-codec/tests/common/inventory.rs b/crates/krishiv-plan-codec/tests/common/inventory.rs index 5d5ca31bf..d76c12c88 100644 --- a/crates/krishiv-plan-codec/tests/common/inventory.rs +++ b/crates/krishiv-plan-codec/tests/common/inventory.rs @@ -1827,6 +1827,7 @@ pub(crate) fn python_relations() -> Vec<(&'static str, PhysNode)> { 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), diff --git a/proto/krishiv/plan/v1/plan.proto b/proto/krishiv/plan/v1/plan.proto index 8e65c91eb..bb4d680d3 100644 --- a/proto/krishiv/plan/v1/plan.proto +++ b/proto/krishiv/plan/v1/plan.proto @@ -211,12 +211,14 @@ message PythonCoGroupMapNode { // feature: python_udf_relations // An aggregate of grouped-aggregate Python UDF calls over `input` clustered -// and sorted by the key columns. +// 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 From 804c92d11f13b314c53640ce37c835f0116d7dcc Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 07:26:52 +0000 Subject: [PATCH 23/29] plan-codec: the python_aggregate golden with key_values (#746 review) The PythonAggregate case now carries key_values (field 5); its golden, new in this PR, is written again with KRISHIV_BLESS_GOLDEN=1. Review thread PRRT_kwDOVMT3-s63aD-J. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../python_udf_relations/python_aggregate.bin | Bin 902 -> 905 bytes 1 file changed, 0 insertions(+), 0 deletions(-) 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 index 74740a6121ff11e92145cb7d786bf1306743aa9b..6918ddd19e31be2f8af2725425a9947ddb297c93 100644 GIT binary patch delta 47 zcmZo;?_}Q~z$h>IiBU|5iIGW)kwJ-tgH3=FNOA~pX)!V?FltR;nz&G8vjd|W69BST B2YLVi delta 44 ycmeBVZ)4vez$hd3iBVXHiIGW)kwJ-tgH3=FNOA~pDKKhHV4Ap4WV17)8xsJcp$A9+ From 559fe5264542a81f531f0d72e82cb19ab051b4ce Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 07:34:10 +0000 Subject: [PATCH 24/29] plan, optimizer, plan-codec: a Python relation UDF carries its determinism (#746 review) A relation node's UDF is a reference, not an expression of the node, so the plan determinism check never saw it: a CTE over a map with a nondeterministic UDF (asNondeterministic, Spark Connect's CommonInlineUserDefinedFunction.deterministic) counted as deterministic and was inlined once per reference, each reference running the UDF on its own. In Spark the node's functionExpr is a PythonUDF whose deterministic is udfDeterministic, and QueryPlan.deterministic and InlineCTE.shouldInline see it (v4.2.0). PythonUdfRef gains deterministic (PythonUdfRef::of takes the call's), and python_relation_is_deterministic says whether a node's relation UDF is. The optimizer's plan determinism check (inline_cte, share_cte, win_magic and the physical planner's shared CTEs through plan_is_deterministic) and the analyzer's asks it. The physical PythonRelation carries it (field 4, python_udf_relations, new in this PR); the relation goldens are written again. Review thread PRRT_kwDOVMT3-s63aD-G. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../physical/instantiate/python_relation.rs | 2 +- .../tests/planner/python_relations.rs | 1 + crates/krishiv-optimizer/src/inline_cte.rs | 5 +- crates/krishiv-optimizer/src/lib.rs | 3 +- .../tests/it/python_relations.rs | 66 ++++++++++++++++++ crates/krishiv-plan-codec/src/decode.rs | 1 + crates/krishiv-plan-codec/src/encode.rs | 1 + .../src/proto/krishiv.plan.v1.pb.rs | 3 + .../tests/common/inventory.rs | 1 + .../python_cogroup_map.bin | Bin 938 -> 940 bytes .../python_udf_relations/python_group_map.bin | Bin 914 -> 917 bytes .../python_udf_relations/python_map.bin | Bin 910 -> 913 bytes .../src/analyzer/subquery/support.rs | 6 +- crates/krishiv-plan/src/cache.rs | 1 + crates/krishiv-plan/src/function.rs | 13 +++- crates/krishiv-plan/src/lib.rs | 2 +- crates/krishiv-plan/src/python.rs | 18 ++++- .../src/tests/python_relations.rs | 36 +++++++++- proto/krishiv/plan/v1/plan.proto | 2 + 19 files changed, 150 insertions(+), 11 deletions(-) diff --git a/crates/krishiv-exec/src/physical/instantiate/python_relation.rs b/crates/krishiv-exec/src/physical/instantiate/python_relation.rs index 1382012f8..8837d8da0 100644 --- a/crates/krishiv-exec/src/physical/instantiate/python_relation.rs +++ b/crates/krishiv-exec/src/physical/instantiate/python_relation.rs @@ -244,7 +244,7 @@ impl Instantiation<'_> { "a Python aggregate call is not a function".to_owned(), )); }; - let Some(udf) = PythonUdfRef::of(&f.id) else { + let Some(udf) = PythonUdfRef::of(&f.id, f.deterministic) else { return Err(internal( "a Python aggregate call is not a Python UDF".to_owned(), )); diff --git a/crates/krishiv-exec/tests/planner/python_relations.rs b/crates/krishiv-exec/tests/planner/python_relations.rs index f3ff94058..accf3d7aa 100644 --- a/crates/krishiv-exec/tests/planner/python_relations.rs +++ b/crates/krishiv-exec/tests/planner/python_relations.rs @@ -297,6 +297,7 @@ fn udf(name: &str, eval_type: PythonEvalType) -> PythonUdfRef { version: 1, eval_type, thread_safe: true, + deterministic: true, } } 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 af0708a12..25d59d830 100644 --- a/crates/krishiv-optimizer/src/lib.rs +++ b/crates/krishiv-optimizer/src/lib.rs @@ -861,7 +861,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/tests/it/python_relations.rs b/crates/krishiv-optimizer/tests/it/python_relations.rs index 4b921cd96..65e3bf119 100644 --- a/crates/krishiv-optimizer/tests/it/python_relations.rs +++ b/crates/krishiv-optimizer/tests/it/python_relations.rs @@ -20,6 +20,7 @@ fn udf(eval_type: PythonEvalType) -> PythonUdfRef { version: 1, eval_type, thread_safe: true, + deterministic: true, } } @@ -130,3 +131,68 @@ fn a_filter_stays_above_a_python_relation_whose_input_keeps_every_column() { ); } } + +/// 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 b6027e1bc..b02506c12 100644 --- a/crates/krishiv-plan-codec/src/decode.rs +++ b/crates/krishiv-plan-codec/src/decode.rs @@ -1733,6 +1733,7 @@ fn python_relation(p: &pb::PythonRelation) -> Result, #[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)] diff --git a/crates/krishiv-plan-codec/tests/common/inventory.rs b/crates/krishiv-plan-codec/tests/common/inventory.rs index d76c12c88..6b16b5b45 100644 --- a/crates/krishiv-plan-codec/tests/common/inventory.rs +++ b/crates/krishiv-plan-codec/tests/common/inventory.rs @@ -1741,6 +1741,7 @@ fn python_udf_ref(eval_type: krishiv_plan::PythonEvalType) -> krishiv_plan::Pyth version: 4, eval_type, thread_safe: true, + deterministic: true, } } 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 index ed5c33e9c5217a515cf3497a18fce6ca1418a937..dced81544d4a809614d12179723defc313c45bc1 100644 GIT binary patch delta 39 vcmZ3*zJ`5+03+{g##fBf8HE@bBp8+YC(ANUU{zq$n!q%1p~&VVjL#SW&!-Ca delta 30 mcmZ3(zKVT=03+{A##f9}8HE@bBp8)?C(ANU*u0+cFCzegSP07i 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 index 6bdc79080b17a1b88fd8bc524233846ffd5da90e..05bd7cdd0f533fc0e32ecfbe8fbb99f964d8d420 100644 GIT binary patch delta 37 tcmbQlK9zlg03&w`<0HlfMj=K9sp`p6jMc0Pj9L?zCN31&yoT`^BLK8R3M&8r delta 34 qcmbQrK8byU03%l;<0Hm8Aw~wN%E{7<)y!HGm?kb1*}R_d86yCqX9@!V 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 index c9d37f7efb6d584620ce51ce1fed4fe5e90156b2..4fac2017d1257b5ee76a8f43dee3609d40facfa2 100644 GIT binary patch delta 24 fcmeBUpUA#JfRVj{@fKsPQ1xUn#_G-ejQ<${T7(D0 delta 23 fcmbQp-p9T{fRU~47Gt$gdjLa|1tsqSHlOu 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 1bb6976f6..3d2ab48fb 100644 --- a/crates/krishiv-plan/src/cache.rs +++ b/crates/krishiv-plan/src/cache.rs @@ -1886,6 +1886,7 @@ mod tests { version: 1, eval_type, thread_safe: true, + deterministic: true, } } diff --git a/crates/krishiv-plan/src/function.rs b/crates/krishiv-plan/src/function.rs index d2938b83b..4ee461532 100644 --- a/crates/krishiv-plan/src/function.rs +++ b/crates/krishiv-plan/src/function.rs @@ -548,6 +548,13 @@ pub struct PythonUdfRef { 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 { @@ -562,10 +569,11 @@ impl PythonUdfRef { } } - /// The reference of a [`FunctionId::PythonUdf`]; `None` for any other + /// 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) -> Option { + pub fn of(id: &FunctionId, deterministic: bool) -> Option { match id { FunctionId::PythonUdf { name, @@ -577,6 +585,7 @@ impl PythonUdfRef { version: *version, eval_type: *eval_type, thread_safe: *thread_safe, + deterministic, }), _ => None, } diff --git a/crates/krishiv-plan/src/lib.rs b/crates/krishiv-plan/src/lib.rs index 120abe8de..07e0a4b79 100644 --- a/crates/krishiv-plan/src/lib.rs +++ b/crates/krishiv-plan/src/lib.rs @@ -176,6 +176,6 @@ pub use plan::{ TABLE_FUNCTION_NAME_SEP, TableArg, TableVersion, UnpivotColumnSet, generator_arg, without_generated_marker, }; -pub use python::{check_python_aggregates, is_python_udaf}; +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/python.rs b/crates/krishiv-plan/src/python.rs index 13a340a6d..82d0b379a 100644 --- a/crates/krishiv-plan/src/python.rs +++ b/crates/krishiv-plan/src/python.rs @@ -7,7 +7,7 @@ use std::sync::Arc; use krishiv_common::error::{InvalidPandasUdfPlacement, KrishivError}; -use crate::{Expr, ExprNode, FunctionId, FunctionKind, PythonUdfKind, ResolvedExpr}; +use crate::{Expr, ExprNode, FunctionId, FunctionKind, PlanNode, PythonUdfKind, ResolvedExpr}; /// Whether `e` is a call of a grouped-aggregate Python UDF (202, 217, 252, /// 254). @@ -22,6 +22,22 @@ pub fn is_python_udaf(e: &Expr) -> bool { ) } +/// 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`, diff --git a/crates/krishiv-plan/src/tests/python_relations.rs b/crates/krishiv-plan/src/tests/python_relations.rs index fe93f5ac7..2e7e14fec 100644 --- a/crates/krishiv-plan/src/tests/python_relations.rs +++ b/crates/krishiv-plan/src/tests/python_relations.rs @@ -47,6 +47,7 @@ fn udf(eval_type: PythonEvalType) -> PythonUdfRef { version: 1, eval_type, thread_safe: true, + deterministic: true, } } @@ -149,8 +150,39 @@ fn eval_types_classify_by_what_the_udf_sees() { fn a_udf_reference_is_its_function_identity() { let r = udf(PythonEvalType::GroupedAggArrow); let id = r.function_id(); - assert_eq!(PythonUdfRef::of(&id), Some(r)); - assert_eq!(PythonUdfRef::of(&FunctionId::Builtin(Builtin::Add)), None); + 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] diff --git a/proto/krishiv/plan/v1/plan.proto b/proto/krishiv/plan/v1/plan.proto index bb4d680d3..2ae508db0 100644 --- a/proto/krishiv/plan/v1/plan.proto +++ b/proto/krishiv/plan/v1/plan.proto @@ -228,6 +228,8 @@ 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 From b7332a4ab80fbe2ab7f06044075e660006ebe86d Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 11:22:44 +0000 Subject: [PATCH 25/29] reviews: T063 N1b re-review (APPROVE_WITH_NITS) (opus-B) Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-exec/CLAUDE.md | 2 +- .../krishiv-exec/src/python_relation/sink.rs | 313 +++++++++++++----- .../krishiv-exec/src/python_relation/tests.rs | 201 ++++++++++- 3 files changed, 422 insertions(+), 94 deletions(-) diff --git a/crates/krishiv-exec/CLAUDE.md b/crates/krishiv-exec/CLAUDE.md index 399a8ad28..cfbfcb2b6 100644 --- a/crates/krishiv-exec/CLAUDE.md +++ b/crates/krishiv-exec/CLAUDE.md @@ -130,7 +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, 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/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/python_relation/sink.rs b/crates/krishiv-exec/src/python_relation/sink.rs index 0ac77e040..6d55ea6f2 100644 --- a/crates/krishiv-exec/src/python_relation/sink.rs +++ b/crates/krishiv-exec/src/python_relation/sink.rs @@ -29,12 +29,16 @@ //! **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. +//! 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 the copies made for the client forms), the results a task +//! 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 @@ -150,7 +154,8 @@ pub const DEFAULT_WINDOW: usize = 4; /// The sink side of a Python relation node. pub struct RelationSink { shape: RelationShape, - lanes: Vec, + /// Shared with the drivers' `Blocked` futures, which frame with them. + lanes: Arc<[RelationLane]>, encoder: Option>, /// The output (canonical carriers). schema: SchemaRef, @@ -244,7 +249,7 @@ impl RelationSink { }; Ok(Arc::new(Self { shape, - lanes, + lanes: lanes.into(), encoder, store: SharedSink::new(Arc::clone(&schema), scope), schema, @@ -300,8 +305,9 @@ struct Driver { /// 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. - backlog: VecDeque, + /// 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, @@ -311,7 +317,15 @@ struct Driver { pending: Pending, } -/// What [`RelationSink::flush`] stopped at. +/// 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. @@ -398,8 +412,8 @@ impl RelationSink { } /// The frame of `lane` over `rows` (one side): its columns evaluated, - /// in their client forms; and the bytes of the columns copied to reach - /// them. + /// 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, @@ -411,16 +425,19 @@ impl RelationSink { 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 = e.evaluate_in(rows, partition)?; let client = krishiv_types::arrow::client_arrow_type(a.data_type()); let a = if a.data_type() == &client { a } else { - let moved = crate::carrier::recarried_array(&a, &client)?; - copied += moved.get_array_memory_size(); - moved + crate::carrier::recarried_array(&a, &client)? }; + copied += fresh_bytes(&a.to_data(), &input); fields.push(Field::new(name, client, e.nullable())); columns.push(a); } @@ -458,6 +475,7 @@ impl RelationSink { let fit = (self.max_bytes.saturating_mul(rows) / bytes).max(1); step = step.min(fit); } + let mut feed = lock(&d.feed); let mut offset = 0; while offset < rows { let len = step.min(rows - offset); @@ -469,7 +487,7 @@ impl RelationSink { } let share = d.held_bytes.size() * len / d.held_bytes_rows.max(len); d.held_bytes_rows = d.held_bytes_rows.saturating_sub(len); - d.backlog.push_back(Cut::Rows { + feed.backlog.push_back(Cut::Rows { side: d.held_side, rows: data.slice(offset, len), memory: d.held_bytes.split(share), @@ -484,7 +502,7 @@ impl RelationSink { /// 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; - d.backlog.push_back(Cut::Rows { + lock(&d.feed).backlog.push_back(Cut::Rows { side, rows: RecordBatch::new_empty(Arc::clone(&self.input)), memory: d.held_bytes.split(0), @@ -502,7 +520,7 @@ impl RelationSink { if self.cogrouped() && !d.sides_seen[1] { self.empty_side(d, 1); } - d.backlog.push_back(Cut::EndGroup); + lock(&d.feed).backlog.push_back(Cut::EndGroup); Ok(()) } @@ -535,62 +553,21 @@ impl RelationSink { Ok(()) } - /// Frames the backlog into the lanes while every lane's window has - /// room. Nothing changes for a frame whose copies the pool refuses. + /// 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 { - loop { - if d.lanes - .iter() - .any(|l| lock(&l.io).in_flight() >= self.window) - { - return Ok(Flushed::Window); - } - let Some(cut) = d.backlog.front() else { - return Ok(Flushed::All); - }; - match cut { - Cut::Rows { side, rows, .. } => { - let mut frames = Vec::with_capacity(self.lanes.len()); - for lane in &self.lanes { - frames.push(Self::frame(lane, *side, rows, d.partition)?); - } - let copied: usize = frames.iter().map(|(_, c)| c).sum(); - let more = copied.saturating_sub(d.spare.size()); - if more > 0 && d.spare.try_grow(more).is_err() { - return Ok(Flushed::Memory(more)); - } - let Some(Cut::Rows { memory, .. }) = d.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 - self.lanes.iter().zip(&d.lanes).zip(frames) - { - let memory = FrameMem { - _shared: Some(Arc::clone(&shared)), - _own: d.spare.split(copied), - }; - state.queue(lane.udf.as_ref(), frame, memory)?; - } - } - Cut::EndGroup | Cut::End => { - let end = matches!(cut, Cut::End); - d.backlog.pop_front(); - for (lane, state) in self.lanes.iter().zip(&d.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: d.spare.split(0), - }; - state.queue(lane.udf.as_ref(), frame, memory)?; - } - } - } + 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 @@ -792,19 +769,162 @@ fn lend(parts: Vec, handle: &Arc, bytes: usize) -> lent } -/// The future of a `Blocked` driver: ready once every lane's window has -/// room, a lane ended (its failure is the next call's), or a result waits -/// for memory (the next call asks for it). -fn blocked(lanes: Vec, window: usize) -> BoxFuture<'static, ()> { +/// 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)?; + } + } + } + } +} + +/// 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| { - let mut room = true; - for lane in &lanes { - if lane.poll(cx).is_some() || lock(&lane.io).refused.is_some() { + 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(()); } - room &= lock(&lane.io).in_flight() < window; + 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 + }; + } + } } - if room { Poll::Ready(()) } else { Poll::Pending } })) } @@ -842,7 +962,10 @@ impl Sink for RelationSink { sides_seen: [false; 2], held_bytes: consumer.reservation(), held_bytes_rows: 0, - backlog: VecDeque::new(), + feed: Arc::new(Mutex::new(Feed { + backlog: VecDeque::new(), + failed: None, + })), spare: consumer.reservation(), wants: 0, pending: Pending::default(), @@ -874,12 +997,22 @@ impl Sink for RelationSink { self.flush(&mut l.driver)?; Self::progress(&l.driver)?; } - let d = &l.driver; - if d.lanes - .iter() - .any(|lane| lock(&lane.io).in_flight() >= self.window) + // 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(d.lanes.clone(), 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) } @@ -899,7 +1032,7 @@ impl Sink for RelationSink { } else { self.cut_held(&mut driver, true)?; } - driver.backlog.push_back(Cut::End); + lock(&driver.feed).backlog.push_back(Cut::End); self.flush(&mut driver)?; Self::progress(&driver)?; lock(&g.leftover).push(driver); @@ -1019,7 +1152,7 @@ impl RelationSink { let _missing = lane.try_take_refused(); } } - Ok(d.backlog.is_empty() + Ok(lock(&d.feed).backlog.is_empty() && d.ready.is_none() && d.lanes .iter() @@ -1070,8 +1203,9 @@ impl RelationSink { } room &= q.in_flight() < self.window; } - ready |= joinable || (room && !d.backlog.is_empty()); - all_done &= d.backlog.is_empty() && d.ready.is_none(); + 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(())) @@ -1098,7 +1232,8 @@ 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(); - Ok((lanes, d.backlog.len())) + let backlog = lock(&d.feed).backlog.len(); + Ok((lanes, backlog)) } /// The global state a relation node's reader takes its list from: the diff --git a/crates/krishiv-exec/src/python_relation/tests.rs b/crates/krishiv-exec/src/python_relation/tests.rs index 42cc8cf79..d3ec32cb4 100644 --- a/crates/krishiv-exec/src/python_relation/tests.rs +++ b/crates/krishiv-exec/src/python_relation/tests.rs @@ -40,6 +40,8 @@ struct Shared { 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>, @@ -131,7 +133,9 @@ impl UdfTask for Task { Box::pin(async move { poll_fn(|cx| { self.shared.gate.register(cx.waker()); - if self.shared.closed.load(Ordering::SeqCst) { + 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(()) @@ -253,6 +257,59 @@ fn rig(shape: RelationShape, yields: Yield, limits: RelationLimits, limit: usize 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, @@ -1084,21 +1141,157 @@ async fn a_batch_s_later_group_keeps_its_share_of_the_input_reservation() { // 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(); - tokio::time::timeout(Duration::from_secs(5), wait) - .await - .unwrap(); + 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); +} From 53e4e3fbc51cf532f7a632f52cf644fba8ad584a Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 18:38:35 +0000 Subject: [PATCH 26/29] exec: maxBytesPerBatch cuts frames by the frame columns' Arrow size, as Spark's writer With spark.sql.execution.arrow.maxBytesPerBatch set, the relation sink cut frames by the held input batch's average bytes per row. Uneven rows (two 30 KiB strings, then empty ones) and computed arguments that grow the payload made frames the average did not see. The cut now sizes the frame columns the tasks get (computed arguments, client forms) as Spark's writer does: rows are written while the frame's Arrow size is under the cap, so a frame ends at the row that brings it to the cap or past (BatchedPythonArrowInput.writeSizedBatch, PythonArrowInput.scala:226-233 at v4.2.0; the size is Arrow Java's getBufferSizeFor per field, ArrowWriter.scala:135-143 and 183-187). The size of rows start..start+k comes from validity, offsets and value spans (no copy), and grows with k, so the cut is a binary search. A map's rows travel as one struct column (MapInBatchEvaluatorFactory.scala:56-71), whose validity counts. Spark's default, Int.MaxValue, reads as no limit (PythonArrowInput.scala:227). Tests: the uneven-strings case (frames 2, 100), rows that each fill most of the cap (two per frame), and a row over the cap by itself (it ends the frame it joins; alone it is a frame of its own). Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../physical/instantiate/python_relation.rs | 5 +- crates/krishiv-exec/src/python_relation.rs | 5 +- .../krishiv-exec/src/python_relation/sink.rs | 187 ++++++++++++++++-- .../krishiv-exec/src/python_relation/tests.rs | 79 ++++++++ 4 files changed, 253 insertions(+), 23 deletions(-) diff --git a/crates/krishiv-exec/src/physical/instantiate/python_relation.rs b/crates/krishiv-exec/src/physical/instantiate/python_relation.rs index 8837d8da0..480c3c021 100644 --- a/crates/krishiv-exec/src/physical/instantiate/python_relation.rs +++ b/crates/krishiv-exec/src/physical/instantiate/python_relation.rs @@ -21,7 +21,7 @@ use crate::python_relation::{ /// Spark's `spark.sql.execution.arrow.maxBytesPerBatch`, read as the /// session wrote it (Krishiv registers no key for it yet): no limit unless -/// set. +/// set to other than Spark's default. const MAX_BYTES_PER_BATCH: &str = "spark.sql.execution.arrow.maxBytesPerBatch"; impl Instantiation<'_> { @@ -112,6 +112,9 @@ impl Instantiation<'_> { 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 { diff --git a/crates/krishiv-exec/src/python_relation.rs b/crates/krishiv-exec/src/python_relation.rs index 87499b3f7..3a6fa6d12 100644 --- a/crates/krishiv-exec/src/python_relation.rs +++ b/crates/krishiv-exec/src/python_relation.rs @@ -110,8 +110,9 @@ 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`: most bytes per frame - /// (cut between rows, at least one row); 0 is no limit. + /// `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, diff --git a/crates/krishiv-exec/src/python_relation/sink.rs b/crates/krishiv-exec/src/python_relation/sink.rs index 6d55ea6f2..1545d6a1b 100644 --- a/crates/krishiv-exec/src/python_relation/sink.rs +++ b/crates/krishiv-exec/src/python_relation/sink.rs @@ -22,9 +22,11 @@ //! 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 -//! `maxBytesPerBatch` bytes when set) or at the end of their group or side, -//! as Spark cuts a group's Arrow stream; a group is never held whole. +//! 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 @@ -184,7 +186,8 @@ impl fmt::Debug for RelationSink { pub struct RelationLimits { /// Most rows per frame; 0 is no limit. pub max_rows: usize, - /// Most bytes per frame (cut between rows); 0 is no limit. + /// 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, @@ -430,15 +433,9 @@ impl RelationSink { buffer_ranges(&c.to_data(), &mut input); } for (e, name) in exprs.iter().zip(names) { - let a = e.evaluate_in(rows, partition)?; - let client = krishiv_types::arrow::client_arrow_type(a.data_type()); - let a = if a.data_type() == &client { - a - } else { - crate::carrier::recarried_array(&a, &client)? - }; + let a = client_column(e.as_ref(), rows, partition)?; copied += fresh_bytes(&a.to_data(), &input); - fields.push(Field::new(name, client, e.nullable())); + 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())); @@ -465,21 +462,40 @@ impl RelationSink { arrow_select::concat::concat_batches(&self.input, &held) .map_err(|e| internal(format!("held Python UDF rows: {e}")))? }; - let mut step = if self.max_rows == 0 { + let step = if self.max_rows == 0 { rows } else { self.max_rows }; - if self.max_bytes > 0 { - let bytes = data.get_array_memory_size().max(1); - let fit = (self.max_bytes.saturating_mul(rows) / bytes).max(1); - step = step.min(fit); - } + // 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 len = step.min(rows - offset); - if len < step && !all { + 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; @@ -830,6 +846,137 @@ fn flush_feed( } } +/// `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| { diff --git a/crates/krishiv-exec/src/python_relation/tests.rs b/crates/krishiv-exec/src/python_relation/tests.rs index d3ec32cb4..5e69f53a6 100644 --- a/crates/krishiv-exec/src/python_relation/tests.rs +++ b/crates/krishiv-exec/src/python_relation/tests.rs @@ -737,6 +737,85 @@ async fn max_bytes_per_batch_cuts_frames_between_rows() { 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. From 7da68d8f73b8e7231927c1966a52a61bb0aec350 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sun, 15 Nov 2026 00:55:34 +0000 Subject: [PATCH 27/29] exec: an aggregate's fast Python UDF lane waits for its peers instead of filling the pool In an aggregate with several UDF lanes a group is stored only once every lane has its value, and lane results are not spillable. A fast lane could reserve results until the pool was full before a slower lane's first value was reserved; nothing could be joined or released, and the slow lane's value waited for memory until arbitration failed (UNABLE_TO_ACQUIRE_MEMORY). A lane now reads no result while it holds more rows than a peer whose task still runs (`lane.rs` `ahead`): it waits on its own task, holding at most one result past the others, and is woken when a peer yields or ends. A lane of a map or grouped shape has no peers. Memory accounting is unchanged: every result is still reserved as it arrives. Test: two lanes, the slow one gated, a 1 MiB pool with spill on: the fast lane reads one value then waits, and finalize completes (before: the slow lane's 104-byte value was refused). Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../krishiv-exec/src/python_relation/lane.rs | 134 ++++++++++++++---- .../krishiv-exec/src/python_relation/sink.rs | 9 +- .../krishiv-exec/src/python_relation/tests.rs | 71 ++++++++++ 3 files changed, 183 insertions(+), 31 deletions(-) diff --git a/crates/krishiv-exec/src/python_relation/lane.rs b/crates/krishiv-exec/src/python_relation/lane.rs index 5449c28eb..1de5814f9 100644 --- a/crates/krishiv-exec/src/python_relation/lane.rs +++ b/crates/krishiv-exec/src/python_relation/lane.rs @@ -14,11 +14,19 @@ //! 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: +//! it waits on its own task, holding at most one result past the others, +//! and a slower lane's value never waits for memory its peers pin. A lane +//! of a map or grouped shape has no peers. use std::collections::VecDeque; use std::future::poll_fn; use std::sync::Arc; -use std::task::{Context, Poll}; +use std::task::{Context, Poll, Waker}; use arrow_array::RecordBatch; use futures::future::BoxFuture; @@ -51,6 +59,11 @@ pub(super) struct LaneIo { /// An empty reservation of the driver's consumer, the results' are /// split from. pub(super) memory: Reservation, + /// The task yielded its last result. + ended: bool, + /// The pump waiting while the lane is ahead of a peer, woken when a + /// peer yields or ends. + waiting: Option, } impl LaneIo { @@ -62,8 +75,21 @@ impl LaneIo { pub(super) fn has_results(&self) -> bool { !self.outputs.is_empty() || self.refused.is_some() } + + /// The rows of the results not taken by the driver, the refused one's + /// with them. + fn rows(&self) -> usize { + self.outputs + .iter() + .map(|(b, _)| b.num_rows()) + .sum::() + + self.refused.as_ref().map_or(0, RecordBatch::num_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. @@ -77,21 +103,40 @@ enum Pump { 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 { - /// A lane whose results reserve in `memory`'s consumer. - pub(super) fn new(memory: Reservation) -> Self { - Self { - io: Arc::new(Mutex::new(LaneIo { - inbound: VecDeque::new(), - sending: None, - outputs: VecDeque::new(), - refused: None, - memory, - })), - pump: Arc::new(Mutex::new(Pump::Idle)), - } + /// The lanes of one driver, one per reservation, whose results reserve + /// in that reservation's consumer; each reads no result while it holds + /// more rows than a peer still running. + pub(super) fn group(memories: impl IntoIterator) -> Vec { + let peers: Peers = memories + .into_iter() + .map(|memory| { + Arc::new(Mutex::new(LaneIo { + inbound: VecDeque::new(), + sending: None, + outputs: VecDeque::new(), + refused: None, + memory, + ended: 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 @@ -129,7 +174,7 @@ impl Lane { let mut p = lock(&self.pump); if matches!(*p, Pump::Idle) { let task = udf.open()?; - *p = Pump::Running(pump(task, Arc::clone(&self.io))); + *p = Pump::Running(pump(task, Arc::clone(&self.peers), self.index)); } } lock(&self.io).inbound.push_back((frame, memory)); @@ -177,10 +222,37 @@ impl Lane { } } -/// Sends the queued frames to `task` while it reads the task's results. -fn pump(task: Box, io: Arc>) -> BoxFuture<'static, Result<()>> { +/// Whether lane `index` holds more rows than a peer whose task still runs. +/// Each queue is locked alone, so two lanes asking at once never deadlock. +fn ahead(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.ended && q.rows() < own + }) +} + +/// Wakes every peer of lane `index` waiting for it to yield or end. +fn wake_peers(peers: &[Arc>], index: usize) { + for (i, io) in peers.iter().enumerate() { + if i != index + && let Some(w) = lock(io).waiting.take() + { + 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. @@ -205,27 +277,39 @@ fn pump(task: Box, io: Arc>) -> BoxFuture<'static, Re }; let receive = async { loop { - // Never woken either: whoever reserves the refused batch - // polls the pump next. - poll_fn(|_| { + // Not woken while refused: whoever reserves the refused + // batch polls the pump next. Ahead of a peer: woken when a + // peer yields or ends (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 ahead(&peers, index) { Poll::Pending } else { + lock(&io).waiting = None; Poll::Ready(()) } }) .await; let Some(batch) = task.next().await? else { + lock(&io).ended = true; + wake_peers(&peers, index); return Ok(()); }; let bytes = batch.get_array_memory_size(); - let mut q = lock(&io); - 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); + { + let mut q = lock(&io); + 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); } }; futures::future::try_join(send, receive).await.map(drop) diff --git a/crates/krishiv-exec/src/python_relation/sink.rs b/crates/krishiv-exec/src/python_relation/sink.rs index 1545d6a1b..982185ff2 100644 --- a/crates/krishiv-exec/src/python_relation/sink.rs +++ b/crates/krishiv-exec/src/python_relation/sink.rs @@ -47,7 +47,8 @@ //! 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. A stored result's bytes pass to the store's +//! other lanes move, which hold at most one result past it (a lane ahead +//! of a running peer reads no more, `lane.rs`). 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 @@ -1094,11 +1095,7 @@ impl Sink for RelationSink { consumer: Arc::clone(consumer), store: self.store.local_state(ctx, &g.store)?, driver: Driver { - lanes: self - .lanes - .iter() - .map(|_| Lane::new(consumer.reservation())) - .collect(), + lanes: Lane::group(self.lanes.iter().map(|_| consumer.reservation())), ready: None, partition: ctx.partition(), current: None, diff --git a/crates/krishiv-exec/src/python_relation/tests.rs b/crates/krishiv-exec/src/python_relation/tests.rs index 5e69f53a6..87aee8265 100644 --- a/crates/krishiv-exec/src/python_relation/tests.rs +++ b/crates/krishiv-exec/src/python_relation/tests.rs @@ -1043,6 +1043,77 @@ async fn an_aggregate_waits_for_its_slowest_lane_without_spinning() { done(r); } +/// An aggregate whose fast lane 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 peer still running, so it holds at most one result past the +/// others and the slow lane always gets its value's memory. +#[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 = 12_000; + 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(); + // Faster than the pool's spill wait: nothing waits for memory. + tokio::time::timeout(Duration::from_secs(4), 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); +} + #[tokio::test] async fn a_result_refused_in_finalize_spills_instead_of_waiting() { let dir = tempfile::tempdir().unwrap(); From dfb132037b978c6c3339f1aed754ab6a8551b40c Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sun, 15 Nov 2026 01:30:25 +0000 Subject: [PATCH 28/29] exec: a fast aggregate lane is held back only while its peer has frames to take Round-5 review M1: the hold of 7da68d8f7 stopped a fast lane's reads while any running peer had fewer rows. A runtime that takes no frame while its results are unread then filled the fast lane's window, which stops the frames to every lane; a peer runtime that needs later frames before it yields (one batching values across groups) never got them, and exec has no call deadline: a hang. The hold now also needs the lagging peer to have frames it has not taken (`lane.rs` `held_back`: `in_flight() > 0`). The peer wakes the waiting pump when it takes its last frame, yields or stops (ended or failed). A peer starved of input never holds anyone back, so progress depends only on a lane taking frames it already has; past that the pool 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. With the production window the lead stays bounded by the window plus what the peer runtime buffers. Also from the review: - m1: only an aggregate's lanes are peers (`Lane::group(.., peered)`); a map or grouped shape's lanes are alone. - m2: a running row count per lane (`LaneIo::rows`, `pop_output`, `took`) instead of re-summing the queue on every result. - m3: "one result batch", not "one result". - m4: a peer's waker is taken under its lock and woken after release. Tests (M2) at the relation sink's production window (DEFAULT_WINDOW, 4): a peer batching values across 40 groups with a fast runtime that takes no frame while a result is unread (hung before, timeout 10 s); a peer ending with fewer values (the mismatch error, hung before); a peer failing while the fast lane is held (its error); the gated slow lane at the production window. The deep-window reproduction of the original stall is kept; once its slow lane opens it takes the whole window at once, so the hold ends and the lanes share the pool (slower, 12 s in a debug build, never stuck). Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../krishiv-exec/src/python_relation/lane.rs | 126 +++++++--- .../krishiv-exec/src/python_relation/pair.rs | 3 +- .../krishiv-exec/src/python_relation/sink.rs | 11 +- .../krishiv-exec/src/python_relation/tests.rs | 230 +++++++++++++++++- 4 files changed, 314 insertions(+), 56 deletions(-) diff --git a/crates/krishiv-exec/src/python_relation/lane.rs b/crates/krishiv-exec/src/python_relation/lane.rs index 1de5814f9..2fb16a0ef 100644 --- a/crates/krishiv-exec/src/python_relation/lane.rs +++ b/crates/krishiv-exec/src/python_relation/lane.rs @@ -18,10 +18,20 @@ //! **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: -//! it waits on its own task, holding at most one result past the others, -//! and a slower lane's value never waits for memory its peers pin. A lane -//! of a map or grouped shape has no peers. +//! 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; @@ -59,10 +69,13 @@ pub(super) struct LaneIo { /// An empty reservation of the driver's consumer, the results' are /// split from. pub(super) memory: Reservation, - /// The task yielded its last result. - ended: bool, - /// The pump waiting while the lane is ahead of a peer, woken when a - /// peer yields or ends. + /// The rows of [`Self::outputs`] and [`Self::refused`], kept as they + /// change (a peer's hold reads it on every result). + rows: usize, + /// The pump ended: the task yielded its last result or failed. + 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, } @@ -76,14 +89,19 @@ impl LaneIo { !self.outputs.is_empty() || self.refused.is_some() } - /// The rows of the results not taken by the driver, the refused one's - /// with them. - fn rows(&self) -> usize { - self.outputs - .iter() - .map(|(b, _)| b.num_rows()) - .sum::() - + self.refused.as_ref().map_or(0, RecordBatch::num_rows) + /// 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) { + self.rows = self.rows.saturating_sub(rows); } } @@ -110,9 +128,19 @@ pub(super) struct Lane { impl Lane { /// The lanes of one driver, one per reservation, whose results reserve - /// in that reservation's consumer; each reads no result while it holds - /// more rows than a peer still running. - pub(super) fn group(memories: impl IntoIterator) -> Vec { + /// 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| { @@ -122,7 +150,8 @@ impl Lane { outputs: VecDeque::new(), refused: None, memory, - ended: false, + rows: 0, + stopped: false, waiting: None, })) }) @@ -199,7 +228,10 @@ impl Lane { /// 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 { - lock(&self.io).refused.take() + 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 @@ -222,24 +254,30 @@ impl Lane { } } -/// Whether lane `index` holds more rows than a peer whose task still runs. -/// Each queue is locked alone, so two lanes asking at once never deadlock. -fn ahead(peers: &[Arc>], index: usize) -> bool { - let Some(own) = peers.get(index).map(|io| lock(io).rows()) else { +/// 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.ended && q.rows() < own + i != index && !q.stopped && q.in_flight() > 0 && q.rows < own }) } -/// Wakes every peer of lane `index` waiting for it to yield or end. +/// 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 - && let Some(w) = lock(io).waiting.take() - { + if i == index { + continue; + } + let waker = lock(io).waiting.take(); + if let Some(w) = waker { w.wake(); } } @@ -269,7 +307,15 @@ fn pump(task: Box, peers: Peers, index: usize) -> BoxFuture<'static .await; let end = matches!(frame, Frame::End); task.send(frame).await?; - lock(&io).sending = None; + 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>(()); } @@ -278,15 +324,16 @@ fn pump(task: Box, peers: Peers, index: usize) -> BoxFuture<'static let receive = async { loop { // Not woken while refused: whoever reserves the refused - // batch polls the pump next. Ahead of a peer: woken when a - // peer yields or ends (the waker is set before the check, - // so a peer that moves after it wakes this pump). + // 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 ahead(&peers, index) { + if held_back(&peers, index) { Poll::Pending } else { lock(&io).waiting = None; @@ -295,13 +342,12 @@ fn pump(task: Box, peers: Peers, index: usize) -> BoxFuture<'static }) .await; let Some(batch) = task.next().await? else { - lock(&io).ended = 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)); @@ -312,6 +358,10 @@ fn pump(task: Box, peers: Peers, index: usize) -> BoxFuture<'static wake_peers(&peers, index); } }; - futures::future::try_join(send, receive).await.map(drop) + 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 index 9dab87333..e6e2b0990 100644 --- a/crates/krishiv-exec/src/python_relation/pair.rs +++ b/crates/krishiv-exec/src/python_relation/pair.rs @@ -153,11 +153,12 @@ fn take_values(state: &Lane, n: usize, memory: &mut Vec) -> Result< parts.push(Arc::clone(front.column(0)).slice(0, len)); want -= len; if len == front.num_rows() { - if let Some((_, r)) = q.outputs.pop_front() { + 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(); diff --git a/crates/krishiv-exec/src/python_relation/sink.rs b/crates/krishiv-exec/src/python_relation/sink.rs index 982185ff2..9a4edea0d 100644 --- a/crates/krishiv-exec/src/python_relation/sink.rs +++ b/crates/krishiv-exec/src/python_relation/sink.rs @@ -47,8 +47,8 @@ //! 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 past it (a lane ahead -//! of a running peer reads no more, `lane.rs`). A stored result's bytes pass to the store's +//! 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 @@ -632,7 +632,7 @@ impl RelationSink { ); } for (lane, state) in self.lanes.iter().zip(&d.lanes) { - let next = lock(&state.io).outputs.pop_front(); + let next = lock(&state.io).pop_output(); if let Some((batch, memory)) = next { return Ok(Some((self.output(lane, &batch)?, vec![memory]))); } @@ -1095,7 +1095,10 @@ impl Sink for RelationSink { consumer: Arc::clone(consumer), store: self.store.local_state(ctx, &g.store)?, driver: Driver { - lanes: Lane::group(self.lanes.iter().map(|_| consumer.reservation())), + lanes: Lane::group( + self.lanes.iter().map(|_| consumer.reservation()), + matches!(self.shape, RelationShape::Aggregate { .. }), + ), ready: None, partition: ctx.partition(), current: None, diff --git a/crates/krishiv-exec/src/python_relation/tests.rs b/crates/krishiv-exec/src/python_relation/tests.rs index 87aee8265..4c7e8cbed 100644 --- a/crates/krishiv-exec/src/python_relation/tests.rs +++ b/crates/krishiv-exec/src/python_relation/tests.rs @@ -22,7 +22,7 @@ use krishiv_common::error::{ErrorClass, InternalError}; use krishiv_memory::{MemoryConfig, MemoryPool, SpillOwner, TreePool}; use krishiv_types::batch::Batch; -use super::sink::{InputColumn, RelationLimits}; +use super::sink::{DEFAULT_WINDOW, InputColumn, RelationLimits}; use super::{Frame, PreparedRelationUdf, RelationLane, RelationShape, RelationSink, UdfTask}; use crate::operator::{Sink, SinkStatus}; use crate::operators::ScalarExpr; @@ -51,6 +51,9 @@ struct Shared { 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>, } /// Polls of a task's `next` that only a busy loop reaches. @@ -80,6 +83,9 @@ enum Yield { 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 { @@ -95,6 +101,7 @@ impl PreparedRelationUdf for Fake { yields: self.yields, state: std::sync::Mutex::new(TaskState::default()), ready: AtomicWaker::new(), + drained: AtomicWaker::new(), })) } } @@ -104,6 +111,8 @@ struct TaskState { out: VecDeque, rows: i64, ended: bool, + /// `Coalesce`: the values not yielded yet. + held: Vec, } struct Task { @@ -111,6 +120,8 @@ struct Task { yields: Yield, state: std::sync::Mutex, ready: AtomicWaker, + /// Woken as `next` takes a result (`Shared::out_cap`). + drained: AtomicWaker, } impl Drop for Task { @@ -120,18 +131,23 @@ impl Drop for Task { } 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(vec![n])) as ArrayRef], - ) - .unwrap() + 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); @@ -164,7 +180,7 @@ impl UdfTask for Task { let copies = match self.yields { Yield::Twice => 2, Yield::Burst(n) => n, - Yield::Nothing | Yield::Count | Yield::CountTwice => 0, + Yield::Nothing | Yield::Count | Yield::CountTwice | Yield::Coalesce(_) => 0, }; for _ in 0..copies { s.out.push_back(batch.clone()); @@ -178,10 +194,23 @@ impl UdfTask for Task { 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 => s.ended = true, + Frame::End => { + if !s.held.is_empty() { + let values = std::mem::take(&mut s.held); + s.out.push_back(counts(values)); + } + s.ended = true; + } } drop(s); self.ready.wake(); @@ -202,6 +231,7 @@ impl UdfTask for Task { 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)), @@ -325,7 +355,7 @@ fn rig_in( 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::Count | Yield::CountTwice | Yield::Coalesce(_) => ( vec![ArrowType::Int64], match shape { RelationShape::Aggregate { .. } => vec![ @@ -1047,8 +1077,9 @@ async fn an_aggregate_waits_for_its_slowest_lane_without_spinning() { /// 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 peer still running, so it holds at most one result past the -/// others and the slow lane always gets its value's memory. +/// than a running peer that has frames to take, so here it holds at most +/// one result batch past the gated lane, which always gets its value's +/// memory. (A deep window: the stall as first seen.) #[tokio::test] async fn a_fast_aggregate_lane_waits_for_its_peers_instead_of_filling_the_pool() { let dir = tempfile::tempdir().unwrap(); @@ -1102,8 +1133,9 @@ async fn a_fast_aggregate_lane_waits_for_its_peers_instead_of_filling_the_pool() assert_eq!(r.shared.taken.load(Ordering::SeqCst), 1); r.sink.combine(l, &g).unwrap(); slow.open_gate(); - // Faster than the pool's spill wait: nothing waits for memory. - tokio::time::timeout(Duration::from_secs(4), r.sink.finalize(&g, &r.ctx)) + // The open slow lane takes its whole deep window at once, so the hold + // ends and the two lanes share the pool: slower, never stuck. + tokio::time::timeout(Duration::from_secs(40), r.sink.finalize(&g, &r.ctx)) .await .unwrap() .unwrap(); @@ -1114,6 +1146,178 @@ async fn a_fast_aggregate_lane_waits_for_its_peers_instead_of_filling_the_pool() 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); +} + +/// 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(); From 1ce019a7d5c56bb4b745878e3437ce78fc32fa6a Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sun, 15 Nov 2026 06:57:59 +0000 Subject: [PATCH 29/29] exec: a Python UDF lane whose output ended holds no peer back; the last-frame wake is tested (#746 review r6) - m2: when a task's `next` ends, its lane is marked stopped and its peers woken at once, before its send side finished, so a runtime that ends its output early while a send hangs cannot hold every peer. - m3: a new test, a quiet peer that batches values across groups and signals its reader only on output, hangs without the "took its last frame" wake (checked by deleting it: the test times out) and passes with it. No earlier test failed without that wake. - m1: the deep-window test's doc states what it checks (the hold while the gated peer has frames) and the bound once the gate opens; 3,000 groups instead of 12,000 (0.25 s instead of 21 s) and a 10 s timeout instead of 40 s. - `LaneIo::took` debug-asserts that the running row count does not drift. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../krishiv-exec/src/python_relation/lane.rs | 13 ++- .../krishiv-exec/src/python_relation/tests.rs | 82 ++++++++++++++++--- 2 files changed, 83 insertions(+), 12 deletions(-) diff --git a/crates/krishiv-exec/src/python_relation/lane.rs b/crates/krishiv-exec/src/python_relation/lane.rs index 2fb16a0ef..6243814e6 100644 --- a/crates/krishiv-exec/src/python_relation/lane.rs +++ b/crates/krishiv-exec/src/python_relation/lane.rs @@ -72,7 +72,8 @@ pub(super) struct LaneIo { /// The rows of [`Self::outputs`] and [`Self::refused`], kept as they /// change (a peer's hold reads it on every result). rows: usize, - /// The pump ended: the task yielded its last result or failed. + /// 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. @@ -101,6 +102,7 @@ impl LaneIo { /// `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); } } @@ -342,6 +344,15 @@ fn pump(task: Box, peers: Peers, index: usize) -> BoxFuture<'static }) .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(); diff --git a/crates/krishiv-exec/src/python_relation/tests.rs b/crates/krishiv-exec/src/python_relation/tests.rs index 4c7e8cbed..dcc097fb5 100644 --- a/crates/krishiv-exec/src/python_relation/tests.rs +++ b/crates/krishiv-exec/src/python_relation/tests.rs @@ -54,6 +54,9 @@ struct Shared { /// 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. @@ -174,6 +177,7 @@ impl UdfTask for Task { .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(); @@ -212,8 +216,11 @@ impl UdfTask for Task { s.ended = true; } } + let output = s.out.len() != before || s.ended; drop(s); - self.ready.wake(); + if output || !self.shared.quiet.load(Ordering::SeqCst) { + self.ready.wake(); + } Ok(()) }) } @@ -1073,13 +1080,23 @@ async fn an_aggregate_waits_for_its_slowest_lane_without_spinning() { done(r); } -/// An aggregate whose fast lane 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 here it holds at most -/// one result batch past the gated lane, which always gets its value's -/// memory. (A deep window: the stall as first seen.) +/// 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(); @@ -1125,7 +1142,7 @@ async fn a_fast_aggregate_lane_waits_for_its_peers_instead_of_filling_the_pool() .unwrap(); let g = r.sink.global_state(&r.ctx).unwrap(); let mut l = r.sink.local_state(&r.ctx, &g).unwrap(); - let groups = 12_000; + 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:?}"); @@ -1134,8 +1151,8 @@ async fn a_fast_aggregate_lane_waits_for_its_peers_instead_of_filling_the_pool() 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: slower, never stuck. - tokio::time::timeout(Duration::from_secs(40), r.sink.finalize(&g, &r.ctx)) + // 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(); @@ -1253,6 +1270,49 @@ async fn a_peer_batching_values_across_groups_never_hangs_the_fast_lane() { 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]