plan, exec, optimizer, plan-codec: Python UDF relations: IR, frame seam, operators and codec feature (P3-T014a, PR-26a)

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.

crates/krishiv-common/krishiv-error-classes.json

@@ -130,6 +130,12 @@
     ],
     "sqlState" : "53400"
   },
+  "KRISHIV_PYTHON_UDF_RESULT_MISMATCH" : {
+    "message" : [
+      "The Python UDF <name> returned a result that does not match what it declares: <reason>. Make the function return the declared columns, in their declared types, and for a grouped aggregate one value per group."
+    ],
+    "sqlState" : "39000"
+  },
   "KRISHIV_PYTHON_WORKER_EXITED" : {
     "message" : [
       "Python worker exited unexpectedly (crashed): <reason>. The next query starts a new worker. If a UDF ran out of memory, raise \"krishiv.python.worker.memoryLimit\"."

crates/krishiv-exec/CLAUDE.md

@@ -130,6 +130,7 @@ Paths are relative to this crate.
 - `src/window/mod.rs` window operator and plan types; `src/window/eval.rs` frames, ranking and frame aggregates; `src/planner/window.rs` plans it over sorted input
 - `src/operators/file_scan.rs` and `file_write.rs` file scan and write operators over `krishiv-io`
 - `src/planner/write.rs` (P3-T023a) `WriteMode` (append, truncate, overwrite by filter, dynamic overwrite), `WriteRequest` (mode, table name, input schema, options, write id) and `TableCapabilities` (Spark's `TableCapability` write values, checked as Spark v4.0.0's `TableCapabilityCheck`: every mode needs `BATCH_WRITE` or `V1_BATCH_WRITE`, a truncate also `TRUNCATE` or `OVERWRITE_BY_FILTER`, dynamic overwrite `BATCH_WRITE` and `OVERWRITE_DYNAMIC`); `TableProvider::capabilities` / `writer(&WriteRequest)` default to no write modes, and `TableCapabilities::check` raises `UNSUPPORTED_FEATURE.TABLE_OPERATION` ("... in batch mode") before anything is written
+- `src/python_relation.rs` (design P3-T085 §3.2) the Python relation seam: `RelationCallSpec`, `PreparedRelationUdf::open`, `UdfTask` fed by `Frame`s (`UdfRuntime::prepare_relation` defaults to `KRISHIV_UDF_NOT_AVAILABLE`); `python_relation/sink.rs` `RelationSink` (map, grouped, cogrouped, aggregate shapes: one task per driver, groups cut where the sorted key changes, frames lead with the keys, a one-sided cogroup key gets a zero-row frame for the other side, the window checked per frame with a reserved backlog, `Blocked` while a lane's frame window is full or cuts wait, its future framing the backlog so no input is taken before it drained; a frame's computed columns reserved, results in the shared plan's list as they arrive), `python_relation/lane.rs` (a lane's frames, its reserved results and the pump; a refused result stops the task's reads), `python_relation/pair.rs` (an aggregate's pending keys joined with the values as they arrive); `planner/python_relation.rs` plans `MapPartitions`/`GroupMap`/`CoGroupMap`/Python-UDAF `Aggregate` onto `PhysNode::PythonMap`/`PythonGroupMap`/`PythonCoGroupMap`/`PythonAggregate` (hash exchange + spillable sort on the keys; a cogroup is one union stream sorted by key then side), refuses barrier mode, Python window aggregates and UDTFs (PR-26c); `physical/instantiate/python_relation.rs` builds them
 - `src/planner/ndv_sample.rs` (P2-T164 U6) `FileTable::sample_distinct`: plan-time distinct-value samples of the file columns an equality or `IN` compares, and of the column group keys while `krishiv.optimizer.groupKeyCode` is on (the session picks them, `catalog/ndv_sampling.rs`) (first rows in file order through `ScanStream`, `krishiv.optimizer.ndvSampling.*`), exact for a table the sample read whole, else GEE (`estimate`); cached in `krishiv_io::DistinctSamples`; `FileTable::statistics` reports them as `StatSource::Sampled`; `FileTable::sample_outcomes` says per table what each sample did (`SampleOutcome`: sampled, cached, timed out, no time left), which `EXPLAIN COST` prints as "Distinct-Value Samples" and `EXPLAIN VERBOSE` as a `Statistics` step
 
 ## Extension points

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.

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,

crates/krishiv-exec/src/physical/instantiate.rs

@@ -89,6 +89,7 @@ mod detached;
 mod distributed;
 mod order;
 mod python_eval;
+mod python_relation;
 pub(crate) use detached::{Detached, LoopRun};
 mod join;
 use distributed::planned_schema_holds;
@@ -927,6 +928,10 @@ impl Instantiation<'_> {
     fn untagged(&mut self, n: &PhysNode) -> Result<Built> {
         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(

crates/krishiv-exec/src/physical/instantiate/python_relation.rs

@@ -0,0 +1,300 @@
+// SPDX-License-Identifier: Apache-2.0
+//! Instantiation of the Python relation nodes (design P3-T085 §3.1-§3.2):
+//! each node's UDFs prepared by the worker's [`UdfRuntime`] as relation call
+//! sites, and a [`RelationSink`] breaker over the input.
+//!
+//! [`UdfRuntime`]: crate::python_eval::UdfRuntime
+
+use super::*;
+use arrow_schema::{DataType as ArrowType, Field};
+use krishiv_common::config::ARROW_MAX_RECORDS_PER_BATCH;
+use krishiv_common::error::{KrishivError, KrishivUdfNotAvailable};
+use krishiv_plan::{ExprNode, PythonUdfRef, ResolvedExpr};
+use krishiv_types::arrow::client_arrow_type;
+
+use crate::operators::ScalarExpr;
+use crate::physical::{PythonCallPlan, PythonRelationPlan};
+use crate::python_relation::{
+    DEFAULT_WINDOW, InputColumn, PreparedRelationUdf, RelationCallSpec, RelationLane,
+    RelationLimits, RelationShape, RelationSink, RunnerConf,
+};
+
+/// Spark's `spark.sql.execution.arrow.maxBytesPerBatch`, read as the
+/// session wrote it (Krishiv registers no key for it yet): no limit unless
+/// set to other than Spark's default.
+const MAX_BYTES_PER_BATCH: &str = "spark.sql.execution.arrow.maxBytesPerBatch";
+
+impl Instantiation<'_> {
+    /// A Python relation node: a breaker whose sink feeds the node's tasks
+    /// and whose source yields what they returned.
+    pub(super) fn python_relation(&mut self, node: &PhysNode) -> Result<Built> {
+        let (input, shape) = match node {
+            PhysNode::PythonMap { input, .. } => (input, RelationShape::Map),
+            PhysNode::PythonGroupMap { input, keys, .. } => {
+                (input, RelationShape::Grouped { keys: keys.clone() })
+            }
+            PhysNode::PythonCoGroupMap {
+                input, keys, side, ..
+            } => (
+                input,
+                RelationShape::CoGrouped {
+                    keys: keys.clone(),
+                    side: *side,
+                },
+            ),
+            PhysNode::PythonAggregate {
+                input,
+                keys,
+                key_values,
+                ..
+            } => (
+                input,
+                RelationShape::Aggregate {
+                    keys: keys.clone(),
+                    key_values: key_values.clone(),
+                },
+            ),
+            _ => return Err(internal("not a Python relation node".to_owned())),
+        };
+        let input = self.node(input)?;
+        let conf = self.runner_conf();
+        let (lanes, schema) = match node {
+            PhysNode::PythonMap { udf, .. }
+            | PhysNode::PythonGroupMap { udf, .. }
+            | PhysNode::PythonCoGroupMap { udf, .. } => {
+                let keys = match &shape {
+                    RelationShape::Grouped { keys } | RelationShape::CoGrouped { keys, .. } => {
+                        keys.len()
+                    }
+                    RelationShape::Map | RelationShape::Aggregate { .. } => 0,
+                };
+                self.relation_lane(&input.schema, udf, keys, &conf)?
+            }
+            PhysNode::PythonAggregate {
+                layout,
+                key_values,
+                calls,
+                ..
+            } => self.aggregate_lanes(&input.schema, layout, key_values, calls, &conf)?,
+            _ => return Err(internal("not a Python relation node".to_owned())),
+        };
+        let sink = RelationSink::try_new(
+            Arc::clone(&input.schema),
+            shape,
+            lanes,
+            schema,
+            RelationLimits {
+                max_rows: conf.max_records_per_batch,
+                max_bytes: conf.max_bytes_per_batch,
+                window: DEFAULT_WINDOW,
+            },
+            self.spill.clone(),
+        )?;
+        let source = SharedSource::new(sink.schema());
+        Ok(Built {
+            schema: sink.schema(),
+            node: Arc::new(BreakerNode {
+                input: input.node,
+                sink,
+                source,
+                ordered: false,
+            }),
+        })
+    }
+
+    /// The runner configuration of this worker's session.
+    fn runner_conf(&self) -> RunnerConf {
+        let config = self.worker.config.as_deref();
+        let max_rows = config
+            .and_then(|c| c.get(&ARROW_MAX_RECORDS_PER_BATCH))
+            .and_then(krishiv_common::config::ConfigValue::as_int)
+            .unwrap_or(10_000);
+        let max_bytes = config
+            .and_then(|c| c.get_unregistered(MAX_BYTES_PER_BATCH))
+            .and_then(|v| v.trim().parse::<i64>().ok())
+            // Spark's default, `Int.MaxValue`, is no limit
+            // (PythonArrowInput.scala:227 at v4.2.0).
+            .filter(|&v| v != i64::from(i32::MAX))
+            .unwrap_or(0);
+        let mut settings = std::collections::BTreeMap::new();
+        if let Some(c) = config {
+            for (k, v) in c.explicit_registered() {
+                settings.insert(k.to_owned(), v);
+            }
+            for (k, v) in c.explicit_unregistered() {
+                settings.insert(k.to_owned(), v.to_owned());
+            }
+        }
+        RunnerConf {
+            max_records_per_batch: usize::try_from(max_rows).unwrap_or(0),
+            max_bytes_per_batch: usize::try_from(max_bytes).unwrap_or(0),
+            settings,
+        }
+    }
+
+    /// The prepared relation call site of `spec`.
+    fn prepare_relation(
+        &self,
+        udf: &PythonUdfRef,
+        spec: &RelationCallSpec,
+    ) -> Result<Arc<dyn PreparedRelationUdf>> {
+        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<RelationLane>, SchemaRef)> {
+        let canonical: Vec<Field> = plan
+            .output
+            .iter()
+            .map(|c| Field::new(&c.name, c.data_type.to_arrow(), c.nullable))
+            .collect();
+        let client: Vec<ArrowType> = canonical
+            .iter()
+            .map(|f| client_arrow_type(f.data_type()))
+            .collect();
+        let client_fields: Vec<Field> = 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<dyn ScalarExpr>)
+                })
+                .collect::<Result<Vec<_>>>()?;
+            sides.push(columns);
+        }
+        let spec = RelationCallSpec {
+            name: plan.udf.name.clone(),
+            version: plan.udf.version,
+            eval_type: plan.udf.eval_type,
+            thread_safe: plan.udf.thread_safe,
+            input_types: plan.sides.iter().map(|s| s.types.clone()).collect(),
+            input_names: plan.sides.iter().map(|s| s.names.clone()).collect(),
+            key_count,
+            arity: plan
+                .sides
+                .first()
+                .map_or(0, |s| s.columns.len().saturating_sub(key_count)),
+            output_names: plan.output.iter().map(|c| c.name.clone()).collect(),
+            output_types: plan.output.iter().map(|c| c.data_type.clone()).collect(),
+            output_arrow: Arc::new(Schema::new(client_fields)),
+            conf: conf.clone(),
+        };
+        let udf = self.prepare_relation(&plan.udf, &spec)?;
+        let lane = RelationLane {
+            name: plan.udf.name.clone(),
+            udf,
+            sides,
+            names: plan.sides.iter().map(|s| s.names.clone()).collect(),
+            client,
+        };
+        Ok((vec![lane], Arc::new(Schema::new(canonical))))
+    }
+
+    /// One lane per call of a Python aggregate, and the node's schema: the
+    /// key columns, then one column per call.
+    fn aggregate_lanes(
+        &self,
+        input: &SchemaRef,
+        layout: &[AttrId],
+        keys: &[usize],
+        calls: &[PythonCallPlan],
+        conf: &RunnerConf,
+    ) -> Result<(Vec<RelationLane>, SchemaRef)> {
+        let values = Values {
+            env: &self.env,
+            subqueries: &[],
+        };
+        let mut fields: Vec<Field> = keys
+            .iter()
+            .map(|&k| {
+                input
+                    .fields()
+                    .get(k)
+                    .map(|f| (**f).clone())
+                    .ok_or_else(|| internal(format!("key column {k} out of range")))
+            })
+            .collect::<Result<_>>()?;
+        let mut lanes = Vec::with_capacity(calls.len());
+        for (i, c) in calls.iter().enumerate() {
+            let ExprNode::Resolved(ResolvedExpr::Function(f)) = &c.call.node else {
+                return Err(internal(
+                    "a Python aggregate call is not a function".to_owned(),
+                ));
+            };
+            let Some(udf) = PythonUdfRef::of(&f.id, f.deterministic) else {
+                return Err(internal(
+                    "a Python aggregate call is not a Python UDF".to_owned(),
+                ));
+            };
+            let args: Vec<&Expr> = f.args.iter().map(AsRef::as_ref).collect();
+            let args = compile_all(&args, layout, input, &values)?;
+            let arg_types = f
+                .args
+                .iter()
+                .map(|a| {
+                    a.resolved_type().ok_or_else(|| {
+                        internal("a Python UDF argument has no resolved type".to_owned())
+                    })
+                })
+                .collect::<Result<Vec<_>>>()?;
+            let canonical = f.data_type.to_arrow();
+            let client = client_arrow_type(&canonical);
+            let name = format!("pythonUDAF{i}");
+            let arg_names: Vec<String> = (0..args.len()).map(|j| format!("_{j}")).collect();
+            let spec = RelationCallSpec {
+                name: udf.name.clone(),
+                version: udf.version,
+                eval_type: udf.eval_type,
+                thread_safe: udf.thread_safe,
+                arity: arg_types.len(),
+                input_types: vec![arg_types],
+                input_names: vec![arg_names.clone()],
+                key_count: 0,
+                output_names: vec![name.clone()],
+                output_types: vec![f.data_type.clone()],
+                output_arrow: Arc::new(Schema::new(vec![Field::new(&name, client.clone(), true)])),
+                conf: conf.clone(),
+            };
+            let prepared = self.prepare_relation(&udf, &spec)?;
+            fields.push(Field::new(&name, canonical, true));
+            lanes.push(RelationLane {
+                name: udf.name.clone(),
+                udf: prepared,
+                names: vec![arg_names],
+                sides: vec![
+                    args.into_iter()
+                        .map(|e| Arc::new(Compiled(e)) as Arc<dyn ScalarExpr>)
+                        .collect(),
+                ],
+                client: vec![client],
+            });
+        }
+        Ok((lanes, Arc::new(Schema::new(fields))))
+    }
+}