From c1ba4d6980ee02cd1ac44825d92d2d6026c0d619 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 19:41:06 +0000 Subject: [PATCH 01/10] plan, exec, optimizer, plan-codec: nested literal values (Array, KeyValueMap, Struct) and codec feature nested_literals (P3-T012a part 2a, step 1) Literal::{Array, KeyValueMap, Struct} with Arc types and checked constructors (element types, nullability and depth against MAX_NESTING_DEPTH), handled at the sites of the design note section 5, literal_array recursive over Arrow, canonical encoding tags 20-22, plan-codec feature nested_literals (requires nested_types) with goldens and reader-refusal tests. Names and explain text for nested literals are golden-pending. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-exec/src/planner/eval/mod.rs | 169 +++++++++++++++ crates/krishiv-exec/src/planner/pushdown.rs | 2 + crates/krishiv-optimizer/src/fold.rs | 42 ++++ crates/krishiv-plan-codec/src/decode.rs | 54 +++++ crates/krishiv-plan-codec/src/encode.rs | 44 ++++ crates/krishiv-plan-codec/src/features.rs | 41 ++++ .../src/proto/krishiv.plan.v1.pb.rs | 48 ++++- .../tests/common/inventory.rs | 45 ++++ crates/krishiv-plan-codec/tests/features.rs | 104 +++++++++ crates/krishiv-plan-codec/tests/golden.rs | 7 +- .../golden/nested_types/nested_literals.bin | Bin 0 -> 1581 bytes crates/krishiv-plan-codec/tests/roundtrip.rs | 3 +- .../src/analyzer/exprs/literals.rs | 29 +-- .../krishiv-plan/src/analyzer/exprs/pretty.rs | 10 +- crates/krishiv-plan/src/analyzer/fold.rs | 11 +- crates/krishiv-plan/src/display.rs | 32 +++ crates/krishiv-plan/src/encode.rs | 30 +++ crates/krishiv-plan/src/literal.rs | 203 +++++++++++++++++- crates/krishiv-plan/src/raises.rs | 12 +- crates/krishiv-plan/src/tests/mod.rs | 1 + .../krishiv-plan/src/tests/nested_literals.rs | 189 ++++++++++++++++ crates/krishiv-plan/src/tests/props.rs | 24 +++ .../krishiv-session/src/execute_immediate.rs | 3 + proto/krishiv/plan/v1/plan.proto | 35 +++ 24 files changed, 1102 insertions(+), 36 deletions(-) create mode 100644 crates/krishiv-plan-codec/tests/golden/nested_types/nested_literals.bin create mode 100644 crates/krishiv-plan/src/tests/nested_literals.rs diff --git a/crates/krishiv-exec/src/planner/eval/mod.rs b/crates/krishiv-exec/src/planner/eval/mod.rs index 1831e73e6..19c0dd870 100644 --- a/crates/krishiv-exec/src/planner/eval/mod.rs +++ b/crates/krishiv-exec/src/planner/eval/mod.rs @@ -1090,10 +1090,88 @@ pub(crate) fn literal_array(l: &Literal) -> Result { .map_err(|e| internal(&e.to_string()))?, ) } + Literal::Array { + data_type, + elements, + } => { + let ArrowType::List(field) = DataType::Array(Arc::clone(data_type)).to_arrow() else { + return Err(internal("an array literal's Arrow type")); + }; + let values = concat_literals(elements.iter(), field.data_type())?; + let len = i32::try_from(values.len()) + .map_err(|_| internal("an array literal past the 32-bit offset limit"))?; + Arc::new( + arrow_array::ListArray::try_new( + field, + OffsetBuffer::new(ScalarBuffer::from(vec![0, len])), + values, + None, + ) + .map_err(|e| internal(&e.to_string()))?, + ) + } + Literal::KeyValueMap { data_type, entries } => { + let ArrowType::Map(entries_field, sorted) = + DataType::Map(Arc::clone(data_type)).to_arrow() + else { + return Err(internal("a map literal's Arrow type")); + }; + let ArrowType::Struct(kv) = entries_field.data_type() else { + return Err(internal("a map literal's entries type")); + }; + let [key_field, value_field] = &kv[..] else { + return Err(internal("a map literal's entries type")); + }; + let keys = concat_literals(entries.iter().map(|(k, _)| k), key_field.data_type())?; + let values = concat_literals(entries.iter().map(|(_, v)| v), value_field.data_type())?; + let len = i32::try_from(keys.len()) + .map_err(|_| internal("a map literal past the 32-bit offset limit"))?; + let pairs = StructArray::try_new(kv.clone(), vec![keys, values], None) + .map_err(|e| internal(&e.to_string()))?; + Arc::new( + MapArray::try_new( + entries_field, + OffsetBuffer::new(ScalarBuffer::from(vec![0, len])), + pairs, + None, + sorted, + ) + .map_err(|e| internal(&e.to_string()))?, + ) + } + Literal::Struct { data_type, fields } => { + let ArrowType::Struct(arrow_fields) = + DataType::Struct((**data_type).clone()).to_arrow() + else { + return Err(internal("a struct literal's Arrow type")); + }; + let columns = fields + .iter() + .map(literal_array) + .collect::>>()?; + Arc::new( + StructArray::try_new(arrow_fields, columns, None) + .map_err(|e| internal(&e.to_string()))?, + ) + } _ => return Err(unsupported("interval literals in physical plans")), }) } +/// The one-row arrays of `values` joined into one array of `data_type` +/// (an empty one when there are none). +fn concat_literals<'a>( + values: impl Iterator, + data_type: &ArrowType, +) -> Result { + let arrays = values.map(literal_array).collect::>>()?; + if arrays.is_empty() { + return Ok(arrow_array::new_empty_array(data_type)); + } + let refs: Vec<&dyn arrow_array::Array> = arrays.iter().map(AsRef::as_ref).collect(); + arrow_select::concat::concat(&refs).map_err(|e| internal(&e.to_string())) +} + /// A compiled expression as the operators' [`Predicate`] and /// [`ScalarExpr`]. #[derive(Debug, Clone)] @@ -1201,6 +1279,97 @@ mod tests { assert_eq!(values, [Some("FAILFAST"), Some("dd/MM/yyyy")]); } + /// Nested literals build one-row list, map and struct arrays of exactly + /// their type's Arrow form: empty, with a null element, and nested. + #[test] + fn nested_literals_build_one_row_arrays_of_their_arrow_type() { + use arrow_array::cast::AsArray; + use arrow_array::{Array, ListArray}; + use krishiv_types::{ArrayType, DataType, MapType, StructField, StructType}; + + let ints = |values: Vec| { + Literal::array(ArrayType::new(DataType::Integer, true), values).unwrap() + }; + let array_type: DataType = ArrayType::new(DataType::Integer, true).into(); + + // Empty. + let empty = super::literal_array(&ints(vec![])).unwrap(); + assert_eq!(empty.data_type(), &array_type.to_arrow()); + assert_eq!((empty.len(), empty.as_list::().value(0).len()), (1, 0)); + + // A null element. + let a = super::literal_array(&ints(vec![ + Literal::Integer(1), + Literal::Null(Some(DataType::Integer)), + Literal::Integer(3), + ])) + .unwrap(); + assert_eq!(a.data_type(), &array_type.to_arrow()); + let list: &ListArray = a.as_list(); + let row = list.value(0); + let values: Vec> = row + .as_primitive::() + .iter() + .collect(); + assert_eq!(values, [Some(1), None, Some(3)]); + + // Array of arrays, strings inside. + let outer_type = ArrayType::new(array_type.clone(), true); + let outer = Literal::array( + outer_type.clone(), + vec![ + ints(vec![Literal::Integer(7)]), + Literal::Null(Some(array_type)), + ints(vec![]), + ], + ) + .unwrap(); + let a = super::literal_array(&outer).unwrap(); + assert_eq!(a.data_type(), &DataType::from(outer_type).to_arrow()); + let inner = a.as_list::().value(0); + assert_eq!(inner.len(), 3); + assert!(inner.is_null(1)); + + // A map with a struct value. + let struct_type = StructType::new([ + StructField::new("a", DataType::Integer, false), + StructField::new("s", DataType::String, true), + ]); + let map_type = MapType::new(DataType::String, struct_type.clone().into(), true); + let entry = |k: &str, v: i32| { + ( + Literal::String(k.to_owned()), + Literal::struct_value( + struct_type.clone(), + vec![Literal::Integer(v), Literal::Null(Some(DataType::String))], + ) + .unwrap(), + ) + }; + let map = + Literal::key_value_map(map_type.clone(), vec![entry("x", 1), entry("y", 2)]).unwrap(); + let a = super::literal_array(&map).unwrap(); + assert_eq!(a.data_type(), &DataType::from(map_type).to_arrow()); + let m = a.as_map(); + assert_eq!((m.len(), m.value_length(0)), (1, 2)); + + // An empty map and a bare struct. + let empty_map = Literal::key_value_map( + MapType::new(DataType::Integer, DataType::Integer, true), + vec![], + ) + .unwrap(); + assert_eq!(super::literal_array(&empty_map).unwrap().len(), 1); + let st = Literal::struct_value( + struct_type.clone(), + vec![Literal::Integer(5), Literal::String("z".to_owned())], + ) + .unwrap(); + let a = super::literal_array(&st).unwrap(); + assert_eq!(a.data_type(), &DataType::from(struct_type).to_arrow()); + assert_eq!(a.as_struct().num_columns(), 2); + } + /// `scatter` puts the i-th value at the i-th set row, NULL elsewhere, /// keeping the values' own NULLs, for booleans and strings alike. #[test] diff --git a/crates/krishiv-exec/src/planner/pushdown.rs b/crates/krishiv-exec/src/planner/pushdown.rs index 945b320c1..8adca3a3c 100644 --- a/crates/krishiv-exec/src/planner/pushdown.rs +++ b/crates/krishiv-exec/src/planner/pushdown.rs @@ -243,6 +243,8 @@ fn value_or_null(e: &Expr) -> Option> { Literal::Binary(v) => ScalarValue::Binary(v.clone()), Literal::Date(v) => ScalarValue::Date32(*v), Literal::Timestamp(v) | Literal::TimestampNtz(v) => ScalarValue::TimestampMicros(*v), + // No scalar form for the file filters, so never pushed down: nested + // values (`Literal::{Array, KeyValueMap, Struct}`), maps, intervals. _ => return None, })) } diff --git a/crates/krishiv-optimizer/src/fold.rs b/crates/krishiv-optimizer/src/fold.rs index 24505be53..1dbbe6c14 100644 --- a/crates/krishiv-optimizer/src/fold.rs +++ b/crates/krishiv-optimizer/src/fold.rs @@ -245,6 +245,8 @@ pub(crate) fn compare_values(x: &Literal, y: &Literal) -> Option { |v: i128, from: u8| v.checked_mul(10_i128.checked_pow(u32::from(scale - from))?); up(*a, s.scale())?.cmp(&up(*b, t.scale())?) } + // Not foldable: nested values (`Literal::{Array, KeyValueMap, + // Struct}`) are not ordered here, nor are mixed types. _ => return None, }) } @@ -363,6 +365,9 @@ pub(crate) fn literal_type(l: &Literal) -> Option { Literal::Date(_) => DataType::Date, Literal::Timestamp(_) => DataType::Timestamp, Literal::TimestampNtz(_) => DataType::TimestampNtz, + Literal::Array { .. } | Literal::KeyValueMap { .. } | Literal::Struct { .. } => { + return l.data_type(); + } _ => return None, }) } @@ -401,6 +406,9 @@ fn to_array(l: &Literal) -> Option { Literal::String(v) => Arc::new(StringArray::from(vec![v.as_str()])), Literal::Date(v) => Arc::new(Date32Array::from(vec![*v])), Literal::TimestampNtz(v) => Arc::new(TimestampMicrosecondArray::from(vec![*v])), + // Not foldable: `Literal::{Array, KeyValueMap, Struct}` have no + // one-row form here, so a cast or call over a nested value stays as + // written. _ => return None, }) } @@ -723,6 +731,8 @@ mod tests { use krishiv_plan::{Builtin, CastMode, Literal}; use krishiv_types::{DataType, DecimalType}; + use std::sync::Arc; + use super::{arith, compare_values, fold_cast, lit, negate}; fn dec(unscaled: i128, p: u32, s: u32) -> Literal { @@ -732,6 +742,38 @@ mod tests { } } + #[test] + fn nested_literals_are_left_intact() { + use krishiv_plan::Expr; + use krishiv_types::ArrayType; + + let array = |v: i32| { + Literal::array( + ArrayType::new(DataType::Integer, true), + vec![Literal::Integer(v)], + ) + .unwrap() + }; + // They have no ordering and no one-row form here, so nothing over + // them folds, but their type is known. + assert_eq!(compare_values(&array(1), &array(1)), None); + assert_eq!( + fold_cast(&array(1), &DataType::String, CastMode::Ansi), + None + ); + assert_eq!( + super::literal_type(&array(1)), + Some(ArrayType::new(DataType::Integer, true).into()) + ); + // `array(1) = array(1)` stays a call; the literals are not touched. + let eq = Arc::new(Expr::builtin( + Builtin::Eq, + vec![Expr::from(array(1)), Expr::from(array(1))], + )); + let folded = super::fold(&eq); + assert!(Arc::ptr_eq(&eq, &folded) || *folded == *eq); + } + #[test] fn casts_fold_through_the_kernels() { let d52 = DataType::Decimal(DecimalType::try_new(5, 2).unwrap()); diff --git a/crates/krishiv-plan-codec/src/decode.rs b/crates/krishiv-plan-codec/src/decode.rs index 494c61d5a..303099b99 100644 --- a/crates/krishiv-plan-codec/src/decode.rs +++ b/crates/krishiv-plan-codec/src/decode.rs @@ -1629,7 +1629,24 @@ fn narrow>(v: i32) -> Result { } fn literal(l: &pb::Literal) -> Result { + literal_at(l, 0) +} + +/// `l`, `depth` literals inside a nested literal: the type bounds the depth of +/// a valid one, and the bound keeps a hostile chain from recursing further. +fn literal_at(l: &pb::Literal, depth: usize) -> Result { use pb::literal::Kind; + if depth > DataType::MAX_NESTING_DEPTH { + return Err(corrupt("a nested literal deeper than a type may be")); + } + let member = |m: &pb::Literal| { + // The options map is its own feature and never a member (the encoder + // refuses it). + if matches!(m.kind, Some(Kind::Map(_))) { + return Err(corrupt("an options map inside a nested literal")); + } + literal_at(m, depth + 1) + }; Ok(match required(&l.kind, "a literal kind")? { Kind::Null(n) => Literal::Null(n.data_type.as_ref().map(spark_type).transpose()?), Kind::Boolean(v) => Literal::Boolean(*v), @@ -1693,6 +1710,43 @@ fn literal(l: &pb::Literal) -> Result { .map(|e| (e.key.clone(), e.value.clone())) .collect(), ), + Kind::Array(v) => { + let DataType::Array(t) = spark_type(required(&v.data_type, "an array literal type")?)? + else { + return Err(corrupt("an array literal whose type is not an array")); + }; + Literal::array( + (*t).clone(), + v.elements.iter().map(member).collect::>()?, + ) + .map_err(|_| corrupt("an array literal that does not fit its type"))? + } + Kind::KeyValueMap(v) => { + let DataType::Map(t) = spark_type(required(&v.data_type, "a map literal type")?)? + else { + return Err(corrupt("a map literal whose type is not a map")); + }; + let entries = v + .entries + .iter() + .map(|e| { + Ok(( + member(required(&e.key, "a map literal key")?)?, + member(required(&e.value, "a map literal value")?)?, + )) + }) + .collect::>()?; + Literal::key_value_map((*t).clone(), entries) + .map_err(|_| corrupt("a map literal that does not fit its type"))? + } + Kind::Struct(v) => { + let DataType::Struct(t) = spark_type(required(&v.data_type, "a struct literal type")?)? + else { + return Err(corrupt("a struct literal whose type is not a struct")); + }; + Literal::struct_value(t, v.fields.iter().map(member).collect::>()?) + .map_err(|_| corrupt("a struct literal that does not fit its type"))? + } }) } diff --git a/crates/krishiv-plan-codec/src/encode.rs b/crates/krishiv-plan-codec/src/encode.rs index 1088ad88c..933558847 100644 --- a/crates/krishiv-plan-codec/src/encode.rs +++ b/crates/krishiv-plan-codec/src/encode.rs @@ -1558,6 +1558,38 @@ fn literal(l: &Literal, nested: &Cell) -> Result { .map(|(key, value)| map_entry(key, value)) .collect::>()?, }), + Literal::Array { + data_type, + elements, + } => Kind::Array(pb::ArrayLiteral { + data_type: Some(spark_type(&DataType::Array(Arc::clone(data_type)), nested)?), + elements: elements + .iter() + .map(|e| nested_member(e, nested)) + .collect::>()?, + }), + Literal::KeyValueMap { data_type, entries } => Kind::KeyValueMap(pb::KeyValueMapLiteral { + data_type: Some(spark_type(&DataType::Map(Arc::clone(data_type)), nested)?), + entries: entries + .iter() + .map(|(k, v)| { + Ok(pb::KeyValueEntry { + key: Some(nested_member(k, nested)?), + value: Some(nested_member(v, nested)?), + }) + }) + .collect::>()?, + }), + Literal::Struct { data_type, fields } => Kind::Struct(pb::StructLiteral { + data_type: Some(spark_type( + &DataType::Struct((**data_type).clone()), + nested, + )?), + fields: fields + .iter() + .map(|e| nested_member(e, nested)) + .collect::>()?, + }), _ => { return Err(CodecError::Unsupported( "a literal kind this version does not know".to_owned(), @@ -1567,6 +1599,18 @@ fn literal(l: &Literal, nested: &Cell) -> Result { Ok(pb::Literal { kind: Some(kind) }) } +/// A value inside a nested literal. The options map is its own feature +/// (`options_map_literals`), which the walk reports for an expression, not for +/// a member, so it does not travel inside one. +fn nested_member(l: &Literal, nested: &Cell) -> Result { + if matches!(l, Literal::Map(_)) { + return Err(CodecError::Unsupported( + "an options map inside an array, map or struct literal".to_owned(), + )); + } + literal(l, nested) +} + fn year_month_field(f: YearMonthField) -> pb::YearMonthField { match f { YearMonthField::Year => pb::YearMonthField::Year, diff --git a/crates/krishiv-plan-codec/src/features.rs b/crates/krishiv-plan-codec/src/features.rs index f811db498..728be1db2 100644 --- a/crates/krishiv-plan-codec/src/features.rs +++ b/crates/krishiv-plan-codec/src/features.rs @@ -148,6 +148,7 @@ registry! { (map_key_dedup_policy, MapKeyDedupPolicy, POST_MIGRATION_READER_VERSION, Semantic, "The session's map-key dedup policy (`LAST_WIN`); without it a worker keeps the first value on duplicate map keys."), (mode_regr_bit_listagg_aggs, ModeRegrBitListaggAggs, POST_MIGRATION_READER_VERSION, Semantic, "The `mode`, regression (`regr_*`), bitwise (`bit_and`, `bit_or`, `bit_xor`), `listagg`/`string_agg` and `try_sum`/`try_avg` aggregates."), (moment_aggregates, MomentAggregates, POST_MIGRATION_READER_VERSION, Semantic, "The moment aggregates (`stddev_samp`, `stddev_pop`, `var_samp`, `var_pop`, `skewness`, `kurtosis`, `corr`, `covar_samp`, `covar_pop`)."), + (nested_literals, NestedLiterals, POST_MIGRATION_READER_VERSION, Semantic, "An array, map or struct literal value (`Literal::{Array, KeyValueMap, Struct}`). The literal's type is a nested type, so a plan with one lists `nested_types` too."), (nested_types, NestedTypes, POST_MIGRATION_READER_VERSION, Semantic, "Array, map and struct types (`DataType.KIND_NESTED`, written as Spark's JSON form of the type)."), (options_map_literals, OptionsMapLiterals, POST_MIGRATION_READER_VERSION, Semantic, "A `MAP` literal: the options of one call (`from_json`, `from_csv`, `from_xml`). Entries are classified as source options are: a sensitive key or a URL with credentials does not encode (ADR 0038)."), (partial_aggregate, PartialAggregate, POST_MIGRATION_READER_VERSION, Optimization, "Partial aggregates below a join, whose states a later aggregate merges (eager aggregation)."), @@ -236,6 +237,10 @@ pub enum Construct<'a> { EvalMode, /// A `WritePlan` naming the Spark types of input columns. WriteColumnTypes, + /// The type of a nested literal (array, map or struct): a nested type + /// the literal carries itself, so `nested_types` is listed with + /// `nested_literals`. + NestedLiteralType, /// A `Generate` node's generator. Generator(krishiv_plan::GeneratorKind), /// A hash join's probe selection (`HashJoinPlan.select`). @@ -480,6 +485,7 @@ pub fn feature_of(construct: Construct<'_>) -> Option { Construct::Lambda => Some(Feature::LambdaFunctions), Construct::EvalMode => Some(Feature::CapturedEvalMode), Construct::WriteColumnTypes => Some(Feature::WriteColumnTypes), + Construct::NestedLiteralType => Some(Feature::NestedTypes), Construct::ProbeSelect => Some(Feature::FusedProbe), Construct::Generator(kind) => match kind { krishiv_plan::GeneratorKind::VariantExplode => Some(Feature::VariantExplode), @@ -519,6 +525,9 @@ pub fn feature_of(construct: Construct<'_>) -> Option { ExprNode::Resolved(ResolvedExpr::Field { .. }) => Some(Feature::StructField), ExprNode::Literal(Literal::CalendarInterval { .. }) => Some(Feature::IntervalTypes), ExprNode::Literal(Literal::Map(_)) => Some(Feature::OptionsMapLiterals), + ExprNode::Literal( + Literal::Array { .. } | Literal::KeyValueMap { .. } | Literal::Struct { .. }, + ) => Some(Feature::NestedLiterals), // The year-month and day-time literal kinds predate the feature, // so a bare one stays baseline: whatever reads its column names // the interval type. As an argument nothing else may name it, and @@ -735,12 +744,44 @@ fn constructs_of<'a>(node: &'a PhysNode, f: &mut impl FnMut(Construct<'a>)) { f(Construct::EvalMode); } f(Construct::Expr(e)); + nested_literal_types(e, f); // A UDF call hides in any expression of the node (`exprs_of` // already yields every descendant, so one node check is enough). functions_in(e, f); } } +/// A nested literal's type as constructs: the literal itself is a nested type +/// (`nested_types`), and the types inside it are visited as any type is +/// (`array` needs `time_type` too). +fn nested_literal_types<'a>(e: &'a Expr, f: &mut impl FnMut(Construct<'a>)) { + let ExprNode::Literal(l) = &e.node else { + return; + }; + match l { + Literal::Array { data_type, .. } => { + f(Construct::NestedLiteralType); + visit_type(data_type.element_type(), f); + } + Literal::KeyValueMap { data_type, .. } => { + f(Construct::NestedLiteralType); + visit_type(data_type.key_type(), f); + visit_type(data_type.value_type(), f); + } + Literal::Struct { data_type, .. } => { + f(Construct::NestedLiteralType); + data_type + .fields() + .iter() + .for_each(|x| visit_type(x.data_type(), f)); + } + // Scalars and the options map have no type of their own to visit; + // `Literal` is non-exhaustive, and a kind added later fails to encode + // (`encode::literal`) until it is handled there. + _ => {} + } +} + /// The Arrow types `node` carries as IPC bytes. fn arrow_types_of(node: &PhysNode) -> Vec<&arrow_schema::DataType> { fn window<'a>(function: &'a WindowFunction, out: &mut Vec<&'a arrow_schema::DataType>) { diff --git a/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs b/crates/krishiv-plan-codec/src/proto/krishiv.plan.v1.pb.rs index df1dcdee4..133eaf351 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 @@ -1531,7 +1531,7 @@ pub mod data_type { pub struct Literal { #[prost( oneof = "literal::Kind", - tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18" + tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21" )] pub kind: ::core::option::Option, } @@ -1577,8 +1577,54 @@ pub mod literal { /// feature: options_map_literals #[prost(message, tag = "18")] Map(super::MapLiteral), + /// feature: nested_literals + #[prost(message, tag = "19")] + Array(super::ArrayLiteral), + /// feature: nested_literals + #[prost(message, tag = "20")] + KeyValueMap(super::KeyValueMapLiteral), + /// feature: nested_literals + #[prost(message, tag = "21")] + Struct(super::StructLiteral), } } +/// feature: nested_literals +/// An array value of any element type: `data_type` is the array type (a +/// nested type, so `nested_types` too), each element a literal of the element +/// type, a null one a typed `NullLiteral`. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct ArrayLiteral { + #[prost(message, optional, tag = "1")] + pub data_type: ::core::option::Option, + #[prost(message, repeated, tag = "2")] + pub elements: ::prost::alloc::vec::Vec, +} +/// feature: nested_literals +/// A map value of any key and value type (not the options-only `MapLiteral`). +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct KeyValueMapLiteral { + #[prost(message, optional, tag = "1")] + pub data_type: ::core::option::Option, + #[prost(message, repeated, tag = "2")] + pub entries: ::prost::alloc::vec::Vec, +} +/// feature: nested_literals +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct KeyValueEntry { + #[prost(message, optional, tag = "1")] + pub key: ::core::option::Option, + #[prost(message, optional, tag = "2")] + pub value: ::core::option::Option, +} +/// feature: nested_literals +/// A struct value: one literal per field of `data_type`, in field order. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct StructLiteral { + #[prost(message, optional, tag = "1")] + pub data_type: ::core::option::Option, + #[prost(message, repeated, tag = "2")] + pub fields: ::prost::alloc::vec::Vec, +} /// feature: options_map_literals /// A `MAP` literal: one call's options (`from_json`, /// `from_csv`, `from_xml` and the like), entries in their given order. diff --git a/crates/krishiv-plan-codec/tests/common/inventory.rs b/crates/krishiv-plan-codec/tests/common/inventory.rs index bb4a262b8..6c0a580a7 100644 --- a/crates/krishiv-plan-codec/tests/common/inventory.rs +++ b/crates/krishiv-plan-codec/tests/common/inventory.rs @@ -2250,6 +2250,50 @@ pub(crate) fn typed() -> Vec<(&'static str, PhysNode)> { nullable: None, subqueries: Vec::new(), }; + // Nested values: an array with a null element, an empty one, a map with + // a struct value and a bare struct. + let nested_literals = { + use krishiv_types::{ArrayType, MapType, StructField, StructType}; + let array_type = ArrayType::new(DataType::Integer, true); + let struct_type = StructType::new([ + StructField::new("a", DataType::Integer, false), + StructField::new("s", DataType::String, true), + ]); + let a_struct = |v: i32| { + Literal::struct_value( + struct_type.clone(), + vec![Literal::Integer(v), Literal::Null(Some(DataType::String))], + ) + .unwrap() + }; + PhysNode::Project { + input: boxed(scan()), + layout: vec![AttrId(1), AttrId(2)], + items: [ + Literal::array( + array_type.clone(), + vec![Literal::Integer(1), Literal::Null(Some(DataType::Integer))], + ) + .unwrap(), + Literal::array(array_type, Vec::new()).unwrap(), + Literal::key_value_map( + MapType::new(DataType::String, struct_type.clone().into(), true), + vec![(Literal::String("k".to_owned()), a_struct(1))], + ) + .unwrap(), + a_struct(2), + ] + .into_iter() + .enumerate() + .map(|(i, l)| ProjectItemPlan::Expr { + name: format!("n{i}"), + expr: lit(l), + }) + .collect(), + nullable: None, + subqueries: Vec::new(), + } + }; let mut interval_columns = file_scan(); interval_columns.scan.output = vec![ column( @@ -2308,6 +2352,7 @@ pub(crate) fn typed() -> Vec<(&'static str, PhysNode)> { vec![ ("interval_literals", interval_literals), ("options_map_literal", options_map), + ("nested_literals", nested_literals), ("interval_scan", PhysNode::Scan(interval_columns)), ("interval_nested", nested_intervals), ( diff --git a/crates/krishiv-plan-codec/tests/features.rs b/crates/krishiv-plan-codec/tests/features.rs index 42745b98b..23e67f8e3 100644 --- a/crates/krishiv-plan-codec/tests/features.rs +++ b/crates/krishiv-plan-codec/tests/features.rs @@ -112,6 +112,11 @@ fn the_header_lists_exactly_the_features_in_use() { message("options_map_literal").features, ["options_map_literals"] ); + // The literal's own type is nested, so a reader needs both. + assert_eq!( + message("nested_literals").features, + ["nested_literals", "nested_types"] + ); assert_eq!( message("aggregate_any_value").features, ["any_value_by_aggregates"] @@ -2122,3 +2127,102 @@ fn a_sample_needs_its_feature() { } ); } + +// ---- nested literals (nested_literals) ------------------------------------ + +#[test] +fn nested_literals_need_their_feature_and_their_type_feature() { + let m = message("nested_literals"); + assert_eq!(m.features, ["nested_literals", "nested_types"]); + assert_eq!(Feature::NestedLiterals.kind(), FeatureKind::Semantic); + let bytes = m.encode_to_vec(); + // A reader that has the nested types but not the literals refuses the + // plan naming the literals; the reverse names the types. + let without_literals = ReaderFeatures::only([Feature::NestedTypes]); + assert_eq!( + decode_with(&bytes, &Limits::default(), &without_literals).unwrap_err(), + CodecError::UnsupportedFeatures { + missing: vec!["nested_literals".to_owned()] + } + ); + let without_types = ReaderFeatures::only([Feature::NestedLiterals]); + assert_eq!( + decode_with(&bytes, &Limits::default(), &without_types).unwrap_err(), + CodecError::UnsupportedFeatures { + missing: vec!["nested_types".to_owned()] + } + ); + let both = ReaderFeatures::only([Feature::NestedLiterals, Feature::NestedTypes]); + let plan = decode_with(&bytes, &Limits::default(), &both).unwrap(); + assert_eq!(encode(&plan).unwrap(), bytes); +} + +#[test] +fn a_nested_literal_is_classified_as_an_expression_and_by_its_type() { + use krishiv_plan::{Expr, Literal}; + use krishiv_types::{ArrayType, DataType}; + let l = Literal::array(ArrayType::new(DataType::Integer, true), Vec::new()).unwrap(); + assert_eq!( + feature_of(Construct::Expr(&Expr::from(l))), + Some(Feature::NestedLiterals) + ); + assert_eq!( + feature_of(Construct::NestedLiteralType), + Some(Feature::NestedTypes) + ); +} + +/// Every literal of the plan's expression table, mutably. +fn literals_mut(m: &mut pb::PhysicalPlan) -> Vec<&mut pb::Literal> { + m.exprs + .iter_mut() + .filter_map(|e| match &mut e.kind { + Some(pb::expr::Kind::Literal(l)) => Some(l), + _ => None, + }) + .collect() +} + +#[test] +fn a_nested_literal_that_does_not_fit_its_type_is_corrupt() { + use pb::literal::Kind; + let int_type = pb::DataType { + kind: pb::data_type::Kind::Integer as i32, + ..pb::DataType::default() + }; + // The array type replaced by a scalar one. + let mut m = message("nested_literals"); + for l in literals_mut(&mut m) { + if let Some(Kind::Array(a)) = &mut l.kind { + a.data_type = Some(int_type.clone()); + } + } + corrupt(&m, "an array literal whose type is not an array"); + // An element of another type. + let mut m = message("nested_literals"); + for l in literals_mut(&mut m) { + if let Some(Kind::Array(a)) = &mut l.kind { + for e in &mut a.elements { + e.kind = Some(Kind::String("x".to_owned())); + } + } + } + corrupt(&m, "an array literal that does not fit its type"); + // A struct with a missing value. + let mut m = message("nested_literals"); + for l in literals_mut(&mut m) { + if let Some(Kind::Struct(s)) = &mut l.kind { + s.fields.pop(); + } + } + corrupt(&m, "a struct literal that does not fit its type"); + // No type at all. + let mut m = message("nested_literals"); + for l in literals_mut(&mut m) { + if let Some(Kind::KeyValueMap(v)) = &mut l.kind { + v.data_type = None; + } + } + let err = decode(&m.encode_to_vec()).unwrap_err(); + assert!(matches!(err, CodecError::Corrupt(_)), "{err:?}"); +} diff --git a/crates/krishiv-plan-codec/tests/golden.rs b/crates/krishiv-plan-codec/tests/golden.rs index dd03a8ed2..27cf9c96f 100644 --- a/crates/krishiv-plan-codec/tests/golden.rs +++ b/crates/krishiv-plan-codec/tests/golden.rs @@ -192,12 +192,15 @@ fn the_registry_is_sorted_and_every_feature_has_goldens() { // aggregate, whose feature sorts after it, so its plans file under // `moment_aggregates` and it has no directory of its own; likewise // `collated_min_max`, whose struct argument files its plans under - // `nested_types`. + // `nested_types`; and `nested_literals`, whose literal's type is a + // nested type too, so its plans file there. assert!( files > 0 || matches!( *f, - Feature::LegacyStatisticalAggregate | Feature::CollatedMinMax + Feature::LegacyStatisticalAggregate + | Feature::CollatedMinMax + | Feature::NestedLiterals ), "{id}: no golden files in {}", dir.display() diff --git a/crates/krishiv-plan-codec/tests/golden/nested_types/nested_literals.bin b/crates/krishiv-plan-codec/tests/golden/nested_types/nested_literals.bin new file mode 100644 index 0000000000000000000000000000000000000000..533bdd59ae9a457133acb59ba6b8af9498f59349 GIT binary patch literal 1581 zcmdT^F>ljA6!x>7hB5;P@qX->u+5wWxE3^jX6?6dMV zPsCQ-PsN3?$(|v=fZICBIbauw&LG82(OIh+6nK?VNNFyJB9YVy)j_w+3N!*@FlDwk zOPZz;oSJ54TGIe`3n&(9OFN~N(WXvWs4FNA3ZBx>p(a`htutd^N5K>KaSVx@r4l(P z%#CsF;D&wCtuZ8F6PBlXbjiwPtIJrX30AiX+3Jc#9H_p~kE5sveM8wrSwvC$2cywQ zQBvoT1$D;1O68GW9;X`mSSzK{R227)`V~|RHB@ROD~s|{ag~pcGAWJrS7|19YZX;0 zHO%px_@%f%X-PT8s28buPqZY1aI)kr22bjg=DeW2C2fHEJT<*ImJHiV@*u>H(DbXu z664NR5J_j7r%uq0;zam&D+H|3SEKUu`RM1DkFWD?fk!+Xq|%!`(p#QgtO3`gKDZI*7_F4-r_PJO9S8==QM?28FFJ&eYFQ=+jIh|o{n#{C^P z;@xqoh-O(I=BEnFABAXd@h}xNw~No`<6ni_*clFMQ|I*1d3yh$ud2%*t`CP9!JEfa z@H0@ugAV%u@8KQLmn!4E!;Jr-Bz}ABqW1~@55{KJp2hWYBvL8-jWCcR!3SP1Wq@jg HbQ697uc2Q~ literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/roundtrip.rs b/crates/krishiv-plan-codec/tests/roundtrip.rs index 8fe380ce2..e1ae05dea 100644 --- a/crates/krishiv-plan-codec/tests/roundtrip.rs +++ b/crates/krishiv-plan-codec/tests/roundtrip.rs @@ -634,7 +634,8 @@ fn readers_that_predate_features_still_refuse_what_they_cannot_read() { | "write_files_partitioned" | "write_files_column_types" | "aggregate_percentiles" - | "options_map_literal" => 6, + | "options_map_literal" + | "nested_literals" => 6, // Features added after the version was frozen: above every // sequential format version. n if n.starts_with("time_") diff --git a/crates/krishiv-plan/src/analyzer/exprs/literals.rs b/crates/krishiv-plan/src/analyzer/exprs/literals.rs index 83633aec5..9f2b4b625 100644 --- a/crates/krishiv-plan/src/analyzer/exprs/literals.rs +++ b/crates/krishiv-plan/src/analyzer/exprs/literals.rs @@ -5,34 +5,7 @@ use super::*; use krishiv_kernels::cast::{TimestampText, timestamp_text}; pub(in crate::analyzer) fn literal_type(v: &Literal) -> Option { - Some(match v { - Literal::Null(t) => return t.clone().filter(|t| *t != DataType::Null), - Literal::Boolean(_) => DataType::Boolean, - Literal::Byte(_) => DataType::Byte, - Literal::Short(_) => DataType::Short, - Literal::Integer(_) => DataType::Integer, - Literal::Long(_) => DataType::Long, - Literal::Float(_) => DataType::Float, - Literal::Double(_) => DataType::Double, - Literal::Decimal { data_type, .. } => DataType::Decimal(*data_type), - Literal::String(_) => DataType::String, - Literal::Map(_) => DataType::Map( - krishiv_types::MapType::new(DataType::String, DataType::String, false).into(), - ), - Literal::Binary(_) => DataType::Binary, - Literal::Date(_) => DataType::Date, - Literal::Timestamp(_) => DataType::Timestamp, - Literal::TimestampNtz(_) => DataType::TimestampNtz, - Literal::YearMonthInterval { start, end, .. } => DataType::YearMonthInterval { - start: *start, - end: *end, - }, - Literal::DayTimeInterval { start, end, .. } => DataType::DayTimeInterval { - start: *start, - end: *end, - }, - Literal::CalendarInterval { .. } => DataType::CalendarInterval, - }) + v.data_type() } /// A typed datetime literal as Spark's `AstBuilder` builds it: the same diff --git a/crates/krishiv-plan/src/analyzer/exprs/pretty.rs b/crates/krishiv-plan/src/analyzer/exprs/pretty.rs index 4d78e0744..7c88b68d3 100644 --- a/crates/krishiv-plan/src/analyzer/exprs/pretty.rs +++ b/crates/krishiv-plan/src/analyzer/exprs/pretty.rs @@ -822,7 +822,15 @@ pub(super) fn pretty_literal(v: &Literal, zone: Option) -> String { crate::spark_fmt::timestamp(SessionZone::local_micros(zone, *us), false) } Literal::TimestampNtz(us) => crate::spark_fmt::timestamp(*us, true), - _ => v.to_string(), + // Nested values name as the display form (`ARRAY(1, 2)`); Spark's own + // name is golden-pending (design P3-T012a, 5). + Literal::Array { .. } + | Literal::KeyValueMap { .. } + | Literal::Struct { .. } + | Literal::Map(_) + | Literal::YearMonthInterval { .. } + | Literal::DayTimeInterval { .. } + | Literal::CalendarInterval { .. } => v.to_string(), } } diff --git a/crates/krishiv-plan/src/analyzer/fold.rs b/crates/krishiv-plan/src/analyzer/fold.rs index 59dd9c8f0..05b4c574a 100644 --- a/crates/krishiv-plan/src/analyzer/fold.rs +++ b/crates/krishiv-plan/src/analyzer/fold.rs @@ -842,7 +842,16 @@ fn literal(l: &Literal) -> Option { Literal::DayTimeInterval { micros, .. } => { Arc::new(arrow_array::DurationMicrosecondArray::from(vec![*micros])) } - _ => return None, + // Not foldable here: nested values never take part in constant + // folding (a call over one stays a call). + Literal::Array { .. } + | Literal::KeyValueMap { .. } + | Literal::Struct { .. } + | Literal::Map(_) + | Literal::Binary(_) + | Literal::Timestamp(_) + | Literal::TimestampNtz(_) + | Literal::CalendarInterval { .. } => return None, }) } diff --git a/crates/krishiv-plan/src/display.rs b/crates/krishiv-plan/src/display.rs index ed7d4c639..dccca829a 100644 --- a/crates/krishiv-plan/src/display.rs +++ b/crates/krishiv-plan/src/display.rs @@ -1277,6 +1277,38 @@ impl Display for Literal { "INTERVAL '{}'", krishiv_types::interval::calendar_string(*months, *days, *micros) ), + // Nested values print as the constructor call that builds them; + // Spark's own text is golden-pending (design P3-T012a note, 5). + Self::Array { elements, .. } => { + f.write_str("ARRAY(")?; + for (i, e) in elements.iter().enumerate() { + if i > 0 { + f.write_str(", ")?; + } + write!(f, "{e}")?; + } + f.write_str(")") + } + Self::KeyValueMap { entries, .. } => { + f.write_str("MAP(")?; + for (i, (k, v)) in entries.iter().enumerate() { + if i > 0 { + f.write_str(", ")?; + } + write!(f, "{k}, {v}")?; + } + f.write_str(")") + } + Self::Struct { data_type, fields } => { + f.write_str("NAMED_STRUCT(")?; + for (i, (field, v)) in data_type.fields().iter().zip(fields).enumerate() { + if i > 0 { + f.write_str(", ")?; + } + write!(f, "'{}', {v}", Escaped(field.name()))?; + } + f.write_str(")") + } } } } diff --git a/crates/krishiv-plan/src/encode.rs b/crates/krishiv-plan/src/encode.rs index 6d28d5b88..82fcd752a 100644 --- a/crates/krishiv-plan/src/encode.rs +++ b/crates/krishiv-plan/src/encode.rs @@ -1550,6 +1550,36 @@ impl<'k, 'a> Encoder<'k, 'a> { self.put(&days.to_le_bytes()); self.i64(*micros); } + // Nested values: the type, then the values (the type bounds the + // depth, `Literal::array` and friends check it). + Literal::Array { + data_type, + elements, + } => { + self.u8(20); + self.data_type(&krishiv_types::DataType::Array(Arc::clone(data_type))); + self.len(elements.len()); + for e in elements { + self.literal(e); + } + } + Literal::KeyValueMap { data_type, entries } => { + self.u8(21); + self.data_type(&krishiv_types::DataType::Map(Arc::clone(data_type))); + self.len(entries.len()); + for (k, v) in entries { + self.literal(k); + self.literal(v); + } + } + Literal::Struct { data_type, fields } => { + self.u8(22); + self.data_type(&krishiv_types::DataType::Struct((**data_type).clone())); + self.len(fields.len()); + for f in fields { + self.literal(f); + } + } } } } diff --git a/crates/krishiv-plan/src/literal.rs b/crates/krishiv-plan/src/literal.rs index 6d69389c4..46398688b 100644 --- a/crates/krishiv-plan/src/literal.rs +++ b/crates/krishiv-plan/src/literal.rs @@ -1,7 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 //! Literal values. -use krishiv_types::{DataType, DecimalType}; +use std::sync::Arc; + +use krishiv_common::error::{FailedToParseTooComplex, InternalError, KrishivError}; +use krishiv_types::{ArrayType, DataType, DecimalType, MapType, StructType}; pub use krishiv_types::{DayTimeField, YearMonthField}; /// A literal value, in Spark's physical representation for its type. @@ -78,6 +81,30 @@ pub enum Literal { /// Microseconds. micros: i64, }, + /// An `ARRAY` value (Spark Connect `Literal.Array`). Build it with + /// [`Literal::array`]; a null element is `Null(Some(element type))`. + Array { + /// The array type, element type and `containsNull`. + data_type: Arc, + /// The elements, each of the element type. + elements: Vec, + }, + /// A `MAP` value of any key and value type (not the options-only + /// [`Literal::Map`]). Build it with [`Literal::key_value_map`]; keys are + /// never null. + KeyValueMap { + /// The map type. + data_type: Arc, + /// The entries, in order, as `(key, value)`. + entries: Vec<(Literal, Literal)>, + }, + /// A `STRUCT` value. Build it with [`Literal::struct_value`]. + Struct { + /// The struct type. + data_type: Arc, + /// One literal per field, in field order. + fields: Vec, + }, } impl std::fmt::Debug for Literal { @@ -130,6 +157,24 @@ impl std::fmt::Debug for Literal { .field("days", days) .field("micros", micros) .finish(), + Self::Array { + data_type, + elements, + } => f + .debug_struct("Array") + .field("data_type", data_type) + .field("elements", elements) + .finish(), + Self::KeyValueMap { data_type, entries } => f + .debug_struct("KeyValueMap") + .field("data_type", data_type) + .field("entries", entries) + .finish(), + Self::Struct { data_type, fields } => f + .debug_struct("Struct") + .field("data_type", data_type) + .field("fields", fields) + .finish(), } } } @@ -155,6 +200,9 @@ enum Key<'a> { YearMonth(i32, YearMonthField, YearMonthField), DayTime(i64, DayTimeField, DayTimeField), Calendar(i32, i32, i64), + Array(&'a ArrayType, &'a [Literal]), + Pairs(&'a MapType, &'a [(Literal, Literal)]), + Struct(&'a StructType, &'a [Literal]), } impl Literal { @@ -185,6 +233,12 @@ impl Literal { days, micros, } => Key::Calendar(*months, *days, *micros), + Self::Array { + data_type, + elements, + } => Key::Array(data_type, elements), + Self::KeyValueMap { data_type, entries } => Key::Pairs(data_type, entries), + Self::Struct { data_type, fields } => Key::Struct(data_type, fields), } } } @@ -196,3 +250,150 @@ impl PartialEq for Literal { } impl Eq for Literal {} + +fn invalid(message: String) -> KrishivError { + InternalError { message }.into() +} + +fn check_depth(t: &DataType) -> Result<(), KrishivError> { + if t.nesting_depth() > DataType::MAX_NESTING_DEPTH { + return Err(FailedToParseTooComplex {}.into()); + } + Ok(()) +} + +/// Checks that `value` is a literal of `expected` (a typed null included), +/// and that a null is allowed here. +fn check_member( + what: &str, + value: &Literal, + expected: &DataType, + nullable: bool, +) -> Result<(), KrishivError> { + if matches!(value, Literal::Null(_)) && !nullable { + return Err(invalid(format!( + "a null {what} in a nested literal whose type forbids nulls" + ))); + } + match value.data_type() { + Some(t) if t == *expected => Ok(()), + Some(t) => Err(invalid(format!( + "a nested literal {what} of type {} where the type says {}", + t.simple_string(), + expected.simple_string() + ))), + None => Err(invalid(format!( + "an untyped null {what} in a nested literal; nulls inside carry their type" + ))), + } +} + +impl Literal { + /// The type of this literal; `None` for an untyped `NULL`. + #[must_use] + pub fn data_type(&self) -> Option { + Some(match self { + Self::Null(t) => return t.clone().filter(|t| *t != DataType::Null), + Self::Boolean(_) => DataType::Boolean, + Self::Byte(_) => DataType::Byte, + Self::Short(_) => DataType::Short, + Self::Integer(_) => DataType::Integer, + Self::Long(_) => DataType::Long, + Self::Float(_) => DataType::Float, + Self::Double(_) => DataType::Double, + Self::Decimal { data_type, .. } => DataType::Decimal(*data_type), + Self::String(_) => DataType::String, + Self::Map(_) => { + DataType::Map(MapType::new(DataType::String, DataType::String, false).into()) + } + Self::Binary(_) => DataType::Binary, + Self::Date(_) => DataType::Date, + Self::Timestamp(_) => DataType::Timestamp, + Self::TimestampNtz(_) => DataType::TimestampNtz, + Self::YearMonthInterval { start, end, .. } => DataType::YearMonthInterval { + start: *start, + end: *end, + }, + Self::DayTimeInterval { start, end, .. } => DataType::DayTimeInterval { + start: *start, + end: *end, + }, + Self::CalendarInterval { .. } => DataType::CalendarInterval, + Self::Array { data_type, .. } => DataType::Array(Arc::clone(data_type)), + Self::KeyValueMap { data_type, .. } => DataType::Map(Arc::clone(data_type)), + Self::Struct { data_type, .. } => DataType::Struct((**data_type).clone()), + }) + } + + /// An array literal: every element is a literal of the element type + /// (a null is `Null(Some(element type))` and needs `contains_null`). + /// + /// # Errors + /// `FAILED_TO_PARSE_TOO_COMPLEX` when the type nests deeper than + /// [`DataType::MAX_NESTING_DEPTH`]; `INTERNAL_ERROR` when an element + /// has another type or is a forbidden or untyped null. + pub fn array(data_type: ArrayType, elements: Vec) -> Result { + check_depth(&DataType::Array(Arc::new(data_type.clone())))?; + for e in &elements { + check_member( + "element", + e, + data_type.element_type(), + data_type.contains_null(), + )?; + } + Ok(Self::Array { + data_type: Arc::new(data_type), + elements, + }) + } + + /// A map literal of `(key, value)` entries of the map's key and value + /// types; keys are never null. + /// + /// # Errors + /// As [`Literal::array`]. + pub fn key_value_map( + data_type: MapType, + entries: Vec<(Literal, Literal)>, + ) -> Result { + check_depth(&DataType::Map(Arc::new(data_type.clone())))?; + for (k, v) in &entries { + check_member("key", k, data_type.key_type(), false)?; + check_member( + "value", + v, + data_type.value_type(), + data_type.value_contains_null(), + )?; + } + Ok(Self::KeyValueMap { + data_type: Arc::new(data_type), + entries, + }) + } + + /// A struct literal: one literal per field, of the field's type (a null + /// needs a nullable field). + /// + /// # Errors + /// As [`Literal::array`]; also when the count of values differs from the + /// count of fields. + pub fn struct_value(data_type: StructType, fields: Vec) -> Result { + check_depth(&DataType::Struct(data_type.clone()))?; + if data_type.fields().len() != fields.len() { + return Err(invalid(format!( + "a struct literal of {} values for a type of {} fields", + fields.len(), + data_type.fields().len() + ))); + } + for (f, v) in data_type.fields().iter().zip(&fields) { + check_member("field", v, f.data_type(), f.nullable())?; + } + Ok(Self::Struct { + data_type: Arc::new(data_type), + fields, + }) + } +} diff --git a/crates/krishiv-plan/src/raises.rs b/crates/krishiv-plan/src/raises.rs index debe73beb..63c8d7df3 100644 --- a/crates/krishiv-plan/src/raises.rs +++ b/crates/krishiv-plan/src/raises.rs @@ -366,7 +366,17 @@ fn literal_type(l: &Literal) -> Option { Literal::Double(_) => DataType::Double, Literal::Decimal { data_type, .. } => DataType::Decimal(*data_type), Literal::String(_) => DataType::String, - _ => return None, + Literal::Map(_) + | Literal::Array { .. } + | Literal::KeyValueMap { .. } + | Literal::Struct { .. } + | Literal::Binary(_) + | Literal::Date(_) + | Literal::Timestamp(_) + | Literal::TimestampNtz(_) + | Literal::YearMonthInterval { .. } + | Literal::DayTimeInterval { .. } + | Literal::CalendarInterval { .. } => return None, }) } diff --git a/crates/krishiv-plan/src/tests/mod.rs b/crates/krishiv-plan/src/tests/mod.rs index 8cc8a572a..3788244c1 100644 --- a/crates/krishiv-plan/src/tests/mod.rs +++ b/crates/krishiv-plan/src/tests/mod.rs @@ -15,6 +15,7 @@ mod files; mod function_call_forms; mod inventory; mod json_tuple; +mod nested_literals; mod nodes; mod pivot; mod prefix_cache; diff --git a/crates/krishiv-plan/src/tests/nested_literals.rs b/crates/krishiv-plan/src/tests/nested_literals.rs new file mode 100644 index 000000000..64840e46d --- /dev/null +++ b/crates/krishiv-plan/src/tests/nested_literals.rs @@ -0,0 +1,189 @@ +// SPDX-License-Identifier: Apache-2.0 +//! Nested literal values (design P3-T012a note, section 5): the checked +//! constructors, equality by bits, the canonical encoding and the display. + +use krishiv_types::{ArrayType, DataType, MapType, StructField, StructType}; + +use crate::{Expr, Literal, LogicalPlan, Plan, StructuralHash}; + +fn ints(values: &[i32]) -> Literal { + Literal::array( + ArrayType::new(DataType::Integer, true), + values.iter().map(|v| Literal::Integer(*v)).collect(), + ) + .unwrap() +} + +fn hash_of(l: Literal) -> StructuralHash { + let plan = LogicalPlan::table(["t"]).project(vec![Expr::from(l)]); + Plan::try_new(plan).unwrap().structural_hash() +} + +fn class(e: &krishiv_common::error::KrishivError) -> String { + e.error_class().name().to_owned() +} + +#[test] +fn equal_values_are_equal_and_different_ones_are_not() { + assert_eq!(ints(&[1, 2]), ints(&[1, 2])); + assert_ne!(ints(&[1, 2]), ints(&[1, 3])); + assert_ne!(ints(&[1, 2]), ints(&[1])); + // The same values under another type differ. + let longs = Literal::array( + ArrayType::new(DataType::Long, true), + vec![Literal::Long(1), Literal::Long(2)], + ) + .unwrap(); + assert_ne!(ints(&[1, 2]), longs); + assert_eq!(ints(&[]), ints(&[])); +} + +#[test] +fn floats_inside_compare_by_bits() { + let of = |v: f64| { + Literal::array( + ArrayType::new(DataType::Double, false), + vec![Literal::Double(v)], + ) + .unwrap() + }; + assert_eq!(of(f64::NAN), of(f64::NAN)); + assert_ne!(of(0.0), of(-0.0)); + assert_ne!(hash_of(of(0.0)), hash_of(of(-0.0))); + assert_eq!(hash_of(of(f64::NAN)), hash_of(of(f64::NAN))); +} + +#[test] +fn the_hash_tells_values_and_kinds_apart() { + assert_eq!(hash_of(ints(&[1, 2])), hash_of(ints(&[1, 2]))); + assert_ne!(hash_of(ints(&[1, 2])), hash_of(ints(&[2, 1]))); + let map = Literal::key_value_map( + MapType::new(DataType::String, DataType::Integer, true), + vec![(Literal::String("a".to_owned()), Literal::Integer(1))], + ) + .unwrap(); + let other = Literal::key_value_map( + MapType::new(DataType::String, DataType::Integer, true), + vec![(Literal::String("a".to_owned()), Literal::Integer(2))], + ) + .unwrap(); + assert_ne!(hash_of(map.clone()), hash_of(other)); + assert_ne!(hash_of(map), hash_of(ints(&[1]))); + let st = |v: i32| { + Literal::struct_value( + StructType::new([StructField::new("a", DataType::Integer, false)]), + vec![Literal::Integer(v)], + ) + .unwrap() + }; + assert_ne!(hash_of(st(1)), hash_of(st(2))); +} + +#[test] +fn nulls_inside_carry_their_type() { + let ok = Literal::array( + ArrayType::new(DataType::Integer, true), + vec![Literal::Integer(1), Literal::Null(Some(DataType::Integer))], + ); + assert!(ok.is_ok()); + // An untyped null, a null of another type, and a null where the type + // forbids one are refused. + for (contains_null, element) in [ + (true, Literal::Null(None)), + (true, Literal::Null(Some(DataType::Long))), + (false, Literal::Null(Some(DataType::Integer))), + ] { + let e = Literal::array( + ArrayType::new(DataType::Integer, contains_null), + vec![element], + ) + .unwrap_err(); + assert_eq!(class(&e), "INTERNAL_ERROR"); + } +} + +#[test] +fn the_constructors_check_element_types_and_counts() { + let e = Literal::array( + ArrayType::new(DataType::Integer, true), + vec![Literal::String("x".to_owned())], + ) + .unwrap_err(); + assert_eq!(class(&e), "INTERNAL_ERROR"); + // A map key is never null. + let e = Literal::key_value_map( + MapType::new(DataType::Integer, DataType::Integer, true), + vec![(Literal::Null(Some(DataType::Integer)), Literal::Integer(1))], + ) + .unwrap_err(); + assert_eq!(class(&e), "INTERNAL_ERROR"); + // One value per field, a null only in a nullable one. + let ty = StructType::new([ + StructField::new("a", DataType::Integer, false), + StructField::new("b", DataType::String, true), + ]); + assert!(Literal::struct_value(ty.clone(), vec![Literal::Integer(1)]).is_err()); + assert!( + Literal::struct_value( + ty.clone(), + vec![ + Literal::Null(Some(DataType::Integer)), + Literal::Null(Some(DataType::String)) + ] + ) + .is_err() + ); + assert!( + Literal::struct_value( + ty, + vec![Literal::Integer(1), Literal::Null(Some(DataType::String))] + ) + .is_ok() + ); +} + +#[test] +fn the_constructors_check_the_depth() { + let nest = |depth: usize| { + let mut t = DataType::Integer; + for _ in 0..depth { + t = ArrayType::new(t, true).into(); + } + t + }; + let deep = |t: DataType| match t { + DataType::Array(a) => Literal::array((*a).clone(), Vec::new()), + _ => unreachable!(), + }; + assert!(deep(nest(DataType::MAX_NESTING_DEPTH)).is_ok()); + let e = deep(nest(DataType::MAX_NESTING_DEPTH + 1)).unwrap_err(); + assert_eq!(class(&e), "FAILED_TO_PARSE_TOO_COMPLEX"); +} + +#[test] +fn a_literal_reports_its_type() { + let l = ints(&[1]); + assert_eq!( + l.data_type(), + Some(ArrayType::new(DataType::Integer, true).into()) + ); + assert_eq!(Literal::Null(None).data_type(), None); +} + +#[test] +fn nested_literals_print_as_constructor_calls() { + // Golden-pending: Spark's text for a nested literal is not established. + assert_eq!(ints(&[1, 2]).to_string(), "ARRAY(1, 2)"); + let st = Literal::struct_value( + StructType::new([StructField::new("a", DataType::Integer, false)]), + vec![Literal::Integer(1)], + ) + .unwrap(); + assert_eq!(st.to_string(), "NAMED_STRUCT('a', 1)"); + let map = Literal::key_value_map( + MapType::new(DataType::String, DataType::Integer, true), + vec![(Literal::String("a".to_owned()), Literal::Integer(1))], + ) + .unwrap(); + assert_eq!(map.to_string(), "MAP('a', 1)"); +} diff --git a/crates/krishiv-plan/src/tests/props.rs b/crates/krishiv-plan/src/tests/props.rs index e9e7e3587..397f3b643 100644 --- a/crates/krishiv-plan/src/tests/props.rs +++ b/crates/krishiv-plan/src/tests/props.rs @@ -800,6 +800,30 @@ fn m_literal(v: &mut Literal, w: &mut Walk) -> bool { site!(w, *days += 1); site!(w, *micros += 1); } + Literal::Array { elements, .. } => { + site!(w, elements.push(Literal::Integer(1))); + for e in elements.iter_mut() { + if m_literal(e, w) { + return true; + } + } + } + Literal::KeyValueMap { entries, .. } => { + site!(w, entries.push((Literal::Integer(1), Literal::Integer(2)))); + for (k, v) in entries.iter_mut() { + if m_literal(k, w) || m_literal(v, w) { + return true; + } + } + } + Literal::Struct { fields, .. } => { + site!(w, fields.push(Literal::Integer(1))); + for f in fields.iter_mut() { + if m_literal(f, w) { + return true; + } + } + } } false } diff --git a/crates/krishiv-session/src/execute_immediate.rs b/crates/krishiv-session/src/execute_immediate.rs index a538d1a27..9ffa0e851 100644 --- a/crates/krishiv-session/src/execute_immediate.rs +++ b/crates/krishiv-session/src/execute_immediate.rs @@ -501,6 +501,9 @@ fn render_literal(value: &Literal, escaped_strings: bool) -> Result Ok(value.to_string()), + Literal::Array { .. } | Literal::KeyValueMap { .. } | Literal::Struct { .. } => { + Err(unsupported("an array, map or struct type")) + } _ => Err(unsupported("this type")), } } diff --git a/proto/krishiv/plan/v1/plan.proto b/proto/krishiv/plan/v1/plan.proto index 52114b23c..d5c6553fa 100644 --- a/proto/krishiv/plan/v1/plan.proto +++ b/proto/krishiv/plan/v1/plan.proto @@ -1259,9 +1259,44 @@ message Literal { CalendarLiteral calendar = 17; // feature: options_map_literals MapLiteral map = 18; + // feature: nested_literals + ArrayLiteral array = 19; + // feature: nested_literals + KeyValueMapLiteral key_value_map = 20; + // feature: nested_literals + StructLiteral struct = 21; } } +// feature: nested_literals +// An array value of any element type: `data_type` is the array type (a +// nested type, so `nested_types` too), each element a literal of the element +// type, a null one a typed `NullLiteral`. +message ArrayLiteral { + DataType data_type = 1; + repeated Literal elements = 2; +} + +// feature: nested_literals +// A map value of any key and value type (not the options-only `MapLiteral`). +message KeyValueMapLiteral { + DataType data_type = 1; + repeated KeyValueEntry entries = 2; +} + +// feature: nested_literals +message KeyValueEntry { + Literal key = 1; + Literal value = 2; +} + +// feature: nested_literals +// A struct value: one literal per field of `data_type`, in field order. +message StructLiteral { + DataType data_type = 1; + repeated Literal fields = 2; +} + // feature: options_map_literals // A `MAP` literal: one call's options (`from_json`, // `from_csv`, `from_xml` and the like), entries in their given order. From 309a79f0b0487959e0c65a8a7be8df93cfa27364 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 21:08:26 +0000 Subject: [PATCH 02/10] plan, plan-codec, session, exec, sql, connect: a declared schema on ReadSource::Files and codec feature file_source_schema (P3-T012a part 2a, step 2) ReadSource::Files gains schema: Option (variant non_exhaustive) with ReadSource::files(..)/with_schema(..); every construction site moved to the constructor. The binder skips inference when a schema is declared, open_files lists with listing_known and the declared table, canonical encoding writes tag 4 only when set (tag 2 bytes unchanged, pinned), cache key and display (schema=) follow. Codec: optional FilesSource.schema, semantic feature file_source_schema (lists nested_types too), golden, reader-refusal tests. Nullability and per-format semantics stay golden-pending. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../krishiv-connect/src/translate/relation.rs | 12 +-- .../src/planner/runtime_filter.rs | 10 +- crates/krishiv-exec/tests/planner/main.rs | 20 ++-- crates/krishiv-plan-codec/src/decode.rs | 17 +++- crates/krishiv-plan-codec/src/encode.rs | 11 ++- crates/krishiv-plan-codec/src/features.rs | 27 +++++- .../src/proto/krishiv.plan.v1.pb.rs | 4 + .../tests/common/inventory.rs | 25 ++++- crates/krishiv-plan-codec/tests/features.rs | 77 +++++++++++++++ crates/krishiv-plan-codec/tests/golden.rs | 1 + .../nested_types/file_source_schema.bin | Bin 0 -> 745 bytes crates/krishiv-plan-codec/tests/roundtrip.rs | 3 +- crates/krishiv-plan/src/analyzer/bind.rs | 45 ++++++--- crates/krishiv-plan/src/cache.rs | 7 ++ crates/krishiv-plan/src/display.rs | 10 ++ crates/krishiv-plan/src/encode.rs | 8 +- crates/krishiv-plan/src/plan.rs | 42 ++++++++ .../krishiv-plan/src/tests/analyzer/part1.rs | 4 + .../krishiv-plan/src/tests/analyzer/part2.rs | 1 + crates/krishiv-plan/src/tests/display.rs | 1 + crates/krishiv-plan/src/tests/files.rs | 91 ++++++++++++++++++ crates/krishiv-plan/src/tests/inventory.rs | 1 + crates/krishiv-plan/src/tests/props.rs | 3 + crates/krishiv-plan/src/tests/secrets.rs | 1 + crates/krishiv-session/src/catalog.rs | 10 +- .../src/catalog/file_tables.rs | 17 ++++ crates/krishiv-session/src/catalog/tests.rs | 49 ++++++++-- crates/krishiv-session/src/session/env.rs | 6 +- crates/krishiv-session/src/table_commands.rs | 8 +- crates/krishiv-sql/src/lower/relations.rs | 10 +- crates/krishiv-testkit/tests/file_reads.rs | 30 +++--- proto/krishiv/plan/v1/plan.proto | 3 + 32 files changed, 456 insertions(+), 98 deletions(-) create mode 100644 crates/krishiv-plan-codec/tests/golden/nested_types/file_source_schema.bin diff --git a/crates/krishiv-connect/src/translate/relation.rs b/crates/krishiv-connect/src/translate/relation.rs index f00ce4ede..861992ab6 100644 --- a/crates/krishiv-connect/src/translate/relation.rs +++ b/crates/krishiv-connect/src/translate/relation.rs @@ -226,18 +226,16 @@ impl Translator<'_> { "reads with predicates", )); } - ReadSource::Files { - format: d - .format + ReadSource::files( + d.format .clone() .unwrap_or_else(|| self.conf.default_data_source.clone()), - paths: d - .paths + d.paths .iter() .map(|p| OptionValue::from_user("path", p.as_str())) .collect(), - options: typed(&d.options), - } + typed(&d.options), + ) } })) } diff --git a/crates/krishiv-exec/src/planner/runtime_filter.rs b/crates/krishiv-exec/src/planner/runtime_filter.rs index e20e823f9..dc51db23e 100644 --- a/crates/krishiv-exec/src/planner/runtime_filter.rs +++ b/crates/krishiv-exec/src/planner/runtime_filter.rs @@ -353,11 +353,11 @@ mod tests { use crate::physical::{ScanId, ScanNode}; fn files(path: &str) -> ReadSource { - ReadSource::Files { - format: "parquet".to_owned(), - paths: vec![OptionValue::from_user("path", path)], - options: BTreeMap::new(), - } + ReadSource::files( + "parquet".to_owned(), + vec![OptionValue::from_user("path", path)], + BTreeMap::new(), + ) } /// Two paths of one format are two sources. diff --git a/crates/krishiv-exec/tests/planner/main.rs b/crates/krishiv-exec/tests/planner/main.rs index 203ab5ece..3d2d38104 100644 --- a/crates/krishiv-exec/tests/planner/main.rs +++ b/crates/krishiv-exec/tests/planner/main.rs @@ -604,11 +604,11 @@ fn tables_and_statistics() { ); // Reading files by path needs a resolver that lists files. let mut files = scan.clone(); - files.source = krishiv_plan::ReadSource::Files { - format: "parquet".to_owned(), - paths: Vec::new(), - options: std::collections::BTreeMap::new(), - }; + files.source = krishiv_plan::ReadSource::files( + "parquet".to_owned(), + Vec::new(), + std::collections::BTreeMap::new(), + ); let e = MemoryTables::new().provider(&files).unwrap_err(); assert_eq!(e.error_class(), ErrorClass::KrishivUnsupportedFeature); let dir = tempfile::tempdir().unwrap(); @@ -676,11 +676,11 @@ fn unoptimized_plans_and_file_sources() { p.for_each_child(&mut |c| find.push(Arc::clone(c))); } let mut scan = scan.unwrap(); - scan.source = krishiv_plan::ReadSource::Files { - format: "parquet".to_owned(), - paths: Vec::new(), - options: std::collections::BTreeMap::new(), - }; + scan.source = krishiv_plan::ReadSource::files( + "parquet".to_owned(), + Vec::new(), + std::collections::BTreeMap::new(), + ); let (schema, batches) = table(5); let any = AnyTable(Arc::new(Bare(MemoryTable::new(schema, batches)))); assert!(any.0.statistics().is_none()); diff --git a/crates/krishiv-plan-codec/src/decode.rs b/crates/krishiv-plan-codec/src/decode.rs index 303099b99..393d3bfc6 100644 --- a/crates/krishiv-plan-codec/src/decode.rs +++ b/crates/krishiv-plan-codec/src/decode.rs @@ -1470,11 +1470,18 @@ fn read_source(m: &pb::ReadSource) -> Result { name: t.name.clone(), options: options(&t.options)?, }, - Kind::Files(f) => ReadSource::Files { - format: f.format.clone(), - paths: f.paths.iter().map(option).collect::>()?, - options: options(&f.options)?, - }, + Kind::Files(f) => { + let files = ReadSource::files( + f.format.clone(), + f.paths.iter().map(option).collect::>()?, + options(&f.options)?, + ); + match f.schema.as_ref().map(spark_type).transpose()? { + None => files, + Some(DataType::Struct(schema)) => files.with_schema(schema), + Some(_) => return Err(corrupt("a declared file schema that is not a struct")), + } + } }) } diff --git a/crates/krishiv-plan-codec/src/encode.rs b/crates/krishiv-plan-codec/src/encode.rs index 933558847..b0ce80cb4 100644 --- a/crates/krishiv-plan-codec/src/encode.rs +++ b/crates/krishiv-plan-codec/src/encode.rs @@ -636,7 +636,7 @@ impl Encoder { .transpose()?; Ok(pb::ScanNode { id: node.id.0, - source: Some(read_source(source)?), + source: Some(read_source(source, &self.nested)?), output: output .iter() .map(|c| bound_column(c, &self.nested)) @@ -1303,7 +1303,7 @@ fn percentile_expr( } } -fn read_source(source: &ReadSource) -> Result { +fn read_source(source: &ReadSource, nested: &Cell) -> Result { use pb::read_source::Kind; let options = |options: &std::collections::BTreeMap| { options @@ -1320,6 +1320,8 @@ fn read_source(source: &ReadSource) -> Result { format, paths, options: o, + schema, + .. } => Kind::Files(pb::FilesSource { format: format.clone(), paths: paths @@ -1327,6 +1329,11 @@ fn read_source(source: &ReadSource) -> Result { .map(|p| option("path", p)) .collect::>()?, options: options(o)?, + // Written only when declared: a plan that infers keeps its bytes. + schema: schema + .as_ref() + .map(|s| spark_type(&DataType::Struct(s.clone()), nested)) + .transpose()?, }), _ => { return Err(CodecError::Unsupported( diff --git a/crates/krishiv-plan-codec/src/features.rs b/crates/krishiv-plan-codec/src/features.rs index 728be1db2..8b421e4d4 100644 --- a/crates/krishiv-plan-codec/src/features.rs +++ b/crates/krishiv-plan-codec/src/features.rs @@ -130,6 +130,7 @@ registry! { (exchange_source, ExchangeSource, POST_MIGRATION_READER_VERSION, Semantic, "The `ExchangeSource` node: a plan fragment's read of an exchange another stage wrote (distributed stages)."), (existence_join, ExistenceJoin, POST_MIGRATION_READER_VERSION, Semantic, "The `Existence` join type (Spark's `ExistenceJoin`: `EXISTS` and `IN` outside a filter conjunct), hash or nested loop."), (expand, Expand, 5, Semantic, "The `Expand` node and its `NullOf` project item (ROLLUP, CUBE, GROUPING SETS)."), + (file_source_schema, FileSourceSchema, POST_MIGRATION_READER_VERSION, Semantic, "A user-declared schema on a file read (`ReadSource::Files.schema`). A reader without it would infer the columns again and misread formats such as CSV. The schema is a struct type, so a plan with one lists `nested_types` too."), (fused_probe, FusedProbe, POST_MIGRATION_READER_VERSION, Optimization, "A hash join's probe selection (`HashJoinPlan.select`, design P2-E3 U2): a filter and a scan's runtime-filter row stage fused into the probe, which joins the selected rows by id; without it the planner keeps the `Filter` node and the scan's own row stage."), (generate, Generate, POST_MIGRATION_READER_VERSION, Semantic, "The `Generate` node (`explode`, `posexplode`, `inline` and their `_outer` forms)."), (geo_types, GeoTypes, POST_MIGRATION_READER_VERSION, Semantic, "The `GEOMETRY` and `GEOGRAPHY` types, with their SRID."), @@ -241,6 +242,11 @@ pub enum Construct<'a> { /// the literal carries itself, so `nested_types` is listed with /// `nested_literals`. NestedLiteralType, + /// The schema a file read declares (`ReadSource::Files.schema`). + FileSourceSchema, + /// A declared schema's struct type, written as one nested type (so + /// `nested_types` is listed with `file_source_schema`). + SchemaType, /// A `Generate` node's generator. Generator(krishiv_plan::GeneratorKind), /// A hash join's probe selection (`HashJoinPlan.select`). @@ -485,7 +491,8 @@ pub fn feature_of(construct: Construct<'_>) -> Option { Construct::Lambda => Some(Feature::LambdaFunctions), Construct::EvalMode => Some(Feature::CapturedEvalMode), Construct::WriteColumnTypes => Some(Feature::WriteColumnTypes), - Construct::NestedLiteralType => Some(Feature::NestedTypes), + Construct::NestedLiteralType | Construct::SchemaType => Some(Feature::NestedTypes), + Construct::FileSourceSchema => Some(Feature::FileSourceSchema), Construct::ProbeSelect => Some(Feature::FusedProbe), Construct::Generator(kind) => match kind { krishiv_plan::GeneratorKind::VariantExplode => Some(Feature::VariantExplode), @@ -705,9 +712,21 @@ fn constructs_of<'a>(node: &'a PhysNode, f: &mut impl FnMut(Construct<'a>)) { .iter() .for_each(|field| visit_arrow_field(field, f)); } - PhysNode::Scan(_) - | PhysNode::LateFetch { .. } - | PhysNode::OneRow + PhysNode::Scan(scan) | PhysNode::LateFetch { scan, .. } => { + if let krishiv_plan::ReadSource::Files { + schema: Some(schema), + .. + } = &scan.scan.source + { + f(Construct::FileSourceSchema); + f(Construct::SchemaType); + schema + .fields() + .iter() + .for_each(|x| visit_type(x.data_type(), f)); + } + } + PhysNode::OneRow | PhysNode::SharedRead { .. } | PhysNode::UnionLoop { .. } | PhysNode::LoopRef { .. } 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 133eaf351..3bc50ce37 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 @@ -1151,6 +1151,10 @@ pub struct FilesSource { ::prost::alloc::string::String, OptionValue, >, + /// feature: file_source_schema + /// The schema the user declared (a struct type); absent: inferred. + #[prost(message, optional, tag = "4")] + pub schema: ::core::option::Option, } /// A data source option: plain text or a reference to a secret. An inline /// secret or a URL carrying credentials is not encodable (ADR 0038). diff --git a/crates/krishiv-plan-codec/tests/common/inventory.rs b/crates/krishiv-plan-codec/tests/common/inventory.rs index 6c0a580a7..67a9c9b9f 100644 --- a/crates/krishiv-plan-codec/tests/common/inventory.rs +++ b/crates/krishiv-plan-codec/tests/common/inventory.rs @@ -181,14 +181,14 @@ pub(crate) fn file_scan() -> ScanNode { ScanNode { id: ScanId(8), scan: Scan { - source: ReadSource::Files { - format: "parquet".to_owned(), - paths: vec![ + source: ReadSource::files( + "parquet".to_owned(), + vec![ OptionValue::Plain("s3://bucket/a/*.parquet".to_owned()), OptionValue::reference("prod", "signed-path"), ], - options: BTreeMap::new(), - }, + BTreeMap::new(), + ), output: scan_columns(), version: None, }, @@ -2349,7 +2349,22 @@ pub(crate) fn typed() -> Vec<(&'static str, PhysNode)> { nullable: None, subqueries: Vec::new(), }; + // A file read with a declared schema (a nested field too). + let file_source_schema = { + use krishiv_types::{ArrayType, StructField, StructType}; + let mut scan = file_scan(); + scan.scan.source = scan.scan.source.with_schema(StructType::new([ + StructField::new("a", DataType::Integer, false), + StructField::new( + "tags", + DataType::Array(std::sync::Arc::new(ArrayType::new(DataType::String, true))), + true, + ), + ])); + PhysNode::Scan(scan) + }; vec![ + ("file_source_schema", file_source_schema), ("interval_literals", interval_literals), ("options_map_literal", options_map), ("nested_literals", nested_literals), diff --git a/crates/krishiv-plan-codec/tests/features.rs b/crates/krishiv-plan-codec/tests/features.rs index 23e67f8e3..5193ff12f 100644 --- a/crates/krishiv-plan-codec/tests/features.rs +++ b/crates/krishiv-plan-codec/tests/features.rs @@ -117,6 +117,11 @@ fn the_header_lists_exactly_the_features_in_use() { message("nested_literals").features, ["nested_literals", "nested_types"] ); + // The declared schema is one struct type, so a reader needs both. + assert_eq!( + message("file_source_schema").features, + ["file_source_schema", "nested_types"] + ); assert_eq!( message("aggregate_any_value").features, ["any_value_by_aggregates"] @@ -2226,3 +2231,75 @@ fn a_nested_literal_that_does_not_fit_its_type_is_corrupt() { let err = decode(&m.encode_to_vec()).unwrap_err(); assert!(matches!(err, CodecError::Corrupt(_)), "{err:?}"); } + +// ---- declared file schema (file_source_schema) ----------------------------- + +#[test] +fn a_declared_file_schema_needs_its_feature_and_the_type_feature() { + let m = message("file_source_schema"); + assert_eq!(Feature::FileSourceSchema.kind(), FeatureKind::Semantic); + let bytes = m.encode_to_vec(); + let without_schema = ReaderFeatures::only([Feature::NestedTypes]); + assert_eq!( + decode_with(&bytes, &Limits::default(), &without_schema).unwrap_err(), + CodecError::UnsupportedFeatures { + missing: vec!["file_source_schema".to_owned()] + } + ); + let without_types = ReaderFeatures::only([Feature::FileSourceSchema]); + assert_eq!( + decode_with(&bytes, &Limits::default(), &without_types).unwrap_err(), + CodecError::UnsupportedFeatures { + missing: vec!["nested_types".to_owned()] + } + ); + let both = ReaderFeatures::only([Feature::FileSourceSchema, Feature::NestedTypes]); + let plan = decode_with(&bytes, &Limits::default(), &both).unwrap(); + assert_eq!(encode(&plan).unwrap(), bytes); + assert_eq!( + feature_of(Construct::FileSourceSchema), + Some(Feature::FileSourceSchema) + ); +} + +/// A read that infers writes no schema field and lists no feature for it. +#[test] +fn a_read_without_a_declared_schema_has_no_field_and_no_feature() { + let m = message("interval_array_aggregate"); + assert!(!m.features.iter().any(|f| f == "file_source_schema")); + let sources: Vec<_> = m + .nodes + .iter() + .filter_map(|n| match &n.kind { + Some(pb::node::Kind::Scan(s)) => s.source.as_ref(), + _ => None, + }) + .collect(); + assert!(!sources.is_empty()); + for s in sources { + if let Some(pb::read_source::Kind::Files(f)) = &s.kind { + assert!(f.schema.is_none()); + } + } +} + +#[test] +fn a_declared_file_schema_that_is_not_a_struct_is_corrupt() { + let mut m = message("file_source_schema"); + let mut seen = false; + for node in &mut m.nodes { + if let Some(pb::node::Kind::Scan(s)) = &mut node.kind + && let Some(pb::ReadSource { + kind: Some(pb::read_source::Kind::Files(f)), + }) = &mut s.source + { + f.schema = Some(pb::DataType { + kind: pb::data_type::Kind::Integer as i32, + ..pb::DataType::default() + }); + seen = true; + } + } + assert!(seen); + corrupt(&m, "a declared file schema that is not a struct"); +} diff --git a/crates/krishiv-plan-codec/tests/golden.rs b/crates/krishiv-plan-codec/tests/golden.rs index 27cf9c96f..304c1383a 100644 --- a/crates/krishiv-plan-codec/tests/golden.rs +++ b/crates/krishiv-plan-codec/tests/golden.rs @@ -201,6 +201,7 @@ fn the_registry_is_sorted_and_every_feature_has_goldens() { Feature::LegacyStatisticalAggregate | Feature::CollatedMinMax | Feature::NestedLiterals + | Feature::FileSourceSchema ), "{id}: no golden files in {}", dir.display() diff --git a/crates/krishiv-plan-codec/tests/golden/nested_types/file_source_schema.bin b/crates/krishiv-plan-codec/tests/golden/nested_types/file_source_schema.bin new file mode 100644 index 0000000000000000000000000000000000000000..d8e53cd2b846cbf0d444b9e80445cfed63f1954f GIT binary patch literal 745 zcmZWn&5qMB5Vq5_b<>3{21JQCL_Su9O5F?U1++(mIIK7=pk+33(ug!Eu`9N!D)k5l zq&@Heya12DVWmAG5HGUtP-fCn@iW#pdN+W`!E4}Q09%bT7%EsIFDO`7C7x<3UwRzJDeSKs(T+4)qfkutWN zg;w=doa!txYV=~Ea+pvHu(~~!=2j=V#6mvHGC0k!A46tzUrn@y2rQ`4V!7<^$QBZ_ zn-eF(z;^=59Uh7`sHgM?6MQ{&Kxu-@{1f8RK41 zMRE<~HE@wRl_$p$7V+tCUX$)16L^!@Y7&d?B~FCsuv-Ym_A6yk?Yo$7v{{#WWNx)K zJcId$cl|$JK!+K3NSkzB*2c(KfIuVJ$^ph<%X6RV*h-#lwh&8xHuSvuQJRdc&-oDN z=%^^HY)1I2(+s~Za~{?XB?M!fi{KrY00*onb@ZewcxQc4>#jbmufBfy+)z!(A)6a} dEaNoOC#IN{A%<|QCvd!%Yh!hELPKU=^IvhrpJM<3 literal 0 HcmV?d00001 diff --git a/crates/krishiv-plan-codec/tests/roundtrip.rs b/crates/krishiv-plan-codec/tests/roundtrip.rs index e1ae05dea..c4e63335f 100644 --- a/crates/krishiv-plan-codec/tests/roundtrip.rs +++ b/crates/krishiv-plan-codec/tests/roundtrip.rs @@ -635,7 +635,8 @@ fn readers_that_predate_features_still_refuse_what_they_cannot_read() { | "write_files_column_types" | "aggregate_percentiles" | "options_map_literal" - | "nested_literals" => 6, + | "nested_literals" + | "file_source_schema" => 6, // Features added after the version was frozen: above every // sequential format version. n if n.starts_with("time_") diff --git a/crates/krishiv-plan/src/analyzer/bind.rs b/crates/krishiv-plan/src/analyzer/bind.rs index 07e52dfa8..c24e130d3 100644 --- a/crates/krishiv-plan/src/analyzer/bind.rs +++ b/crates/krishiv-plan/src/analyzer/bind.rs @@ -1009,11 +1009,11 @@ async fn path_read( } SourceKind::Unknown => return Ok(BoundTable::Missing), } - let source = ReadSource::Files { - format: format.clone(), - paths: vec![OptionValue::from_user("path", path)], - options: BTreeMap::new(), - }; + let source = ReadSource::files( + format.clone(), + vec![OptionValue::from_user("path", path)], + BTreeMap::new(), + ); Ok(BoundTable::Files( bind_files(ctx, b, checked, &source, origin).await?, )) @@ -1079,6 +1079,7 @@ async fn bind_files( format, paths, options, + schema, }, Some(files), ) = (source, &ctx.files) @@ -1105,23 +1106,35 @@ async fn bind_files( if server_scoped(options) { authorize_store_options(ctx, b, checked, origin).await?; } - let columns = files - .infer(format, &qualified, options) - .await - .map_err(fail)? - .into_iter() - // Spark reads every file column as nullable. - .map(|c| TableColumn { - nullable: true, - ..c - }) - .collect(); + let columns = if let Some(schema) = schema { + // A declared schema replaces inference. The nullability rule is the + // inferred one (every file column nullable); whether Spark keeps a + // declared `NOT NULL` per format is golden-pending (P3-T012a note, 3). + schema + .fields() + .iter() + .map(|f| TableColumn::new(f.name(), f.data_type().clone(), true)) + .collect() + } else { + files + .infer(format, &qualified, options) + .await + .map_err(fail)? + .into_iter() + // Spark reads every file column as nullable. + .map(|c| TableColumn { + nullable: true, + ..c + }) + .collect() + }; b.files.push(BoundFiles { source: source.clone(), read: ReadSource::Files { format: format.to_lowercase(), paths: qualified, options: options.clone(), + schema: schema.clone(), }, columns, }); diff --git a/crates/krishiv-plan/src/cache.rs b/crates/krishiv-plan/src/cache.rs index 404a4a1e5..f3d7774a9 100644 --- a/crates/krishiv-plan/src/cache.rs +++ b/crates/krishiv-plan/src/cache.rs @@ -281,6 +281,7 @@ impl Encoder { format, paths, options, + schema, } => { self.tag("files"); self.text(format); @@ -293,6 +294,11 @@ impl Encoder { self.text(key); self.option_value(value); } + // Only when declared, so a read that infers keeps its key. + if let Some(schema) = schema { + self.tag("schema"); + self.debug(schema); + } } } } @@ -1751,6 +1757,7 @@ mod tests { OptionValue::from_user("password", "pw"), ), ]), + schema: None, })), ), ( diff --git a/crates/krishiv-plan/src/display.rs b/crates/krishiv-plan/src/display.rs index dccca829a..5f349ec17 100644 --- a/crates/krishiv-plan/src/display.rs +++ b/crates/krishiv-plan/src/display.rs @@ -547,12 +547,22 @@ impl Display for SourceText<'_> { format, paths, options, + schema, } => { write!(f, "{format} [")?; for (i, p) in paths.iter().enumerate() { write!(f, "{}{p}", sep(i))?; } f.write_str("]")?; + if let Some(schema) = schema { + f.write_str(" schema=")?; + for (i, field) in schema.fields().iter().enumerate() { + write!(f, "{}{} {}", sep(i), field.name(), field.data_type().sql())?; + if !field.nullable() { + f.write_str(" NOT NULL")?; + } + } + } options } }; diff --git a/crates/krishiv-plan/src/encode.rs b/crates/krishiv-plan/src/encode.rs index 82fcd752a..2a0e1e5fd 100644 --- a/crates/krishiv-plan/src/encode.rs +++ b/crates/krishiv-plan/src/encode.rs @@ -1115,12 +1115,18 @@ impl<'k, 'a> Encoder<'k, 'a> { format, paths, options, + schema, } => { - self.u8(2); + // Tag 2 stays byte-identical for a read that infers; a + // declared schema is tag 4 (tag-2 fields, then the schema). + self.u8(if schema.is_some() { 4 } else { 2 }); self.str(format); self.len(paths.len()); paths.iter().for_each(|p| self.option_value(p)); self.options(options); + if let Some(schema) = schema { + self.str(&krishiv_types::DataType::Struct(schema.clone()).sql()); + } } } } diff --git a/crates/krishiv-plan/src/plan.rs b/crates/krishiv-plan/src/plan.rs index 5ab5fac78..718dab821 100644 --- a/crates/krishiv-plan/src/plan.rs +++ b/crates/krishiv-plan/src/plan.rs @@ -688,6 +688,7 @@ pub enum ReadSource { options: BTreeMap, }, /// Files of a format (`spark.read.parquet(path)`). + #[non_exhaustive] Files { /// Format name, e.g. `parquet`, as written (compared /// case-insensitively; a resolved [`Scan`] holds it in lower case). @@ -699,9 +700,50 @@ pub enum ReadSource { paths: Vec, /// Data source options. options: BTreeMap, + /// The schema the user declared (`spark.read.schema(..)`); `None` + /// infers it from the files. Build with [`ReadSource::files`] and + /// [`ReadSource::with_schema`]. + schema: Option, }, } +impl ReadSource { + /// A read of files of `format` at `paths` whose schema is inferred. + #[must_use] + pub fn files( + format: impl Into, + paths: Vec, + options: BTreeMap, + ) -> Self { + Self::Files { + format: format.into(), + paths, + options, + schema: None, + } + } + + /// This read with the user-declared `schema`; a table read is returned + /// unchanged (Spark's `NamedTable` has no schema). + #[must_use] + pub fn with_schema(self, schema: krishiv_types::StructType) -> Self { + match self { + Self::Files { + format, + paths, + options, + .. + } => Self::Files { + format, + paths, + options, + schema: Some(schema), + }, + table @ Self::Table { .. } => table, + } + } +} + /// A resolved read. #[derive(Debug, Clone, PartialEq, Eq)] pub struct Scan { diff --git a/crates/krishiv-plan/src/tests/analyzer/part1.rs b/crates/krishiv-plan/src/tests/analyzer/part1.rs index 9d83adb80..bdd6b314d 100644 --- a/crates/krishiv-plan/src/tests/analyzer/part1.rs +++ b/crates/krishiv-plan/src/tests/analyzer/part1.rs @@ -467,6 +467,7 @@ fn no_relation_bypasses_the_policy() { format: "parquet".to_owned(), paths, options: std::collections::BTreeMap::new(), + schema: None, })) }; let key = [("l_orderkey", DataType::Long)]; @@ -512,6 +513,7 @@ fn no_relation_bypasses_the_policy() { format: "parquet".to_owned(), paths: path(), options: std::collections::BTreeMap::new(), + schema: None, }, output: Vec::new(), version: None, @@ -814,6 +816,7 @@ fn unsupported_inputs_fail_explicitly() { format: "parquet".to_owned(), paths: vec![krishiv_common::secret::OptionValue::from_user("path", "/d")], options: std::collections::BTreeMap::new(), + schema: None, })); let sha = LogicalPlan::table(["t"]).project(vec![Expr::function("sentences", vec![c("b")])]); for plan in [files, sha] { @@ -874,6 +877,7 @@ fn unsupported_inputs_fail_explicitly() { format: "parquet".to_owned(), paths: vec![krishiv_common::secret::OptionValue::from_user("path", "/d")], options: std::collections::BTreeMap::new(), + schema: None, }, output: Vec::new(), version: None, diff --git a/crates/krishiv-plan/src/tests/analyzer/part2.rs b/crates/krishiv-plan/src/tests/analyzer/part2.rs index 2372403e6..bd454af31 100644 --- a/crates/krishiv-plan/src/tests/analyzer/part2.rs +++ b/crates/krishiv-plan/src/tests/analyzer/part2.rs @@ -312,6 +312,7 @@ fn stars_in_calls_and_elsewhere() { format: "parquet".to_owned(), paths: vec![krishiv_common::secret::OptionValue::from_user("path", "/d")], options: std::collections::BTreeMap::new(), + schema: None, }, output: Vec::new(), version: None, diff --git a/crates/krishiv-plan/src/tests/display.rs b/crates/krishiv-plan/src/tests/display.rs index 245ad1bda..9768fd15c 100644 --- a/crates/krishiv-plan/src/tests/display.rs +++ b/crates/krishiv-plan/src/tests/display.rs @@ -158,6 +158,7 @@ fn options_print_as_typed_values() { format: "csv".to_owned(), paths: vec![OptionValue::from_user("path", "/d")], options, + schema: None, })); assert_eq!( plan.to_string(), diff --git a/crates/krishiv-plan/src/tests/files.rs b/crates/krishiv-plan/src/tests/files.rs index 8b50a5269..80698c074 100644 --- a/crates/krishiv-plan/src/tests/files.rs +++ b/crates/krishiv-plan/src/tests/files.rs @@ -120,6 +120,7 @@ fn read(format: &str, paths: &[&str]) -> LogicalPlan { .map(|p| OptionValue::from_user("path", *p)) .collect(), options: BTreeMap::new(), + schema: None, })) } @@ -548,6 +549,7 @@ fn read_with_secret(table: bool) -> LogicalPlan { format: "parquet".to_owned(), paths: vec![OptionValue::from_user("path", "/data/one")], options, + schema: None, })) } } @@ -584,3 +586,92 @@ fn secret_references_need_usage_on_their_scope() { run(read_with_secret(false), &allowed).unwrap(); assert_eq!(*policy.asked.lock().unwrap(), ["prod"]); } + +// ---- a declared schema (P3-T012a part 2a, design note section 3) ----------- + +fn declared() -> krishiv_types::StructType { + use krishiv_types::StructField; + krishiv_types::StructType::new([ + StructField::new("k", DataType::Long, false), + StructField::new("v", DataType::String, true), + ]) +} + +fn read_declared(path: &str) -> LogicalPlan { + LogicalPlan::new(PlanNode::Read( + ReadSource::files( + "parquet", + vec![OptionValue::from_user("path", path)], + BTreeMap::new(), + ) + .with_schema(declared()), + )) +} + +/// A declared schema replaces inference: the resolver is never asked, and +/// the declared fields are the columns. Nullability keeps the rule of an +/// inferred read (every file column nullable); whether Spark keeps a declared +/// `NOT NULL` per format is golden-pending. +#[test] +fn a_declared_schema_skips_inference() { + let files = Arc::new(Files::default()); + // The path does not exist for inference: only a declared schema reads it. + let r = run(read_declared("/data/missing"), &ctx(&files)); + assert!(files.calls.lock().unwrap().is_empty()); + assert_eq!( + output(&r), + vec![ + ("k".to_owned(), DataType::Long, true), + ("v".to_owned(), DataType::String, true), + ] + ); + let a = r.unwrap(); + let PlanNode::Scan(scan) = &a.plan.root().node else { + panic!("not a scan: {}", a.plan.root()); + }; + // The scan keeps the declared schema, so the session opens it as declared. + assert!(matches!(&scan.source, ReadSource::Files { schema: Some(s), .. } if *s == declared())); +} + +/// Two reads of one path differ by their schema: each is bound once. +#[test] +fn a_declared_schema_is_part_of_the_read() { + let files = Arc::new(Files::default()); + let r = run(read("parquet", &["/data/one"]), &ctx(&files)); + assert_eq!(files.calls.lock().unwrap().len(), 1); + assert_eq!(output(&r), nullable_one()); + let r = run(read_declared("/data/one"), &ctx(&files)); + assert_eq!(files.calls.lock().unwrap().len(), 1, "no second inference"); + assert_eq!(output(&r)[0].0, "k"); +} + +/// Display shows the schema; the canonical encoding of a read that infers is +/// the tag-2 bytes (unchanged), a declared schema is tag 4 and differs. +#[test] +fn a_declared_schema_prints_and_encodes_apart() { + let plan = read_declared("/d"); + assert_eq!( + plan.to_string(), + "Read parquet [/d] schema=k BIGINT NOT NULL, v STRING\n" + ); + let inferred = LogicalPlan::new(PlanNode::Read(ReadSource::files( + "parquet", + vec![OptionValue::from_user("path", "/d")], + BTreeMap::new(), + ))); + assert_eq!(inferred.to_string(), "Read parquet [/d]\n"); + let enc = |p: LogicalPlan| Plan::try_new(p).unwrap().canonical_encoding(); + let (a, b) = (enc(inferred), enc(plan)); + assert_ne!(a, b); + // Pinned: the bytes of a read that infers (source tag 2) never change. + assert_eq!( + a, + [ + 11, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 45, 0, 0, 0, 0, 0, 0, 0, 80, 1, 2, 7, 0, 0, + 0, 0, 0, 0, 0, 112, 97, 114, 113, 117, 101, 116, 1, 0, 0, 0, 0, 0, 0, 0, 1, 2, 0, 0, 0, + 0, 0, 0, 0, 47, 100, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 + ] + ); + assert!(b.len() > a.len()); + assert!(b.windows(3).any(|w| w == b"BIG")); +} diff --git a/crates/krishiv-plan/src/tests/inventory.rs b/crates/krishiv-plan/src/tests/inventory.rs index c92066c72..ec44e1adf 100644 --- a/crates/krishiv-plan/src/tests/inventory.rs +++ b/crates/krishiv-plan/src/tests/inventory.rs @@ -186,6 +186,7 @@ fn slice_ops() -> Vec<(&'static str, LogicalPlan)> { "path", "/data", )], options: BTreeMap::new(), + schema: None, })); let sql = LogicalPlan::new(PlanNode::Sql { query: "SELECT 1".to_owned(), diff --git a/crates/krishiv-plan/src/tests/props.rs b/crates/krishiv-plan/src/tests/props.rs index 397f3b643..31665a3f2 100644 --- a/crates/krishiv-plan/src/tests/props.rs +++ b/crates/krishiv-plan/src/tests/props.rs @@ -378,6 +378,7 @@ impl Gen<'_> { format: self.name(), paths: self.many(1, 2, Gen::option), options, + schema: None, } } } @@ -1084,6 +1085,7 @@ fn m_source(s: &mut ReadSource, w: &mut Walk) -> bool { format: "parquet".to_owned(), paths: name.iter().cloned().map(OptionValue::Plain).collect(), options: options.clone(), + schema: None, }, ReadSource::Files { paths, options, .. } => ReadSource::Table { name: paths.iter().map(ToString::to_string).collect(), @@ -1097,6 +1099,7 @@ fn m_source(s: &mut ReadSource, w: &mut Walk) -> bool { format, paths, options, + .. } => m_string(format, w) || m_values(paths, w) || m_options(options, w), } } diff --git a/crates/krishiv-plan/src/tests/secrets.rs b/crates/krishiv-plan/src/tests/secrets.rs index 548ca195b..11b14ac4b 100644 --- a/crates/krishiv-plan/src/tests/secrets.rs +++ b/crates/krishiv-plan/src/tests/secrets.rs @@ -45,6 +45,7 @@ fn read(secret: &str) -> LogicalPlan { format!("https://b.s3.amazonaws.com/k?X-Amz-Signature={secret}&X-Amz-Expires=60"), )], options: options(secret), + schema: None, })) } diff --git a/crates/krishiv-session/src/catalog.rs b/crates/krishiv-session/src/catalog.rs index a83543671..da720d3f8 100644 --- a/crates/krishiv-session/src/catalog.rs +++ b/crates/krishiv-session/src/catalog.rs @@ -1087,13 +1087,13 @@ pub(crate) fn file_metadata(provider: &str) -> FileMetadata { /// The read of a table's files. pub(crate) fn table_source(table: &crate::store::CatalogTable) -> ReadSource { - ReadSource::Files { - format: table.provider.clone(), - paths: vec![krishiv_common::secret::OptionValue::from_user( + ReadSource::files( + table.provider.clone(), + vec![krishiv_common::secret::OptionValue::from_user( "path", table.location.display().to_string(), )], - options: table + table .options .iter() .map(|(k, v)| { @@ -1103,7 +1103,7 @@ pub(crate) fn table_source(table: &crate::store::CatalogTable) -> ReadSource { ) }) .collect(), - } + ) } fn unsupported(feature: &str) -> KrishivError { diff --git a/crates/krishiv-session/src/catalog/file_tables.rs b/crates/krishiv-session/src/catalog/file_tables.rs index ad473b1f3..8d1ef2420 100644 --- a/crates/krishiv-session/src/catalog/file_tables.rs +++ b/crates/krishiv-session/src/catalog/file_tables.rs @@ -46,11 +46,28 @@ impl QueryCatalog { format, paths, options, + schema, + .. } = source else { continue; }; let listing = match key { + // A declared schema is the table: nothing is inferred. Which + // directories a declared-schema read descends into (partition + // discovery per format) is golden-pending, so none beyond the + // root files, as a catalog table without partitions. + FilesKey::Read(_) if schema.is_some() => { + self.files + .listing_known( + format, + paths, + options, + Arc::clone(table), + &krishiv_exec::files::KnownPartitions::default(), + ) + .await? + } FilesKey::Table(_) | FilesKey::Reference { known: true, .. } => { let Some(partitions) = partitions else { return Err(krishiv_common::error::KrishivUnsupportedFeature { diff --git a/crates/krishiv-session/src/catalog/tests.rs b/crates/krishiv-session/src/catalog/tests.rs index e142938eb..d98de1114 100644 --- a/crates/krishiv-session/src/catalog/tests.rs +++ b/crates/krishiv-session/src/catalog/tests.rs @@ -53,11 +53,7 @@ fn files(col: &str) -> TempTable { known: false, permanent: false, partitions: Some(Default::default()), - source: ReadSource::Files { - format: "parquet".to_owned(), - paths: vec![], - options: Default::default(), - }, + source: ReadSource::files("parquet".to_owned(), vec![], Default::default()), columns: vec![TableColumn::new(col.to_owned(), DataType::Integer, false)], } } @@ -917,11 +913,11 @@ fn task_specs_rejects_what_it_cannot_split() { // A scan of files that were never opened. let e = catalog .task_specs( - &split_scan_plan(ReadSource::Files { - format: "parquet".to_owned(), - paths: Vec::new(), - options: BTreeMap::new(), - }), + &split_scan_plan(ReadSource::files( + "parquet".to_owned(), + Vec::new(), + BTreeMap::new(), + )), 2, 2, ) @@ -1087,3 +1083,36 @@ fn task_specs_refuse_a_disk_cached_scan_and_accept_a_memory_one() { rt.block_on(cached.unpersist()).unwrap(); } } + +/// A read with a declared schema opens with the declared table, nothing +/// inferred: a headerless CSV would be named `_c0`, `_c1` by inference. The +/// declared nullability follows the inferred-read rule (golden-pending). +#[test] +fn a_file_read_with_a_declared_schema_uses_it() { + use krishiv_common::secret::OptionValue; + use krishiv_types::{StructField, StructType}; + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("t.csv"); + std::fs::write(&path, "1,a\n2,b\n").unwrap(); + let s = Session::new().unwrap(); + let source = ReadSource::files( + "csv", + vec![OptionValue::from_user("path", path.to_string_lossy())], + Default::default(), + ) + .with_schema(StructType::new([ + StructField::new("k", DataType::Long, true), + StructField::new("v", DataType::String, true), + ])); + let df = crate::dataframe::DataFrame::new( + s, + Arc::new(krishiv_plan::LogicalPlan::new( + krishiv_plan::PlanNode::Read(source), + )), + ); + let batches = block_on(df.collect()).unwrap(); + let schema = batches[0].schema(); + let names: Vec<_> = schema.fields().iter().map(|f| f.name().clone()).collect(); + assert_eq!(names, ["k", "v"]); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 2); +} diff --git a/crates/krishiv-session/src/session/env.rs b/crates/krishiv-session/src/session/env.rs index afe6c612e..b733b2b0c 100644 --- a/crates/krishiv-session/src/session/env.rs +++ b/crates/krishiv-session/src/session/env.rs @@ -158,9 +158,9 @@ pub(crate) fn read_files( paths: Vec, options: BTreeMap, ) -> LogicalPlan { - LogicalPlan::new(PlanNode::Read(ReadSource::Files { - format: format.to_owned(), + LogicalPlan::new(PlanNode::Read(ReadSource::files( + format.to_owned(), paths, options, - })) + ))) } diff --git a/crates/krishiv-session/src/table_commands.rs b/crates/krishiv-session/src/table_commands.rs index 86e205f8e..0d8329c26 100644 --- a/crates/krishiv-session/src/table_commands.rs +++ b/crates/krishiv-session/src/table_commands.rs @@ -826,17 +826,17 @@ impl Session { .into()); } } - let source = krishiv_plan::ReadSource::Files { - format: provider, + let source = krishiv_plan::ReadSource::files( + provider, paths, - options: options + options .into_iter() .map(|(k, val)| { let val = OptionValue::from_user(&k, val); (k, val) }) .collect(), - }; + ); self.state.catalog.put( name, crate::catalog::TempTable::Files { diff --git a/crates/krishiv-sql/src/lower/relations.rs b/crates/krishiv-sql/src/lower/relations.rs index e5bec2aab..1131d58d0 100644 --- a/crates/krishiv-sql/src/lower/relations.rs +++ b/crates/krishiv-sql/src/lower/relations.rs @@ -617,12 +617,12 @@ impl Lowerer { let source = match &name.parts[..] { // `parquet.`s3://...``: a path, kept typed so that its // credentials never sit in a name. - [format, path] if path.quoted && is_path(&path.value) => ReadSource::Files { + [format, path] if path.quoted && is_path(&path.value) => ReadSource::files( // As written: Spark names it so in errors. - format: format.value.clone(), - paths: vec![OptionValue::from_user("path", path.value.clone())], - options: opts, - }, + format.value.clone(), + vec![OptionValue::from_user("path", path.value.clone())], + opts, + ), _ => ReadSource::Table { name: parts, options: opts, diff --git a/crates/krishiv-testkit/tests/file_reads.rs b/crates/krishiv-testkit/tests/file_reads.rs index 0911616bc..a2e2b1f4e 100644 --- a/crates/krishiv-testkit/tests/file_reads.rs +++ b/crates/krishiv-testkit/tests/file_reads.rs @@ -206,11 +206,11 @@ fn directories_list_like_spark() { let analyzed = env.block_on(DefaultAnalyzer.analyze(&plan, &ctx)).unwrap(); assert_eq!(analyzed.output[0].name, "a"); // An unknown `USING` format is Spark's DATA_SOURCE_NOT_FOUND (42K02). - let read = LogicalPlan::new(PlanNode::Read(ReadSource::Files { - format: "delta".to_owned(), - paths: paths(&["/d"]), - options: BTreeMap::new(), - })); + let read = LogicalPlan::new(PlanNode::Read(ReadSource::files( + "delta".to_owned(), + paths(&["/d"]), + BTreeMap::new(), + ))); let e = env .block_on(DefaultAnalyzer.analyze(&Plan::try_new(read).unwrap(), &ctx)) .unwrap_err(); @@ -325,11 +325,11 @@ fn the_analyzer_reads_parquet_paths() { let analyze = |plan: LogicalPlan| { env.block_on(DefaultAnalyzer.analyze(&Plan::try_new(plan).unwrap(), &ctx)) }; - let read = LogicalPlan::new(PlanNode::Read(ReadSource::Files { - format: "parquet".to_owned(), - paths: paths(&["/t"]), - options: BTreeMap::new(), - })); + let read = LogicalPlan::new(PlanNode::Read(ReadSource::files( + "parquet".to_owned(), + paths(&["/t"]), + BTreeMap::new(), + ))); let analyzed = analyze(read).unwrap(); assert_eq!(analyzed.output.len(), 1); // Non-null in the file, nullable as Spark reads it. @@ -471,11 +471,11 @@ fn agrees_with_spark_on_the_file_read_probe() { .map(|p| sub(p.as_str().unwrap())) .collect(); let refs: Vec<&str> = paths.iter().map(String::as_str).collect(); - Plan::try_new(LogicalPlan::new(PlanNode::Read(ReadSource::Files { - format: "parquet".to_owned(), - paths: self::paths(&refs), - options: BTreeMap::new(), - }))) + Plan::try_new(LogicalPlan::new(PlanNode::Read(ReadSource::files( + "parquet".to_owned(), + self::paths(&refs), + BTreeMap::new(), + )))) .unwrap() } else { let sql = sub(case["input"].as_str().unwrap()); diff --git a/proto/krishiv/plan/v1/plan.proto b/proto/krishiv/plan/v1/plan.proto index d5c6553fa..ce6b79d10 100644 --- a/proto/krishiv/plan/v1/plan.proto +++ b/proto/krishiv/plan/v1/plan.proto @@ -999,6 +999,9 @@ message FilesSource { string format = 1; repeated OptionValue paths = 2; map options = 3; + // feature: file_source_schema + // The schema the user declared (a struct type); absent: inferred. + DataType schema = 4; } // A data source option: plain text or a reference to a secret. An inline From 6109bdb86677ea5073306ed67300491f86e65f75 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 22:21:49 +0000 Subject: [PATCH 03/10] plan, memory, plan-codec, session, exec: ReadSource::Local, the LocalRelation carrier and codec feature local_relation_source (P3-T012a part 2a, step 3) krishiv-plan: LocalRelation (hash, counts, optional RelationLease; built only from a lease or by hash after a decode; Debug never prints values), ReadSource::Local resolved to a Scan over the Arrow schema's columns (resolve, hints, bind, prefix cache cacheable), canonical tag 3 and cache key, display `LocalRelation [a#0, b#1] (N rows)`. krishiv-plan gains a krishiv-memory edge (L3 to L1). krishiv-memory: RelationStore::lease (finds the entry and adds an owner under one lock, None when missing) and RelationStore::decode_ipc (decode before insert). krishiv-session: Session::local_relation (unstable-ir) with eager casts through the session executor and a rename path that shares buffers and checks nested non-nullable nulls; QueryCatalog::provider and task_specs Local arms (INTERNAL_ERROR naming the hash when the entry is gone); result cache Local arm. krishiv-exec: source_name. Codec: LocalSource and semantic feature local_relation_source, golden, reader refusal. Explain text, the schemaRequired class and the cast rules stay golden-pending. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- Cargo.lock | 1 + crates/krishiv-exec/src/planner/build/util.rs | 1 + crates/krishiv-exec/src/planner/table.rs | 3 +- crates/krishiv-memory/src/relation.rs | 72 ++++ crates/krishiv-plan-codec/src/decode.rs | 16 + crates/krishiv-plan-codec/src/encode.rs | 7 + crates/krishiv-plan-codec/src/features.rs | 25 +- .../src/proto/krishiv.plan.v1.pb.rs | 24 +- .../tests/common/inventory.rs | 17 + crates/krishiv-plan-codec/tests/features.rs | 93 +++++ .../local_relation_source.bin | Bin 0 -> 666 bytes crates/krishiv-plan-codec/tests/roundtrip.rs | 3 +- crates/krishiv-plan/CLAUDE.md | 1 + crates/krishiv-plan/Cargo.toml | 2 + crates/krishiv-plan/src/analyzer/bind.rs | 4 + crates/krishiv-plan/src/analyzer/hints.rs | 5 +- crates/krishiv-plan/src/analyzer/ops.rs | 31 ++ crates/krishiv-plan/src/analyzer/resolve.rs | 7 +- crates/krishiv-plan/src/cache.rs | 17 + crates/krishiv-plan/src/display.rs | 25 ++ crates/krishiv-plan/src/encode.rs | 17 + crates/krishiv-plan/src/lib.rs | 8 +- crates/krishiv-plan/src/local.rs | 132 ++++++ crates/krishiv-plan/src/plan.rs | 6 +- crates/krishiv-plan/src/tests/local.rs | 189 +++++++++ crates/krishiv-plan/src/tests/mod.rs | 1 + crates/krishiv-plan/src/tests/props.rs | 9 + crates/krishiv-session/CLAUDE.md | 1 + crates/krishiv-session/src/catalog.rs | 47 +++ crates/krishiv-session/src/lib.rs | 2 + crates/krishiv-session/src/local_relation.rs | 381 ++++++++++++++++++ .../src/local_relation/tests.rs | 326 +++++++++++++++ crates/krishiv-session/src/result_cache.rs | 4 + crates/krishiv-session/src/session.rs | 1 + proto/krishiv/plan/v1/plan.proto | 16 + 35 files changed, 1480 insertions(+), 14 deletions(-) create mode 100644 crates/krishiv-plan-codec/tests/golden/local_relation_source/local_relation_source.bin create mode 100644 crates/krishiv-plan/src/local.rs create mode 100644 crates/krishiv-plan/src/tests/local.rs create mode 100644 crates/krishiv-session/src/local_relation.rs create mode 100644 crates/krishiv-session/src/local_relation/tests.rs diff --git a/Cargo.lock b/Cargo.lock index 67ad3cc06..33b63805b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2665,6 +2665,7 @@ dependencies = [ "krishiv-functions", "krishiv-io", "krishiv-kernels", + "krishiv-memory", "krishiv-types", "krishiv-udf-sdk", "proptest", diff --git a/crates/krishiv-exec/src/planner/build/util.rs b/crates/krishiv-exec/src/planner/build/util.rs index dc2a881e1..99e561cab 100644 --- a/crates/krishiv-exec/src/planner/build/util.rs +++ b/crates/krishiv-exec/src/planner/build/util.rs @@ -232,6 +232,7 @@ pub(crate) fn source_name(scan: &krishiv_plan::Scan) -> String { match &scan.source { krishiv_plan::ReadSource::Table { name, .. } => name.join("."), krishiv_plan::ReadSource::Files { format, .. } => format!("{format} files"), + krishiv_plan::ReadSource::Local(_) => "LocalRelation".to_owned(), _ => "source".to_owned(), } } diff --git a/crates/krishiv-exec/src/planner/table.rs b/crates/krishiv-exec/src/planner/table.rs index 1b379dc10..8de71ed9b 100644 --- a/crates/krishiv-exec/src/planner/table.rs +++ b/crates/krishiv-exec/src/planner/table.rs @@ -253,7 +253,8 @@ impl TableProviderResolver for MemoryTables { fn provider(&self, scan: &Scan) -> Result> { let ReadSource::Table { name, .. } = &scan.source else { return Err(KrishivUnsupportedFeature { - feature: "reading files by path from in-memory tables".to_owned(), + feature: "reading files by path or a local relation from in-memory tables" + .to_owned(), } .into()); }; diff --git a/crates/krishiv-memory/src/relation.rs b/crates/krishiv-memory/src/relation.rs index 03a8c966f..9e4013f97 100644 --- a/crates/krishiv-memory/src/relation.rs +++ b/crates/krishiv-memory/src/relation.rs @@ -53,6 +53,12 @@ impl RelationHash { Self(bytes) } + /// The 32 hash bytes. + #[must_use] + pub fn as_bytes(&self) -> &[u8; 32] { + &self.0 + } + /// The hash as 64 hexadecimal characters (as Spark Connect addresses its /// cached local relations). #[must_use] @@ -279,6 +285,33 @@ impl RelationStore { self.insert(&schema, &batches) } + /// Decodes wire bytes (an Arrow IPC stream) into their schema and + /// unweighted batches **without** storing them, so a caller can transform + /// the batches (cast, relabel) and insert only the result. + /// + /// # Errors + /// `INTERNAL_ERROR` for undecodable bytes. + pub fn decode_ipc(ipc: &[u8]) -> Result<(SchemaRef, Vec), KrishivError> { + let (schema, records) = decode_records(ipc)?; + Ok((schema, records.into_iter().map(Batch::new).collect())) + } + + /// A new lease on the entry `hash`, or `None` if the store does not hold + /// it. The entry is found and an owner added under one lock, so a + /// concurrent last-lease drop cannot interleave; the batch count comes + /// from the entry. + #[must_use] + pub fn lease(self: &Arc, hash: &RelationHash) -> Option { + let mut entries = self.entries.lock().unwrap_or_else(PoisonError::into_inner); + let (entry, owners) = entries.map.get_mut(hash)?; + *owners += 1; + Some(RelationLease { + store: Arc::clone(self), + hash: *hash, + batches: entry.batches.len(), + }) + } + /// Releases one owner of `hash`, dropping its entry (and its bytes) /// with the last one. fn release(&self, hash: &RelationHash) { @@ -718,4 +751,43 @@ mod tests { assert_eq!(store.reserved(), 0, "the last lease returns the bytes"); assert!(store.is_empty()); } + + #[test] + fn lease_finds_a_held_entry_and_adds_an_owner() { + let store = store(); + let schema = schema(); + let held = store + .insert(&schema, &[batch(&schema, &[1]), batch(&schema, &[2])]) + .expect("insert stores"); + let found = store.lease(&held.hash()).expect("the entry is held"); + assert_eq!(found.hash(), held.hash()); + assert_eq!(found.batches(), 2, "the count comes from the entry"); + drop(held); + assert_eq!(store.len(), 1, "the found lease keeps the entry alive"); + drop(found); + assert!(store.is_empty()); + } + + #[test] + fn lease_of_a_missing_hash_is_none_and_creates_nothing() { + let store = store(); + assert!(store.lease(&RelationHash::from_bytes([3; 32])).is_none()); + assert!(store.is_empty()); + } + + #[test] + fn decode_ipc_returns_the_batches_without_storing_them() { + let store = store(); + let schema = schema(); + let lease = store + .insert(&schema, &[batch(&schema, &[1, 2, 3])]) + .expect("insert stores"); + let ipc = store.encode_ipc(&lease.hash()).expect("encodes"); + let (decoded, batches) = RelationStore::decode_ipc(&ipc).expect("decodes"); + assert_eq!(decoded, schema); + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 3); + assert_eq!(store.len(), 1, "decoding stores nothing more"); + assert!(RelationStore::decode_ipc(b"not ipc").is_err()); + } } diff --git a/crates/krishiv-plan-codec/src/decode.rs b/crates/krishiv-plan-codec/src/decode.rs index 393d3bfc6..928dcc653 100644 --- a/crates/krishiv-plan-codec/src/decode.rs +++ b/crates/krishiv-plan-codec/src/decode.rs @@ -1482,6 +1482,22 @@ fn read_source(m: &pb::ReadSource) -> Result { Some(_) => return Err(corrupt("a declared file schema that is not a struct")), } } + Kind::Local(l) => { + let hash: [u8; 32] = l + .hash + .as_slice() + .try_into() + .map_err(|_| corrupt("a local relation hash that is not 32 bytes"))?; + let batches = usize::try_from(l.batches) + .map_err(|_| corrupt("a local relation batch count that does not fit"))?; + ReadSource::Local(krishiv_plan::LocalRelation::from_decoded( + schema_from(&l.schema)?, + hash, + batches, + l.rows, + l.bytes, + )) + } }) } diff --git a/crates/krishiv-plan-codec/src/encode.rs b/crates/krishiv-plan-codec/src/encode.rs index b0ce80cb4..ca2203b7e 100644 --- a/crates/krishiv-plan-codec/src/encode.rs +++ b/crates/krishiv-plan-codec/src/encode.rs @@ -1335,6 +1335,13 @@ fn read_source(source: &ReadSource, nested: &Cell) -> Result Kind::Local(pb::LocalSource { + hash: l.relation().hash.as_bytes().to_vec(), + schema: flagged_schema_bytes(l.schema(), nested), + rows: l.rows(), + batches: l.batches() as u64, + bytes: l.bytes(), + }), _ => { return Err(CodecError::Unsupported( "a read source this version does not know".to_owned(), diff --git a/crates/krishiv-plan-codec/src/features.rs b/crates/krishiv-plan-codec/src/features.rs index 8b421e4d4..9aeec483b 100644 --- a/crates/krishiv-plan-codec/src/features.rs +++ b/crates/krishiv-plan-codec/src/features.rs @@ -146,6 +146,7 @@ registry! { (legacy_statistical_aggregate, LegacyStatisticalAggregate, POST_MIGRATION_READER_VERSION, Semantic, "The `legacy_statistical_aggregate` expression environment flag (`spark.sql.legacy.statisticalAggregate`): the sample moment aggregates give NaN instead of NULL where Spark divides by zero."), (legacy_time_parser, LegacyTimeParser, POST_MIGRATION_READER_VERSION, Semantic, "The `legacy_time_parser` expression environment flag (`spark.sql.legacy.timeParserPolicy=LEGACY`): the datetime pattern functions use `SimpleDateFormat`, and the JSON, CSV and XML functions the lenient `FastDateFormat`."), (limit_offset, LimitOffset, 3, Semantic, "A `Limit` that skips leading rows (a nonzero offset)."), + (local_relation_source, LocalRelationSource, POST_MIGRATION_READER_VERSION, Semantic, "A read of Arrow data held in the session's relation store (`ReadSource::Local`, Spark's `LocalRelation`), named by content hash; the data is not in the plan. A reader without it has no way to read the relation."), (map_key_dedup_policy, MapKeyDedupPolicy, POST_MIGRATION_READER_VERSION, Semantic, "The session's map-key dedup policy (`LAST_WIN`); without it a worker keeps the first value on duplicate map keys."), (mode_regr_bit_listagg_aggs, ModeRegrBitListaggAggs, POST_MIGRATION_READER_VERSION, Semantic, "The `mode`, regression (`regr_*`), bitwise (`bit_and`, `bit_or`, `bit_xor`), `listagg`/`string_agg` and `try_sum`/`try_avg` aggregates."), (moment_aggregates, MomentAggregates, POST_MIGRATION_READER_VERSION, Semantic, "The moment aggregates (`stddev_samp`, `stddev_pop`, `var_samp`, `var_pop`, `skewness`, `kurtosis`, `corr`, `covar_samp`, `covar_pop`)."), @@ -244,6 +245,8 @@ pub enum Construct<'a> { NestedLiteralType, /// The schema a file read declares (`ReadSource::Files.schema`). FileSourceSchema, + /// A scan of a local relation (`ReadSource::Local`). + LocalRelationSource, /// A declared schema's struct type, written as one nested type (so /// `nested_types` is listed with `file_source_schema`). SchemaType, @@ -493,6 +496,7 @@ pub fn feature_of(construct: Construct<'_>) -> Option { Construct::WriteColumnTypes => Some(Feature::WriteColumnTypes), Construct::NestedLiteralType | Construct::SchemaType => Some(Feature::NestedTypes), Construct::FileSourceSchema => Some(Feature::FileSourceSchema), + Construct::LocalRelationSource => Some(Feature::LocalRelationSource), Construct::ProbeSelect => Some(Feature::FusedProbe), Construct::Generator(kind) => match kind { krishiv_plan::GeneratorKind::VariantExplode => Some(Feature::VariantExplode), @@ -712,12 +716,11 @@ fn constructs_of<'a>(node: &'a PhysNode, f: &mut impl FnMut(Construct<'a>)) { .iter() .for_each(|field| visit_arrow_field(field, f)); } - PhysNode::Scan(scan) | PhysNode::LateFetch { scan, .. } => { - if let krishiv_plan::ReadSource::Files { + PhysNode::Scan(scan) | PhysNode::LateFetch { scan, .. } => match &scan.scan.source { + krishiv_plan::ReadSource::Files { schema: Some(schema), .. - } = &scan.scan.source - { + } => { f(Construct::FileSourceSchema); f(Construct::SchemaType); schema @@ -725,7 +728,19 @@ fn constructs_of<'a>(node: &'a PhysNode, f: &mut impl FnMut(Construct<'a>)) { .iter() .for_each(|x| visit_type(x.data_type(), f)); } - } + // The stored batches' Arrow schema travels in the plan, so its + // fields are classified like any schema the plan writes. + krishiv_plan::ReadSource::Local(l) => { + f(Construct::LocalRelationSource); + l.schema() + .fields() + .iter() + .for_each(|field| visit_arrow_field(field, f)); + } + // Inferred file reads and table reads need nothing here (the + // encoder refuses a table read it cannot write). + _ => {} + }, PhysNode::OneRow | PhysNode::SharedRead { .. } | PhysNode::UnionLoop { .. } 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 3bc50ce37..1bdf2cc2e 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 @@ -1117,7 +1117,7 @@ pub struct RowStageSettings { } #[derive(Clone, PartialEq, ::prost::Message)] pub struct ReadSource { - #[prost(oneof = "read_source::Kind", tags = "1, 2")] + #[prost(oneof = "read_source::Kind", tags = "1, 2, 3")] pub kind: ::core::option::Option, } /// Nested message and enum types in `ReadSource`. @@ -1128,6 +1128,9 @@ pub mod read_source { Table(super::TableSource), #[prost(message, tag = "2")] Files(super::FilesSource), + /// feature: local_relation_source + #[prost(message, tag = "3")] + Local(super::LocalSource), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -1156,6 +1159,25 @@ pub struct FilesSource { #[prost(message, optional, tag = "4")] pub schema: ::core::option::Option, } +/// Arrow data a plan reads (Spark's LocalRelation), named by content hash. The +/// data is never in the plan: it lives in the session's relation store of the +/// process that runs the plan, found by `hash`; a process without it fails. +/// feature: local_relation_source +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct LocalSource { + /// SHA-256 content hash (32 bytes). + #[prost(bytes = "vec", tag = "1")] + pub hash: ::prost::alloc::vec::Vec, + /// Arrow IPC `Schema` flatbuffer of the stored batches. + #[prost(bytes = "vec", tag = "2")] + pub schema: ::prost::alloc::vec::Vec, + #[prost(uint64, tag = "3")] + pub rows: u64, + #[prost(uint64, tag = "4")] + pub batches: u64, + #[prost(uint64, tag = "5")] + pub bytes: u64, +} /// A data source option: plain text or a reference to a secret. An inline /// secret or a URL carrying credentials is not encodable (ADR 0038). #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] diff --git a/crates/krishiv-plan-codec/tests/common/inventory.rs b/crates/krishiv-plan-codec/tests/common/inventory.rs index 67a9c9b9f..0b137f172 100644 --- a/crates/krishiv-plan-codec/tests/common/inventory.rs +++ b/crates/krishiv-plan-codec/tests/common/inventory.rs @@ -2363,7 +2363,24 @@ pub(crate) fn typed() -> Vec<(&'static str, PhysNode)> { ])); PhysNode::Scan(scan) }; + // A scan of a local relation, named by hash (no data in the plan). + let local_relation_source = { + use arrow_schema::{Field, Schema}; + let mut scan = file_scan(); + scan.scan.source = ReadSource::Local(krishiv_plan::LocalRelation::from_decoded( + std::sync::Arc::new(Schema::new(vec![ + Field::new("a", arrow_schema::DataType::Int32, false), + Field::new("b", arrow_schema::DataType::Utf8, true), + ])), + [0x5a; 32], + 2, + 1_000, + 4_096, + )); + PhysNode::Scan(scan) + }; vec![ + ("local_relation_source", local_relation_source), ("file_source_schema", file_source_schema), ("interval_literals", interval_literals), ("options_map_literal", options_map), diff --git a/crates/krishiv-plan-codec/tests/features.rs b/crates/krishiv-plan-codec/tests/features.rs index 5193ff12f..849d1c0de 100644 --- a/crates/krishiv-plan-codec/tests/features.rs +++ b/crates/krishiv-plan-codec/tests/features.rs @@ -2303,3 +2303,96 @@ fn a_declared_file_schema_that_is_not_a_struct_is_corrupt() { assert!(seen); corrupt(&m, "a declared file schema that is not a struct"); } + +// ---- local relation scans (local_relation_source) -------------------------- + +fn local_source_mut(m: &mut pb::PhysicalPlan, f: impl Fn(&mut pb::LocalSource)) -> bool { + let mut seen = false; + for node in &mut m.nodes { + if let Some(pb::node::Kind::Scan(s)) = &mut node.kind + && let Some(pb::ReadSource { + kind: Some(pb::read_source::Kind::Local(l)), + }) = &mut s.source + { + f(l); + seen = true; + } + } + seen +} + +#[test] +fn a_local_relation_scan_needs_its_feature() { + let m = message("local_relation_source"); + assert_eq!(Feature::LocalRelationSource.kind(), FeatureKind::Semantic); + assert_eq!(m.features, ["local_relation_source"]); + let bytes = m.encode_to_vec(); + assert_eq!( + decode_with(&bytes, &Limits::default(), &ReaderFeatures::only([])).unwrap_err(), + CodecError::UnsupportedFeatures { + missing: vec!["local_relation_source".to_owned()] + } + ); + let reader = ReaderFeatures::only([Feature::LocalRelationSource]); + let plan = decode_with(&bytes, &Limits::default(), &reader).unwrap(); + assert_eq!(encode(&plan).unwrap(), bytes); + assert_eq!( + feature_of(Construct::LocalRelationSource), + Some(Feature::LocalRelationSource) + ); +} + +/// The plan carries the hash and counts, never rows, and decodes without a lease. +#[test] +fn a_decoded_local_relation_holds_the_hash_and_counts_and_no_lease() { + let bytes = message("local_relation_source").encode_to_vec(); + let plan = decode(&bytes).unwrap(); + let mut found = false; + let mut stack = vec![plan.root()]; + while let Some(n) = stack.pop() { + if let PhysNode::Scan(s) = n + && let krishiv_plan::ReadSource::Local(l) = &s.scan.source + { + assert_eq!(l.relation().hash.as_bytes(), &[0x5a; 32]); + assert_eq!((l.batches(), l.rows(), l.bytes()), (2, 1_000, 4_096)); + assert!(l.lease().is_none()); + assert_eq!(l.schema().fields().len(), 2); + found = true; + } + stack.extend(n.children()); + } + assert!(found); +} + +#[test] +fn a_local_relation_hash_of_the_wrong_length_is_corrupt() { + let mut m = message("local_relation_source"); + assert!(local_source_mut(&mut m, |l| l.hash.truncate(5))); + corrupt(&m, "a local relation hash that is not 32 bytes"); +} + +#[test] +fn a_local_relation_with_a_nested_column_lists_the_type_feature_too() { + use arrow_schema::{DataType as A, Field, Schema}; + let mut scan = inventory::file_scan(); + scan.scan.source = krishiv_plan::ReadSource::Local(krishiv_plan::LocalRelation::from_decoded( + std::sync::Arc::new(Schema::new(vec![Field::new( + "xs", + A::List(std::sync::Arc::new(Field::new("item", A::Int32, true))), + true, + )])), + [1; 32], + 1, + 1, + 1, + )); + let m = to_message(&plan_of(PhysNode::Scan(scan))).unwrap(); + assert_eq!(m.features, ["local_relation_source", "nested_types"]); + let only_local = ReaderFeatures::only([Feature::LocalRelationSource]); + assert_eq!( + decode_with(&m.encode_to_vec(), &Limits::default(), &only_local).unwrap_err(), + CodecError::UnsupportedFeatures { + missing: vec!["nested_types".to_owned()] + } + ); +} diff --git a/crates/krishiv-plan-codec/tests/golden/local_relation_source/local_relation_source.bin b/crates/krishiv-plan-codec/tests/golden/local_relation_source/local_relation_source.bin new file mode 100644 index 0000000000000000000000000000000000000000..7f9ac99a46134a21d47116d4a01c959fd29d3546 GIT binary patch literal 666 zcma)2Jxc>Y5Z&3_yi`8fEDmpFY zvmP1~|HDyyFIPcbfCoC@!d4K}1ot4rq}HfN(5+RG=hnvigfKm5-0@{2(dS>0H^@UI z<)=ka&^P9zNQL_nRs2#F7lF(uL8BSu{ZE}kJqR?8z|LeQsi-e$`pDnxDBouOG8?mR z-J<;pYuNo?Ha2V0It~ZMgSg!4fd^MjDeZPN9GL*1xhbu7ni|jPhek&+y36=}rR8Zg ziV~wPb*7?Wl8(dvkuSh4UMhZmTTGvy9;>^KPIP>ona*~c_H^7Ghp{$Mnsl@DV%!VQ F;0vfuQ 6, + | "file_source_schema" + | "local_relation_source" => 6, // Features added after the version was frozen: above every // sequential format version. n if n.starts_with("time_") diff --git a/crates/krishiv-plan/CLAUDE.md b/crates/krishiv-plan/CLAUDE.md index d79693662..11417369e 100644 --- a/crates/krishiv-plan/CLAUDE.md +++ b/crates/krishiv-plan/CLAUDE.md @@ -13,6 +13,7 @@ decoder in `krishiv-connect`, native DataFrame APIs) lower into the IR. ## Key files Paths are relative to this crate. - `src/plan.rs` relation nodes (`LogicalPlan`); `src/expr.rs` expression nodes +- `src/local.rs` `LocalRelation`, the carrier of `ReadSource::Local` (Spark's `LocalRelation`): the content hash, counts and an optional `RelationLease` over the session's `krishiv_memory::RelationStore` (the batches live there once; built only from a lease, or by hash after a codec decode); `Debug` never prints values; equality ignores the lease. The analyzer resolves a `Read(Local)` to a `Scan` over its Arrow schema's columns (no catalog lookup, no securable; prefix-cache cacheable) - `src/origin.rs` origins (excluded from the encoding); `src/literal.rs` literal values; `src/function.rs` function references - `src/encode.rs` canonical encoding and `StructuralHash`; `src/dag.rs` DAG-aware equality and transforms - `src/walk.rs` child mapping and tree rewrites; `src/display.rs` pretty-printer (insta snapshots) diff --git a/crates/krishiv-plan/Cargo.toml b/crates/krishiv-plan/Cargo.toml index 6b74a7ee3..847ad5c23 100644 --- a/crates/krishiv-plan/Cargo.toml +++ b/crates/krishiv-plan/Cargo.toml @@ -19,6 +19,8 @@ krishiv-functions = { path = "../krishiv-functions", version = "=0.2.0" } krishiv-io = { path = "../krishiv-io", version = "=0.2.0" } # Spark's decimal result types, shared with the kernels that compute them. krishiv-kernels = { path = "../krishiv-kernels", version = "=0.2.0" } +# The relation store a local relation's data lives in (layer L1). +krishiv-memory = { path = "../krishiv-memory", version = "=0.2.0" } # Session UDFs, for constant folding through the registry. krishiv-udf-sdk = { path = "../krishiv-udf-sdk", version = "=0.2.0" } krishiv-types = { path = "../krishiv-types", version = "=0.2.0" } diff --git a/crates/krishiv-plan/src/analyzer/bind.rs b/crates/krishiv-plan/src/analyzer/bind.rs index c24e130d3..488b95a74 100644 --- a/crates/krishiv-plan/src/analyzer/bind.rs +++ b/crates/krishiv-plan/src/analyzer/bind.rs @@ -215,6 +215,9 @@ fn visit<'a>( PlanNode::Read(source @ ReadSource::Files { .. }) => { found.push(Found::Files(source, p.origin)); } + // `Read(Local)` falls to the wildcard below on purpose: + // content-addressed data in the session's store has nothing + // to load and no securable to check. // A plan resolved elsewhere is checked like a fresh one. PlanNode::Scan(s) => match &s.source { ReadSource::Table { name, options } => { @@ -230,6 +233,7 @@ fn visit<'a>( source @ ReadSource::Files { .. } => { found.push(Found::FileScan(source, &s.output, p.origin)); } + ReadSource::Local(_) => {} }, PlanNode::WithCte { input, ctes } => { p.for_each_expr(&mut |e| next.push((Node::Expr(e), Arc::clone(env)))); diff --git a/crates/krishiv-plan/src/analyzer/hints.rs b/crates/krishiv-plan/src/analyzer/hints.rs index f7983baee..9983008b5 100644 --- a/crates/krishiv-plan/src/analyzer/hints.rs +++ b/crates/krishiv-plan/src/analyzer/hints.rs @@ -160,7 +160,8 @@ impl Hints { PlanNode::SubqueryAlias { alias, .. } => Some(alias_parts(alias)), PlanNode::Scan(s) => match &s.source { ReadSource::Table { name, .. } => Some(name.clone()), - ReadSource::Files { .. } => Some(Vec::new()), + // Neither a file read nor a local relation has a table name. + ReadSource::Files { .. } | ReadSource::Local(_) => Some(Vec::new()), }, PlanNode::CteRef(r) => Some(vec![r.name.clone()]), PlanNode::JoinHint { @@ -171,7 +172,7 @@ impl Hints { PlanNode::SubqueryAlias { alias, .. } => Some(alias_parts(alias)), PlanNode::Scan(s) => match &s.source { ReadSource::Table { name, .. } => Some(name.clone()), - ReadSource::Files { .. } => None, + ReadSource::Files { .. } | ReadSource::Local(_) => None, }, PlanNode::CteRef(r) => Some(vec![r.name.clone()]), _ => None, diff --git a/crates/krishiv-plan/src/analyzer/ops.rs b/crates/krishiv-plan/src/analyzer/ops.rs index 7abd3204e..ae3236ddd 100644 --- a/crates/krishiv-plan/src/analyzer/ops.rs +++ b/crates/krishiv-plan/src/analyzer/ops.rs @@ -246,6 +246,37 @@ impl Resolver<'_> { rel } + /// A read of a local relation: a scan without a qualifier over the + /// columns of its Arrow schema, with fresh attributes. + pub(super) fn local( + &mut self, + p: &LogicalPlan, + l: &crate::LocalRelation, + ) -> Result { + let mut out = Vec::with_capacity(l.schema().fields().len()); + for f in l.schema().fields() { + let data_type = krishiv_types::DataType::from_arrow_field(f) + .map_err(|e| AnalysisError::new(e, p.origin))?; + out.push(Attr { + col: BoundColumn { + id: self.fresh(), + name: f.name().clone(), + data_type, + nullable: f.is_nullable(), + }, + qual: no_qual(), + source: None, + duplicate: false, + }); + } + let node = PlanNode::Scan(Scan { + source: ReadSource::Local(l.clone()), + output: out.iter().map(|a| a.col.clone()).collect(), + version: None, + }); + Ok(self.make(p.origin, node, out, &[])) + } + /// `c` as an instance: every id inside shifted by `offset`; `plan` is /// the `Shift` node when it exists already. Ids past `MAX_ID` are an /// error, never an overflow. diff --git a/crates/krishiv-plan/src/analyzer/resolve.rs b/crates/krishiv-plan/src/analyzer/resolve.rs index ccff9b9e7..2c2b229a3 100644 --- a/crates/krishiv-plan/src/analyzer/resolve.rs +++ b/crates/krishiv-plan/src/analyzer/resolve.rs @@ -477,6 +477,8 @@ impl Resolver<'_> { PlanNode::Scan(_) | PlanNode::Sql { .. } | PlanNode::Read(ReadSource::Files { .. }) => self.uncacheable = true, + // `Read(Local)` is cacheable on purpose (content-addressed, + // so a kept analysis of it stays valid): it falls to the wildcard. PlanNode::Read(ReadSource::Table { name, .. }) => { self.read_tables.push(name.clone()); } @@ -944,7 +946,7 @@ impl Resolver<'_> { PlanNode::TableFunction { name, args } => self.table_function(p, name, args, env), // Listed, not wildcarded: a new `PlanNode` variant must be // placed here or in `leaf_node` to compile. - PlanNode::Read(ReadSource::Files { .. }) + PlanNode::Read(ReadSource::Files { .. } | ReadSource::Local(_)) | PlanNode::Sql { .. } | PlanNode::Scan(_) | PlanNode::CteRef(_) @@ -968,6 +970,7 @@ impl Resolver<'_> { None => Err(unsupported("reading files by path", o)), } } + PlanNode::Read(ReadSource::Local(l)) => self.local(p, l), PlanNode::Sql { .. } => { let root = self.b.sql.get(&addr(&**p)).cloned(); let env0 = self.root_env.clone(); @@ -979,7 +982,7 @@ impl Resolver<'_> { PlanNode::Scan(s) => { let name = match &s.source { ReadSource::Table { name, .. } => name.clone(), - ReadSource::Files { .. } => Vec::new(), + ReadSource::Files { .. } | ReadSource::Local(_) => Vec::new(), }; let source = self.b.scans.get(&name).copied(); // Bind checked every scan of a table; one it did not bind diff --git a/crates/krishiv-plan/src/cache.rs b/crates/krishiv-plan/src/cache.rs index f3d7774a9..604cc00bb 100644 --- a/crates/krishiv-plan/src/cache.rs +++ b/crates/krishiv-plan/src/cache.rs @@ -204,6 +204,16 @@ impl Encoder { self.out.extend_from_slice(s.as_bytes()); } + /// An Arrow schema as names, Spark types and nullability (never data). + fn arrow_schema(&mut self, schema: &arrow_schema::Schema) { + self.num(schema.fields().len() as u64); + for f in schema.fields() { + self.text(f.name()); + self.text(&crate::local::type_text(f)); + self.flag(f.is_nullable()); + } + } + fn debug(&mut self, v: &T) { let s = format!("{v:?}"); self.text(&s); @@ -277,6 +287,13 @@ impl Encoder { self.option_value(value); } } + ReadSource::Local(l) => { + self.tag("local"); + self.text(&l.relation().hash.hex()); + self.arrow_schema(l.schema()); + self.num(l.batches() as u64); + self.num(l.rows()); + } ReadSource::Files { format, paths, diff --git a/crates/krishiv-plan/src/display.rs b/crates/krishiv-plan/src/display.rs index 5f349ec17..4d7435c24 100644 --- a/crates/krishiv-plan/src/display.rs +++ b/crates/krishiv-plan/src/display.rs @@ -229,6 +229,26 @@ struct NodeLine<'a>(&'a LogicalPlan); impl Display for NodeLine<'_> { fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { match &self.0.node { + // Spark's `LocalRelation [a#1, b#2]`; the row count is Krishiv's + // own, as for an inline table. + PlanNode::Read(ReadSource::Local(l)) => { + f.write_str("LocalRelation [")?; + for (i, field) in l.schema().fields().iter().enumerate() { + write!(f, "{}{}", sep(i), field.name())?; + } + write!(f, "] ({} rows)", l.rows()) + } + PlanNode::Scan(scan) if matches!(scan.source, ReadSource::Local(_)) => { + f.write_str("LocalRelation [")?; + for (i, c) in scan.output.iter().enumerate() { + write!(f, "{}{}#{}", sep(i), c.name, c.id.0)?; + } + f.write_str("]")?; + if let ReadSource::Local(l) = &scan.source { + write!(f, " ({} rows)", l.rows())?; + } + Ok(()) + } PlanNode::Read(source) => write!(f, "Read {}", SourceText(source)), PlanNode::OneRowRelation => f.write_str("OneRowRelation"), PlanNode::Sql { query, args } => { @@ -565,6 +585,11 @@ impl Display for SourceText<'_> { } options } + // Printed by the node line, which has the columns. + ReadSource::Local(l) => { + write!(f, "LocalRelation ({} rows)", l.rows())?; + return Ok(()); + } }; if !options.is_empty() { f.write_str(" {")?; diff --git a/crates/krishiv-plan/src/encode.rs b/crates/krishiv-plan/src/encode.rs index 2a0e1e5fd..a0d48e07b 100644 --- a/crates/krishiv-plan/src/encode.rs +++ b/crates/krishiv-plan/src/encode.rs @@ -704,6 +704,16 @@ impl<'k, 'a> Encoder<'k, 'a> { } } + /// An Arrow schema as names, Spark types and nullability (never data). + fn arrow_schema(&mut self, schema: &arrow_schema::Schema) { + self.len(schema.fields().len()); + for f in schema.fields() { + self.str(f.name()); + self.str(&crate::local::type_text(f)); + self.bool(f.is_nullable()); + } + } + fn data_type(&mut self, t: &krishiv_types::DataType) { self.str(&t.simple_string()); } @@ -1111,6 +1121,13 @@ impl<'k, 'a> Encoder<'k, 'a> { self.strs(name); self.options(options); } + ReadSource::Local(l) => { + self.u8(3); + self.put(l.relation().hash.as_bytes()); + self.arrow_schema(l.schema()); + self.len(l.batches()); + self.put(&l.rows().to_le_bytes()); + } ReadSource::Files { format, paths, diff --git a/crates/krishiv-plan/src/lib.rs b/crates/krishiv-plan/src/lib.rs index 69b982b44..f62394f8c 100644 --- a/crates/krishiv-plan/src/lib.rs +++ b/crates/krishiv-plan/src/lib.rs @@ -114,9 +114,13 @@ //! | `ResolvedExpr::StateInput` | R | - | - | //! | `ResolvedExpr::EvalMode` | R | - | - | //! +//! `ReadSource::Local` (Spark's `LocalRelation`, see [`LocalRelation`]) is a +//! `PlanNode::Read` of data held in the session's relation store; it resolves +//! to a `PlanNode::Scan` and is not a node of its own. +//! //! Not in v1 (spec 13 §9; added with the features that need them): //! `Generate`, `Hint`, `Pivot`, `Command`, -//! `Extension`, `ToDF`, `Offset`, `Deduplicate`, `LocalRelation` and `Range` +//! `Extension`, `ToDF`, `Offset`, `Deduplicate` and `Range` //! relations; lambda, extract-value and extension expressions. #![forbid(unsafe_code)] @@ -129,6 +133,7 @@ mod encode; mod expr; pub mod function; mod literal; +mod local; mod origin; mod partial; mod plan; @@ -151,6 +156,7 @@ pub use expr::{ }; pub use function::{Builtin, FunctionId, FunctionKind, PythonEvalType, QuantifiedLike}; pub use literal::{DayTimeField, Literal, YearMonthField}; +pub use local::LocalRelation; pub use origin::{Origin, PlanId, SourceId, Span}; pub use partial::partial_state_columns; /// The qualifier prefix of the columns a select-list generator adds: an diff --git a/crates/krishiv-plan/src/local.rs b/crates/krishiv-plan/src/local.rs new file mode 100644 index 000000000..add449705 --- /dev/null +++ b/crates/krishiv-plan/src/local.rs @@ -0,0 +1,132 @@ +// SPDX-License-Identifier: Apache-2.0 +//! [`LocalRelation`]: Arrow data a plan reads (Spark's `LocalRelation`). +//! +//! The batches live once, in the session's [`krishiv_memory::RelationStore`]; +//! the plan holds the content hash, the counts and (when built in this +//! process) a lease that keeps the entry alive while the plan does. A plan +//! decoded from the plan codec holds no lease: the provider resolves it by +//! hash in the same process's store or fails loudly. + +use std::fmt; + +use arrow_schema::SchemaRef; +use krishiv_memory::{RelationHash, RelationLease, RelationRef}; + +/// Arrow data a plan reads. See the module docs. +#[derive(Clone)] +pub struct LocalRelation { + schema: SchemaRef, + relation: RelationRef, + batches: usize, + rows: u64, + bytes: u64, + lease: Option, +} + +impl LocalRelation { + /// A relation over the entry `lease` holds. `schema` is the Arrow + /// schema of the stored batches; the hash is the lease's, so it is + /// correct by construction. + #[must_use] + pub fn from_lease(schema: SchemaRef, lease: &RelationLease, rows: u64, bytes: u64) -> Self { + Self { + schema, + relation: lease.reference(), + batches: lease.batches(), + rows, + bytes, + lease: Some(lease.clone()), + } + } + + /// A relation named by its 32-byte content hash, holding no lease (the + /// plan codec's decode: the entry is found by hash where the plan runs). + #[must_use] + pub fn from_decoded( + schema: SchemaRef, + hash: [u8; 32], + batches: usize, + rows: u64, + bytes: u64, + ) -> Self { + Self { + schema, + relation: RelationRef { + hash: RelationHash::from_bytes(hash), + }, + batches, + rows, + bytes, + lease: None, + } + } + + /// The Arrow schema of the stored batches. + #[must_use] + pub fn schema(&self) -> &SchemaRef { + &self.schema + } + + /// The content hash the data is found by. + #[must_use] + pub fn relation(&self) -> RelationRef { + self.relation + } + + /// The number of stored batches. + #[must_use] + pub fn batches(&self) -> usize { + self.batches + } + + /// The number of rows. + #[must_use] + pub fn rows(&self) -> u64 { + self.rows + } + + /// The decoded size in bytes at insertion. + #[must_use] + pub fn bytes(&self) -> u64 { + self.bytes + } + + /// The lease keeping the entry alive; `None` after a codec decode. + #[must_use] + pub fn lease(&self) -> Option<&RelationLease> { + self.lease.as_ref() + } +} + +impl PartialEq for LocalRelation { + fn eq(&self, other: &Self) -> bool { + self.relation == other.relation + && self.schema == other.schema + && self.batches == other.batches + && self.rows == other.rows + } +} + +impl Eq for LocalRelation {} + +impl fmt::Debug for LocalRelation { + /// The shape, counts and a hash prefix; never values. + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let hex = self.relation.hash.hex(); + f.debug_struct("LocalRelation") + .field("columns", &self.schema.fields().len()) + .field("batches", &self.batches) + .field("rows", &self.rows) + .field("hash", &hex.get(..12).unwrap_or(&hex)) + .finish_non_exhaustive() + } +} + +/// The Spark type of `field` as text for the canonical encodings; the Arrow +/// type's name when it has no Spark type (such a relation does not resolve). +pub(crate) fn type_text(field: &arrow_schema::Field) -> String { + krishiv_types::DataType::from_arrow_field(field).map_or_else( + |_| format!("arrow:{:?}", field.data_type()), + |t| t.simple_string(), + ) +} diff --git a/crates/krishiv-plan/src/plan.rs b/crates/krishiv-plan/src/plan.rs index 718dab821..718aa68ae 100644 --- a/crates/krishiv-plan/src/plan.rs +++ b/crates/krishiv-plan/src/plan.rs @@ -705,6 +705,10 @@ pub enum ReadSource { /// [`ReadSource::with_schema`]. schema: Option, }, + /// Arrow data held in the session's relation store (Spark's + /// `LocalRelation`). Built only from a lease; the analyzer resolves it + /// to a [`Scan`] with the Arrow schema's columns. + Local(crate::LocalRelation), } impl ReadSource { @@ -739,7 +743,7 @@ impl ReadSource { options, schema: Some(schema), }, - table @ Self::Table { .. } => table, + other @ (Self::Table { .. } | Self::Local(_)) => other, } } } diff --git a/crates/krishiv-plan/src/tests/local.rs b/crates/krishiv-plan/src/tests/local.rs new file mode 100644 index 000000000..1de438959 --- /dev/null +++ b/crates/krishiv-plan/src/tests/local.rs @@ -0,0 +1,189 @@ +// SPDX-License-Identifier: Apache-2.0 +//! `ReadSource::Local` (P3-T012a part 2a): the carrier's equality, `Debug` +//! and canonical encoding, its analysis to a `Scan`, and the prefix cache. + +use std::sync::Arc; + +use arrow_array::{ArrayRef, Int32Array, RecordBatch, StringArray}; +use arrow_schema::{DataType as ArrowType, Field, Schema, SchemaRef}; +use krishiv_memory::RelationStore; +use krishiv_types::DataType; +use krishiv_types::batch::Batch; + +use super::dsl::{add, int}; +use crate::analyzer::testing::{block_on, context}; +use crate::analyzer::{ + Analyzer as _, CachingAnalyzer, DefaultAnalyzer, MemoryCatalog, ResolutionConf, +}; +use crate::{ + ColumnRef, Expr, ExprNode, LocalRelation, LogicalPlan, Origin, Plan, PlanId, PlanNode, + ReadSource, normalized_plan_hash, +}; + +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", ArrowType::Int32, false), + Field::new("b", ArrowType::Utf8, true), + ])) +} + +fn relation(store: &Arc, ids: &[i32]) -> LocalRelation { + let s = schema(); + let cols: Vec = vec![ + Arc::new(Int32Array::from(ids.to_vec())), + Arc::new(StringArray::from(vec![Some("secret-value"); ids.len()])), + ]; + let batch = RecordBatch::try_new(Arc::clone(&s), cols).unwrap(); + let lease = store.insert(&s, &[Batch::new(batch)]).unwrap(); + LocalRelation::from_lease(s, &lease, ids.len() as u64, 0) +} + +fn read(l: &LocalRelation) -> LogicalPlan { + LogicalPlan::new(PlanNode::Read(ReadSource::Local(l.clone()))) +} + +fn plan(p: LogicalPlan) -> Plan { + Plan::try_new(p).unwrap() +} + +#[test] +fn debug_prints_the_shape_and_never_values() { + let store = RelationStore::new(); + let l = relation(&store, &[1, 2, 3]); + let text = format!("{l:?}"); + assert!(text.contains("rows: 3"), "{text}"); + assert!(!text.contains("secret-value"), "{text}"); + let in_plan = format!("{:?}", read(&l)); + assert!(!in_plan.contains("secret-value"), "{in_plan}"); +} + +#[test] +fn equal_content_is_structurally_equal_and_different_content_is_not() { + let store = RelationStore::new(); + let (a, b) = (relation(&store, &[1, 2]), relation(&store, &[1, 2])); + let other = relation(&store, &[1, 2, 3]); + assert_eq!(a, b); + assert_ne!(a, other); + let (pa, pb, po) = (plan(read(&a)), plan(read(&b)), plan(read(&other))); + assert!(pa.structurally_eq(&pb)); + assert_eq!(pa.canonical_encoding(), pb.canonical_encoding()); + assert_eq!(pa.structural_hash(), pb.structural_hash()); + assert_eq!(normalized_plan_hash(&pa), normalized_plan_hash(&pb)); + assert!(!pa.structurally_eq(&po)); + assert_ne!(pa.canonical_encoding(), po.canonical_encoding()); + assert_ne!(normalized_plan_hash(&pa), normalized_plan_hash(&po)); + // A decoded carrier (no lease) equals the leased one. + let decoded = LocalRelation::from_decoded( + Arc::clone(a.schema()), + *a.relation().hash.as_bytes(), + a.batches(), + a.rows(), + a.bytes(), + ); + assert_eq!(decoded, a); + assert!(plan(read(&decoded)).structurally_eq(&pa)); +} + +#[test] +fn a_read_resolves_to_a_scan_with_the_arrow_columns() { + let store = RelationStore::new(); + let l = relation(&store, &[1]); + let c = context( + MemoryCatalog::new("spark_catalog"), + ResolutionConf::default(), + ); + let analyzed = block_on(DefaultAnalyzer.analyze(&plan(read(&l)), &c)).unwrap(); + assert_eq!(analyzed.output.len(), 2); + assert_eq!(analyzed.output[0].name, "a"); + assert_eq!(analyzed.output[0].data_type, DataType::Integer); + assert!(!analyzed.output[0].nullable); + assert_eq!(analyzed.output[1].data_type, DataType::String); + assert!(analyzed.output[1].nullable); + assert!(analyzed.securables.is_empty()); + let PlanNode::Scan(scan) = &analyzed.plan.root().node else { + panic!("not a scan: {:?}", analyzed.plan.root().node); + }; + assert!(matches!(scan.source, ReadSource::Local(_))); + // Re-analysing the resolved plan keeps the source. + let again = block_on(DefaultAnalyzer.analyze(&analyzed.plan, &c)).unwrap(); + assert_eq!(again.output, analyzed.output); +} + +#[test] +fn an_arrow_type_without_a_spark_type_is_unsupported() { + let store = RelationStore::new(); + let s: SchemaRef = Arc::new(Schema::new(vec![Field::new( + "t", + ArrowType::Time64(arrow_schema::TimeUnit::Nanosecond), + true, + )])); + let batch = RecordBatch::new_empty(Arc::clone(&s)); + let lease = store.insert(&s, &[Batch::new(batch)]).unwrap(); + let l = LocalRelation::from_lease(s, &lease, 0, 0); + let c = context( + MemoryCatalog::new("spark_catalog"), + ResolutionConf::default(), + ); + let err = block_on(DefaultAnalyzer.analyze(&plan(read(&l)), &c)).unwrap_err(); + assert_eq!( + err.error.error_class(), + krishiv_common::error::ErrorClass::UnsupportedTimeType + ); +} + +#[test] +fn the_prefix_cache_reuses_a_local_relation() { + let store = RelationStore::new(); + let l = relation(&store, &[1, 2]); + let c = context( + MemoryCatalog::new("spark_catalog"), + ResolutionConf::default(), + ); + let col = |name: &str, id: i64| { + Expr::new(ExprNode::Column(ColumnRef { + plan_id: Some(PlanId(id)), + ..ColumnRef::new(vec![name.to_owned()]) + })) + }; + let chain = |steps: i64| { + let mut p = read(&l).with_origin(Origin::plan(PlanId(0))); + for i in 1..=steps { + let prev = if i == 1 { + "a".to_owned() + } else { + format!("c{}", i - 1) + }; + p = LogicalPlan::new(PlanNode::WithColumns { + input: Arc::new(p), + columns: vec![(format!("c{i}"), Arc::new(add(col(&prev, i - 1), int(1))))], + }) + .with_origin(Origin::plan(PlanId(i))); + } + plan(p) + }; + let cache = CachingAnalyzer::new(); + block_on(cache.analyze(&chain(2), &c)).unwrap(); + for depth in 3..=5 { + let before = cache.warm_starts(); + let warm = block_on(cache.analyze(&chain(depth), &c)).unwrap(); + assert!(cache.warm_starts() > before, "depth {depth} starts warm"); + assert_eq!( + warm, + block_on(DefaultAnalyzer.analyze(&chain(depth), &c)).unwrap() + ); + } +} + +#[test] +fn display_names_the_columns_and_rows() { + let store = RelationStore::new(); + let l = relation(&store, &[1, 2]); + assert_eq!(read(&l).to_string(), "LocalRelation [a, b] (2 rows)\n"); + let c = context( + MemoryCatalog::new("spark_catalog"), + ResolutionConf::default(), + ); + let analyzed = block_on(DefaultAnalyzer.analyze(&plan(read(&l)), &c)).unwrap(); + // Spark-format text is golden-pending; the row count is Krishiv's own. + insta::assert_snapshot!(analyzed.plan.root().to_string(), @"LocalRelation [a#0, b#1] (2 rows)"); +} diff --git a/crates/krishiv-plan/src/tests/mod.rs b/crates/krishiv-plan/src/tests/mod.rs index 3788244c1..2cd453404 100644 --- a/crates/krishiv-plan/src/tests/mod.rs +++ b/crates/krishiv-plan/src/tests/mod.rs @@ -15,6 +15,7 @@ mod files; mod function_call_forms; mod inventory; mod json_tuple; +mod local; mod nested_literals; mod nodes; mod pivot; diff --git a/crates/krishiv-plan/src/tests/props.rs b/crates/krishiv-plan/src/tests/props.rs index 31665a3f2..26e79f19c 100644 --- a/crates/krishiv-plan/src/tests/props.rs +++ b/crates/krishiv-plan/src/tests/props.rs @@ -1091,6 +1091,13 @@ fn m_source(s: &mut ReadSource, w: &mut Walk) -> bool { name: paths.iter().map(ToString::to_string).collect(), options: options.clone(), }, + ReadSource::Local(l) => ReadSource::Local(crate::LocalRelation::from_decoded( + std::sync::Arc::clone(l.schema()), + [7; 32], + l.batches(), + l.rows() + 1, + l.bytes(), + )), }; site!(w, *s = other); match s { @@ -1101,6 +1108,8 @@ fn m_source(s: &mut ReadSource, w: &mut Walk) -> bool { options, .. } => m_string(format, w) || m_values(paths, w) || m_options(options, w), + // Nothing inside to mutate: the hash and counts are the whole value. + ReadSource::Local(_) => false, } } diff --git a/crates/krishiv-session/CLAUDE.md b/crates/krishiv-session/CLAUDE.md index c5fd9dfbe..4e725c55e 100644 --- a/crates/krishiv-session/CLAUDE.md +++ b/crates/krishiv-session/CLAUDE.md @@ -21,6 +21,7 @@ Paths are relative to this crate. `snapshot` is the stored configuration, which `SET` lists); `local_zone` reads the machine's zone as Java does - `src/catalog.rs` session catalog (temporary tables and views; `catalog/scans.rs` `reads_stored_table` (the scans of a plan and its subqueries, each shared expression node once); `catalog/file_tables.rs` `QueryCatalog::open_files` lists and opens the file tables a plan scans, then `catalog/ndv_sampling.rs` samples the distinct values of the file columns an equality or `IN` in its `Filter` conditions compares, P2-T164 U6 (and, while `krishiv.optimizer.groupKeyCode` is on, the group keys of an aggregate over a join that the rule could code: codable types wider than the code, a table at most half the largest other table below it; design P2b D2), within one time budget per query (`krishiv.optimizer.ndvSampling.timeout`), and `sample_notes` gives `EXPLAIN COST`/`VERBOSE` a line per table, `tests/ndv_sampling.rs`; the permanent catalog `store` and `warehouse` are the runtime's, shared by its sessions as Spark's `newSession()` shares them, `SessionCatalog::over`) - `src/carrier.rs` decimals and calendar intervals at the session's edges: `register_batches` converts to the canonical carrier (`Decimal64`, the calendar struct), results and `schema()` read `Decimal128` and `Interval(MonthDayNano)` (`edge_arrow_type`) +- `src/local_relation.rs` (feature `unstable-ir`) `Session::local_relation(ipc, user_schema)`: Arrow IPC to a `krishiv_plan::LocalRelation` over the session's `RelationStore` (Spark Connect `LocalRelation`). A user schema is applied eagerly (rename by position; a cast projection through this session's executor, conf and pool unless the types already match; non-nullable fields checked at any depth); only the result is stored. A scan of `ReadSource::Local` resolves through its own lease, else by hash in the store (`FileScans::local_lease`), and fails with `INTERNAL_ERROR` naming the hash when the entry is gone - `src/view_schema.rs` schema binding modes of permanent views (`WITH SCHEMA BINDING|COMPENSATION|TYPE EVOLUTION|EVOLUTION`, Spark's `ViewSchemaMode`): the read projection per mode, the checks against the analyzed view query, the evolved schema stored after analysis; `tests/views_schema_binding.rs` - `src/dataframe.rs` DataFrame transformations (IR only) and async actions; `src/column.rs` and `src/functions.rs` user API - `src/query_fold.rs` the bookkeeping around one running query: `QueryStart` (counters read at start), `QueryFold`/`MetricStream`/`PendingFold` (the one fold of its metrics into the engine series and `system.queries` when the stream ends, fails or is dropped; holds the admission permit), `QueuedRow` (the history row of a query waiting for admission) diff --git a/crates/krishiv-session/src/catalog.rs b/crates/krishiv-session/src/catalog.rs index da720d3f8..27fd7f61d 100644 --- a/crates/krishiv-session/src/catalog.rs +++ b/crates/krishiv-session/src/catalog.rs @@ -469,6 +469,32 @@ pub(crate) struct FileScans { pub(crate) slots: u64, /// The executor's I/O runtime, which opens the file tables. pub(crate) runtime: tokio::runtime::Handle, + /// The session's relation store, which holds the data of local relations. + pub(crate) relations: Arc, +} + +impl FileScans { + /// A lease on the entry of the local relation `l`: its own lease, else one + /// found by hash in the session's store. + /// + /// # Errors + /// `INTERNAL_ERROR` naming the hash when the store does not hold it (a + /// decoded plan run where its data is not): never an empty result. + pub(crate) fn local_lease( + &self, + l: &krishiv_plan::LocalRelation, + ) -> Result { + if let Some(lease) = l.lease() { + return Ok(lease.clone()); + } + let hash = l.relation().hash; + self.relations.lease(&hash).ok_or_else(|| { + InternalError { + message: format!("the local relation {hash} is not in this session's store"), + } + .into() + }) + } } /// A snapshot pinned by hidden name, and how many actions pinned it. @@ -806,6 +832,19 @@ impl QueryCatalog { })?; assigned.push((scan.id, self.file_splits(&table, total_slots)?)); } + ReadSource::Local(l) => { + // The batches live in the session's store; the splits name + // them by hash. A relation the store lost is an engine + // bug, never an empty scan. + let lease = self.scans.local_lease(l)?; + assigned.push(( + scan.id, + ScanSplitsKind::Batches { + relation: lease.reference(), + batches: lease.batches(), + }, + )); + } _ => { return Err(unsupported("this scan")); } @@ -1278,6 +1317,14 @@ impl TableProviderResolver for QueryCatalog { source @ ReadSource::Files { .. } => { self.files_provider(&FilesKey::Read(source.clone())) } + ReadSource::Local(l) => { + let lease = self.scans.local_lease(l)?; + Ok(Arc::new(krishiv_exec::planner::MemoryTable::from_lease( + Arc::clone(l.schema()), + l.rows(), + lease, + ))) + } _ => Err(unsupported("this scan")), } } diff --git a/crates/krishiv-session/src/lib.rs b/crates/krishiv-session/src/lib.rs index 21e9c9f51..475919e82 100644 --- a/crates/krishiv-session/src/lib.rs +++ b/crates/krishiv-session/src/lib.rs @@ -64,6 +64,8 @@ pub mod functions; mod history; mod identifier_clause; mod introspection; +#[cfg(feature = "unstable-ir")] +mod local_relation; mod locals; #[cfg(feature = "telemetry-logs")] pub mod logging; diff --git a/crates/krishiv-session/src/local_relation.rs b/crates/krishiv-session/src/local_relation.rs new file mode 100644 index 000000000..7b0ac57d0 --- /dev/null +++ b/crates/krishiv-session/src/local_relation.rs @@ -0,0 +1,381 @@ +// SPDX-License-Identifier: Apache-2.0 +//! [`Session::local_relation`]: Arrow IPC data (Spark Connect's `LocalRelation`) as a +//! [`LocalRelation`] carrier over the session's relation store (design P3-T012a part 2a, section 1). +//! +//! A user schema is applied **eagerly**, as Spark does (`SparkConnectPlanner.buildLocalRelationFromRows` +//! evaluates `Dataset.toDF(names).to(schema)` while the plan is built): the columns are renamed by +//! position, cast through the session's own executor when their types differ, and checked against the +//! schema's non-nullable fields; only the result is stored. The cast and nullability rules and their +//! error classes follow the current style and are pending Spark goldens. + +use std::sync::Arc; + +use arrow_array::cast::AsArray as _; +use arrow_array::{ + Array, ArrayRef, FixedSizeListArray, LargeListArray, ListArray, MapArray, RecordBatch, + RecordBatchOptions, StructArray, +}; +use arrow_schema::{DataType as ArrowType, Field, Schema, SchemaRef}; +use krishiv_common::error::UnableToInferSchema; +use krishiv_common::error::{KrishivError, NotNullAssertViolation, NumColumnsMismatch}; +use krishiv_memory::RelationStore; +use krishiv_plan::{Expr, LocalRelation, LogicalPlan, PlanNode, ReadSource}; +use krishiv_types::batch::Batch; +use krishiv_types::{DataType, StructType}; + +use crate::carrier::{canonical_schema, recarried}; +use crate::dataframe::DataFrame; +use crate::session::{Session, internal}; + +impl Session { + /// The relation of `data` (an Arrow IPC stream with its own schema) as a plan carrier, with + /// `user_schema` applied eagerly when given (Spark Connect's `LocalRelation`): + /// + /// - no user schema: the data as it is; + /// - a user schema: the columns are renamed by position and cast to its types through this + /// session's executor (its configuration and memory pool) unless the types already match; a + /// non-nullable field holding a null is an error. Only the result is stored; + /// - no data: an empty relation of the user schema. + /// + /// The batches live once in the session's relation store; the carrier holds a lease. + /// + /// # Errors + /// `UNABLE_TO_INFER_SCHEMA` for neither data nor schema; `NUM_COLUMNS_MISMATCH` when the + /// schema's column count differs from the data's; `NOT_NULL_ASSERT_VIOLATION` for a null in a + /// non-nullable field; the cast errors of the executor; `INTERNAL_ERROR` for undecodable data. + pub async fn local_relation( + &self, + data: Option<&[u8]>, + user_schema: Option, + ) -> Result { + let store = &self.state.relations; + let target = user_schema.as_ref().map(arrow_schema_of); + let Some(data) = data else { + let Some(target) = target else { + return Err(UnableToInferSchema { + format: "LocalRelation".to_owned(), + } + .into()); + }; + return insert(store, canonical_schema(Arc::new(target)), Vec::new()); + }; + let (decoded, batches) = RelationStore::decode_ipc(data)?; + let schema = canonical_schema(decoded); + let records = batches + .iter() + .map(|b| recarried(b.data(), &schema)) + .collect::, _>>()?; + let Some(target) = target else { + return insert(store, schema, records); + }; + let target = canonical_schema(Arc::new(target)); + if schema.fields().len() != target.fields().len() { + return Err(NumColumnsMismatch { + operator: "LocalRelation".to_owned(), + first_num_columns: target.fields().len().to_string(), + invalid_ordinal_num: "second".to_owned(), + invalid_num_columns: schema.fields().len().to_string(), + } + .into()); + } + let renamed = renamed(&schema, &target); + let same = renamed + .fields() + .iter() + .zip(target.fields()) + .all(|(a, b)| same_shape(a.data_type(), b.data_type())); + let converted = if same || records.is_empty() { + records + } else { + self.cast_batches(&renamed, &target, records).await? + }; + let conformed = converted + .iter() + .map(|b| conform(b, &target)) + .collect::, _>>()?; + insert(store, target, conformed) + } + + /// `records` (of `renamed`, the data's types under the target's names) cast to the types of + /// `target` by a projection run through this session's executor over a temporary relation. + async fn cast_batches( + &self, + renamed: &SchemaRef, + target: &SchemaRef, + records: Vec, + ) -> Result, KrishivError> { + let store = &self.state.relations; + // The temporary entry lives while this function does. The plan names it by hash only, so + // no cache of an analysis keeps it alive afterwards. + let named = records + .iter() + .map(|b| { + RecordBatch::try_new_with_options( + Arc::clone(renamed), + b.columns().to_vec(), + &RecordBatchOptions::new().with_row_count(Some(b.num_rows())), + ) + .map_err(|e| internal(e.to_string())) + }) + .collect::, _>>()?; + let temp = insert(store, Arc::clone(renamed), named)?; + let local = LocalRelation::from_decoded( + Arc::clone(temp.schema()), + *temp.relation().hash.as_bytes(), + temp.batches(), + temp.rows(), + temp.bytes(), + ); + let mut exprs = Vec::with_capacity(target.fields().len()); + for (have, want) in renamed.fields().iter().zip(target.fields()) { + let column = Expr::column([have.name().clone()]); + let to = DataType::from_arrow_field(want)?; + exprs.push( + if same_shape(have.data_type(), want.data_type()) { + column + } else { + column.cast(to) + } + .alias(want.name().clone()), + ); + } + let plan = LogicalPlan::new(PlanNode::Read(ReadSource::Local(local))).project(exprs); + let batches = DataFrame::new(self.clone(), Arc::new(plan)) + .collect() + .await?; + batches + .iter() + .map(|b| recarried(b, &canonical_schema(Arc::clone(b.schema_ref())))) + .collect() + } +} + +/// The Arrow schema of a user-declared struct type. +fn arrow_schema_of(schema: &StructType) -> Schema { + Schema::new( + schema + .fields() + .iter() + .map(|f| f.data_type().to_arrow_field(f.name(), f.nullable())) + .collect::>(), + ) +} + +/// `schema`'s types under `names`' field names (Spark's `toDF(names)`). +fn renamed(schema: &SchemaRef, names: &SchemaRef) -> SchemaRef { + Arc::new(Schema::new_with_metadata( + schema + .fields() + .iter() + .zip(names.fields()) + .map(|(f, n)| f.as_ref().clone().with_name(n.name().clone())) + .collect::>(), + schema.metadata().clone(), + )) +} + +/// Stores `records` (of `schema`) and returns the carrier over the entry. +fn insert( + store: &Arc, + schema: SchemaRef, + records: Vec, +) -> Result { + let rows: u64 = records.iter().map(|b| b.num_rows() as u64).sum(); + let bytes: u64 = records + .iter() + .map(|b| b.get_array_memory_size() as u64) + .sum(); + let batches: Vec = records.into_iter().map(Batch::new).collect(); + let lease = store.insert(&schema, &batches)?; + Ok(LocalRelation::from_lease(schema, &lease, rows, bytes)) +} + +/// Whether `a` and `b` are the same type once nested field names are ignored (Spark renames +/// nested fields `col_i` before comparing); nullability is checked against the data. +fn same_shape(a: &ArrowType, b: &ArrowType) -> bool { + match (a, b) { + (ArrowType::List(x), ArrowType::List(y)) + | (ArrowType::LargeList(x), ArrowType::LargeList(y)) => { + same_shape(x.data_type(), y.data_type()) + } + (ArrowType::FixedSizeList(x, n), ArrowType::FixedSizeList(y, m)) => { + n == m && same_shape(x.data_type(), y.data_type()) + } + (ArrowType::Map(x, s), ArrowType::Map(y, t)) => { + s == t && same_shape(x.data_type(), y.data_type()) + } + (ArrowType::Struct(x), ArrowType::Struct(y)) => { + x.len() == y.len() + && x.iter() + .zip(y.iter()) + .all(|(p, q)| same_shape(p.data_type(), q.data_type())) + } + _ => a == b, + } +} + +/// `batch` relabelled to `target` (names, nested field names, nullability flags) over the same +/// buffers, after checking that no non-nullable field, nested ones included, holds a null. +fn conform(batch: &RecordBatch, target: &SchemaRef) -> Result { + let columns = batch + .columns() + .iter() + .zip(target.fields()) + .map(|(c, f)| { + check_non_null(c, f, f.name())?; + relabel(c, f.data_type()) + }) + .collect::, _>>()?; + RecordBatch::try_new_with_options( + Arc::clone(target), + columns, + &RecordBatchOptions::new().with_row_count(Some(batch.num_rows())), + ) + .map_err(|e| internal(e.to_string())) +} + +/// The error for a null in the non-nullable field at `path`. +fn null_in(path: &str) -> KrishivError { + NotNullAssertViolation { + walked_type_path: path.to_owned(), + } + .into() +} + +/// Fails when `array` holds a null where `field` (and the fields under it) is non-nullable. A struct +/// child's nulls count only where the struct itself is valid; a list's elements are checked over +/// the whole values array (a null parent owns an empty range in the arrays Arrow writes). +fn check_non_null(array: &ArrayRef, field: &Field, path: &str) -> Result<(), KrishivError> { + if !field.is_nullable() && array.logical_null_count() > 0 { + return Err(null_in(path)); + } + match (array.data_type(), field.data_type()) { + (ArrowType::Struct(_), ArrowType::Struct(fields)) => { + let s = array.as_struct(); + let parent = array.logical_nulls(); + for (child, f) in s.columns().iter().zip(fields.iter()) { + let at = format!("{path}.{}", f.name()); + if !f.is_nullable() + && let Some(nulls) = child.logical_nulls() + && (0..child.len()) + .any(|i| nulls.is_null(i) && parent.as_ref().is_none_or(|p| p.is_valid(i))) + { + return Err(null_in(&at)); + } + check_non_null_below(child, f, &at)?; + } + Ok(()) + } + (ArrowType::List(_), ArrowType::List(f)) => check_non_null( + array.as_list::().values(), + f, + &format!("{path}.element"), + ), + (ArrowType::LargeList(_), ArrowType::LargeList(f)) => check_non_null( + array.as_list::().values(), + f, + &format!("{path}.element"), + ), + (ArrowType::FixedSizeList(..), ArrowType::FixedSizeList(f, _)) => check_non_null( + array.as_fixed_size_list().values(), + f, + &format!("{path}.element"), + ), + (ArrowType::Map(..), ArrowType::Map(f, _)) => { + let entries: ArrayRef = Arc::new(array.as_map().entries().clone()); + check_non_null(&entries, f, path) + } + _ => Ok(()), + } +} + +/// [`check_non_null`] for the fields under `field`, whose own nullability the caller checked. +fn check_non_null_below(array: &ArrayRef, field: &Field, path: &str) -> Result<(), KrishivError> { + let lenient = Field::new(field.name(), field.data_type().clone(), true); + check_non_null(array, &lenient, path) +} + +/// `array` under the nested field names and nullability of `to` (same shape; the buffers are +/// shared). +fn relabel(array: &ArrayRef, to: &ArrowType) -> Result { + if array.data_type() == to { + return Ok(Arc::clone(array)); + } + let wrong = || { + internal(format!( + "a local relation column of {} for {to}", + array.data_type() + )) + }; + let relabelled: ArrayRef = match to { + ArrowType::Struct(fields) => { + let s = array.as_struct_opt().ok_or_else(wrong)?; + let children = s + .columns() + .iter() + .zip(fields.iter()) + .map(|(c, f)| relabel(c, f.data_type())) + .collect::, _>>()?; + Arc::new( + StructArray::try_new(fields.clone(), children, s.nulls().cloned()) + .map_err(|e| internal(e.to_string()))?, + ) + } + ArrowType::List(f) => { + let l = array.as_list_opt::().ok_or_else(wrong)?; + Arc::new( + ListArray::try_new( + Arc::clone(f), + l.offsets().clone(), + relabel(l.values(), f.data_type())?, + l.nulls().cloned(), + ) + .map_err(|e| internal(e.to_string()))?, + ) + } + ArrowType::LargeList(f) => { + let l = array.as_list_opt::().ok_or_else(wrong)?; + Arc::new( + LargeListArray::try_new( + Arc::clone(f), + l.offsets().clone(), + relabel(l.values(), f.data_type())?, + l.nulls().cloned(), + ) + .map_err(|e| internal(e.to_string()))?, + ) + } + ArrowType::FixedSizeList(f, size) => { + let l = array.as_fixed_size_list_opt().ok_or_else(wrong)?; + Arc::new( + FixedSizeListArray::try_new( + Arc::clone(f), + *size, + relabel(l.values(), f.data_type())?, + l.nulls().cloned(), + ) + .map_err(|e| internal(e.to_string()))?, + ) + } + ArrowType::Map(f, sorted) => { + let m = array.as_map_opt().ok_or_else(wrong)?; + let entries = relabel(&(Arc::new(m.entries().clone()) as ArrayRef), f.data_type())?; + let entries = entries.as_struct_opt().ok_or_else(wrong)?.clone(); + Arc::new( + MapArray::try_new( + Arc::clone(f), + m.offsets().clone(), + entries, + m.nulls().cloned(), + *sorted, + ) + .map_err(|e| internal(e.to_string()))?, + ) + } + _ => return Err(wrong()), + }; + Ok(relabelled) +} + +#[cfg(test)] +mod tests; diff --git a/crates/krishiv-session/src/local_relation/tests.rs b/crates/krishiv-session/src/local_relation/tests.rs new file mode 100644 index 000000000..182341a44 --- /dev/null +++ b/crates/krishiv-session/src/local_relation/tests.rs @@ -0,0 +1,326 @@ +// SPDX-License-Identifier: Apache-2.0 +//! `Session::local_relation` (P3-T012a part 2a): IPC in, a carrier over the session's relation +//! store out; eager casts; the lease keeping the entry alive; a decoded carrier resolved by hash. + +use std::sync::Arc; + +use arrow_array::cast::AsArray as _; +use arrow_array::types::{Int32Type, Int64Type}; +use arrow_array::{Array, ArrayRef, Int32Array, RecordBatch, StringArray, StructArray}; +use arrow_schema::{DataType as A, Field, Fields, Schema, SchemaRef}; +use futures::executor::block_on; +use krishiv_common::error::ErrorClass; +use krishiv_memory::RelationStore; +use krishiv_plan::{LocalRelation, LogicalPlan, Plan, PlanNode, ReadSource}; +use krishiv_types::batch::Batch; +use krishiv_types::{DataType, StructField, StructType}; + +use super::{conform, relabel, same_shape}; +use crate::functions::col; +use crate::session::Session; + +fn ints(name: &str, nullable: bool) -> Field { + Field::new(name, A::Int32, nullable) +} + +fn batch(schema: &SchemaRef, cols: Vec) -> RecordBatch { + RecordBatch::try_new(Arc::clone(schema), cols).unwrap() +} + +/// `batches` as the IPC bytes a client sends. +fn ipc(schema: &SchemaRef, batches: Vec) -> Vec { + let scratch = RelationStore::new(); + let batches: Vec = batches.into_iter().map(Batch::new).collect(); + let lease = scratch.insert(schema, &batches).unwrap(); + scratch.encode_ipc(&lease.hash()).unwrap() +} + +fn two_ints() -> (SchemaRef, Vec) { + let schema = Arc::new(Schema::new(vec![ints("x", false), ints("y", true)])); + let b = |from: i32| { + batch( + &schema, + vec![ + Arc::new(Int32Array::from(vec![from, from + 1])), + Arc::new(Int32Array::from(vec![Some(from), None])), + ], + ) + }; + let bytes = ipc(&schema, vec![b(1), b(10)]); + (schema, bytes) +} + +fn dataframe(s: &Session, l: &LocalRelation) -> crate::dataframe::DataFrame { + let plan = LogicalPlan::new(PlanNode::Read(ReadSource::Local(l.clone()))); + s.from_plan(&Plan::try_new(plan).unwrap()) +} + +fn rows(batches: &[RecordBatch]) -> usize { + batches.iter().map(RecordBatch::num_rows).sum() +} + +fn user(fields: &[(&str, DataType, bool)]) -> StructType { + StructType::new( + fields + .iter() + .map(|(n, t, nullable)| StructField::new(*n, t.clone(), *nullable)), + ) +} + +#[test] +fn data_round_trips_through_collect() { + let s = Session::new().unwrap(); + let (_, bytes) = two_ints(); + let l = block_on(s.local_relation(Some(&bytes), None)).unwrap(); + assert_eq!((l.rows(), l.batches()), (4, 2)); + let out = block_on(dataframe(&s, &l).collect()).unwrap(); + assert_eq!(rows(&out), 4); + let names: Vec<_> = out[0] + .schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .collect(); + assert_eq!(names, ["x", "y"]); + let x: Vec = out + .iter() + .flat_map(|b| b.column(0).as_primitive::().values().to_vec()) + .collect(); + assert_eq!(x, [1, 2, 10, 11]); +} + +#[test] +fn schema_only_data_and_a_schema_alone_give_zero_rows_and_neither_is_an_error() { + let s = Session::new().unwrap(); + let schema = Arc::new(Schema::new(vec![ints("x", true)])); + let empty = ipc(&schema, Vec::new()); + let l = block_on(s.local_relation(Some(&empty), None)).unwrap(); + assert_eq!(l.rows(), 0); + assert_eq!(rows(&block_on(dataframe(&s, &l).collect()).unwrap()), 0); + + let want = user(&[("a", DataType::Long, true)]); + let l = block_on(s.local_relation(None, Some(want))).unwrap(); + assert_eq!((l.rows(), l.batches()), (0, 0)); + assert_eq!(l.schema().field(0).data_type(), &A::Int64); + assert_eq!(rows(&block_on(dataframe(&s, &l).collect()).unwrap()), 0); + + let err = block_on(s.local_relation(None, None)).unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::UnableToInferSchema); +} + +#[test] +fn a_pure_rename_skips_execution_and_stores_one_copy() { + let s = Session::new().unwrap(); + let (_, bytes) = two_ints(); + let plain = block_on(s.local_relation(Some(&bytes), None)).unwrap(); + let one_copy = s.relations().reserved(); + assert_eq!(s.relations().len(), 1); + drop(plain); + assert_eq!(s.relations().len(), 0); + let want = user(&[ + ("a", DataType::Integer, false), + ("b", DataType::Integer, true), + ]); + let l = block_on(s.local_relation(Some(&bytes), Some(want))).unwrap(); + assert_eq!( + s.relations().len(), + 1, + "only the renamed relation is stored" + ); + assert_eq!(s.relations().reserved(), one_copy); + let names: Vec<_> = l + .schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .collect(); + assert_eq!(names, ["a", "b"]); +} + +#[test] +fn relabelling_shares_the_buffers_of_nested_columns() { + let inner = Fields::from(vec![ints("p", true)]); + let outer = Fields::from(vec![ints("q", true)]); + let values: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let have: ArrayRef = Arc::new(StructArray::new(inner, vec![Arc::clone(&values)], None)); + let want = A::Struct(outer); + assert!(same_shape(have.data_type(), &want)); + let out = relabel(&have, &want).unwrap(); + assert_eq!(out.data_type(), &want); + let before = have.as_struct().column(0).to_data().buffers()[0].as_ptr(); + let after = out.as_struct().column(0).to_data().buffers()[0].as_ptr(); + assert_eq!(before, after, "the values buffer is shared"); +} + +#[test] +fn a_null_in_a_non_nullable_field_fails_at_any_depth() { + let s = Session::new().unwrap(); + // Top level: `y` holds a null. + let (_, bytes) = two_ints(); + let want = user(&[ + ("a", DataType::Integer, false), + ("b", DataType::Integer, false), + ]); + let err = block_on(s.local_relation(Some(&bytes), Some(want))).unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::NotNullAssertViolation); + assert_eq!(s.relations().len(), 0, "nothing is left behind"); + + // Nested: a struct child declared non-nullable holds a null under a valid struct. + let child = Fields::from(vec![ints("p", true)]); + let schema = Arc::new(Schema::new(vec![Field::new( + "s", + A::Struct(child.clone()), + true, + )])); + let col: ArrayRef = Arc::new(StructArray::new( + child, + vec![Arc::new(Int32Array::from(vec![Some(1), None])) as ArrayRef], + None, + )); + let bytes = ipc(&schema, vec![batch(&schema, vec![col])]); + let strict = StructType::new([StructField::new("p", DataType::Integer, false)]); + let want = user(&[("s", DataType::Struct(strict), true)]); + let err = block_on(s.local_relation(Some(&bytes), Some(want))).unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::NotNullAssertViolation); + + // The same shape with the child nullable is fine. + let loose = StructType::new([StructField::new("q", DataType::Integer, true)]); + let want = user(&[("t", DataType::Struct(loose), true)]); + assert!(block_on(s.local_relation(Some(&bytes), Some(want))).is_ok()); +} + +#[test] +fn the_cast_path_stores_only_the_result() { + let s = Session::new().unwrap(); + let (_, bytes) = two_ints(); + let want = user(&[("a", DataType::Long, false), ("b", DataType::Long, true)]); + let l = block_on(s.local_relation(Some(&bytes), Some(want))).unwrap(); + assert_eq!(s.relations().len(), 1, "the temporary input was released"); + assert_eq!(l.schema().field(0).data_type(), &A::Int64); + let out = block_on(dataframe(&s, &l).collect()).unwrap(); + let a: Vec = out + .iter() + .flat_map(|b| b.column(0).as_primitive::().values().to_vec()) + .collect(); + assert_eq!(a, [1, 2, 10, 11]); + // One copy of the result: the same batches stored on their own account for the same bytes. + let held = s.relations().reserved(); + let scratch = RelationStore::new(); + let stored: Vec = s + .relations() + .resolve(&l.relation().hash, 0..l.batches()) + .unwrap(); + let _kept = scratch.insert(l.schema(), &stored).unwrap(); + assert_eq!(held, scratch.reserved()); +} + +#[test] +fn a_failing_cast_or_a_wrong_count_errors_from_local_relation() { + let s = Session::new().unwrap(); + let schema = Arc::new(Schema::new(vec![Field::new("w", A::Utf8, true)])); + let col: ArrayRef = Arc::new(StringArray::from(vec!["abc"])); + let bytes = ipc(&schema, vec![batch(&schema, vec![col])]); + let want = user(&[("n", DataType::Integer, false)]); + assert!(block_on(s.local_relation(Some(&bytes), Some(want))).is_err()); + assert_eq!(s.relations().len(), 0, "nothing is left behind"); + + let (_, bytes) = two_ints(); + let want = user(&[("only", DataType::Integer, true)]); + let err = block_on(s.local_relation(Some(&bytes), Some(want))).unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::NumColumnsMismatch); +} + +#[test] +fn identical_data_shares_one_entry_and_the_last_clone_releases_it() { + let s = Session::new().unwrap(); + let (_, bytes) = two_ints(); + let a = block_on(s.local_relation(Some(&bytes), None)).unwrap(); + let b = block_on(s.local_relation(Some(&bytes), None)).unwrap(); + assert_eq!(a, b); + assert_eq!(s.relations().len(), 1); + let held = s.relations().reserved(); + assert!(held > 0); + let clone = a.clone(); + drop((a, b)); + assert_eq!(s.relations().reserved(), held, "a clone still holds it"); + drop(clone); + assert_eq!(s.relations().len(), 0); + assert_eq!(s.relations().reserved(), 0); +} + +#[test] +fn a_scan_deals_its_batches_as_splits() { + let rt = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap(); + let s = Session::new().unwrap(); + let (_, bytes) = two_ints(); + let l = block_on(s.local_relation(Some(&bytes), None)).unwrap(); + let df = dataframe(&s, &l); + let analysis = rt.block_on(df.analysis()).unwrap(); + let (_, physical) = rt + .block_on(s.physical(&analysis.analyzed, &analysis.env)) + .unwrap(); + let specs = analysis.env.catalog.task_specs(&physical, 2, 2).unwrap(); + let mut seen = 0; + for task in &specs { + for splits in task.scans.values() { + if let krishiv_exec::physical::ScanSplits::Batches { relation, range } = splits { + assert_eq!(relation, &l.relation()); + seen += range.len(); + } + } + } + assert_eq!(seen, 2, "both batches are dealt: {specs:?}"); + let mut total = 0; + for task in &specs { + total += rows( + &rt.block_on(s.run_task(&physical, &analysis.env, task)) + .unwrap(), + ); + } + assert_eq!(total, 4); +} + +#[test] +fn the_same_query_returns_the_same_rows_under_plan_round_trip() { + let (_, bytes) = two_ints(); + let run = |round_trip: &str| { + let s = Session::builder() + .config("krishiv.exec.debug.planRoundTrip", round_trip) + .build() + .unwrap(); + let l = block_on(s.local_relation(Some(&bytes), None)).unwrap(); + let df = dataframe(&s, &l) + .filter(col("x").gt(1)) + .order_by([col("x")]); + block_on(df.collect()).unwrap() + }; + let (plain, through) = (run("false"), run("true")); + assert_eq!(rows(&plain), 3); + assert_eq!(plain, through); +} + +#[test] +fn a_decoded_carrier_whose_entry_is_gone_fails_loudly() { + let s = Session::new().unwrap(); + let schema = Arc::new(Schema::new(vec![ints("x", true)])); + let l = LocalRelation::from_decoded(schema, [9; 32], 1, 1, 8); + let err = block_on(dataframe(&s, &l).collect()).unwrap_err(); + assert_eq!(err.error_class(), ErrorClass::InternalError); + assert!(err.to_string().contains(&"09".repeat(32)), "{err}"); +} + +#[test] +fn conform_relabels_names_without_copying() { + let schema = Arc::new(Schema::new(vec![ints("x", true)])); + let have = batch(&schema, vec![Arc::new(Int32Array::from(vec![1, 2]))]); + let target = Arc::new(Schema::new(vec![ints("renamed", true)])); + let out = conform(&have, &target).unwrap(); + assert_eq!(out.schema(), target); + assert_eq!( + have.column(0).to_data().buffers()[0].as_ptr(), + out.column(0).to_data().buffers()[0].as_ptr() + ); +} diff --git a/crates/krishiv-session/src/result_cache.rs b/crates/krishiv-session/src/result_cache.rs index 5a7cad7cd..dac4dec99 100644 --- a/crates/krishiv-session/src/result_cache.rs +++ b/crates/krishiv-session/src/result_cache.rs @@ -403,6 +403,10 @@ fn survey(plan: &Arc) -> Found { if let PlanNode::Scan(s) = &p.node { match &s.source { ReadSource::Table { name, .. } => found.tables.push(name.clone()), + // Content-addressed input: no file or table dependency. + ReadSource::Local(_) => {} + // Files by path, and any source added later, are assumed + // to read files. _ => found.files = true, } } diff --git a/crates/krishiv-session/src/session.rs b/crates/krishiv-session/src/session.rs index caa30362b..33eccf9fa 100644 --- a/crates/krishiv-session/src/session.rs +++ b/crates/krishiv-session/src/session.rs @@ -1264,6 +1264,7 @@ impl Session { u64::try_from(self.state.runtime.executor.parallelism().max(1)).unwrap_or(1) }, runtime: self.state.runtime.executor.io_handle().clone(), + relations: Arc::clone(&self.state.relations), }; let catalog = Arc::new(QueryCatalog::new( Arc::clone(&self.state.catalog), diff --git a/proto/krishiv/plan/v1/plan.proto b/proto/krishiv/plan/v1/plan.proto index ce6b79d10..51d2167f2 100644 --- a/proto/krishiv/plan/v1/plan.proto +++ b/proto/krishiv/plan/v1/plan.proto @@ -987,6 +987,8 @@ message ReadSource { oneof kind { TableSource table = 1; FilesSource files = 2; + // feature: local_relation_source + LocalSource local = 3; } } @@ -1004,6 +1006,20 @@ message FilesSource { DataType schema = 4; } +// Arrow data a plan reads (Spark's LocalRelation), named by content hash. The +// data is never in the plan: it lives in the session's relation store of the +// process that runs the plan, found by `hash`; a process without it fails. +// feature: local_relation_source +message LocalSource { + // SHA-256 content hash (32 bytes). + bytes hash = 1; + // Arrow IPC `Schema` flatbuffer of the stored batches. + bytes schema = 2; + uint64 rows = 3; + uint64 batches = 4; + uint64 bytes = 5; +} + // A data source option: plain text or a reference to a secret. An inline // secret or a URL carrying credentials is not encodable (ADR 0038). message OptionValue { From 6d3a7063aa5cc642fed4cbfb48536b2174a8e7c8 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 23:03:28 +0000 Subject: [PATCH 04/10] changelog: PR fragment Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- changelog.d/.added.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changelog.d/.added.md diff --git a/changelog.d/.added.md b/changelog.d/.added.md new file mode 100644 index 000000000..f80c8a185 --- /dev/null +++ b/changelog.d/.added.md @@ -0,0 +1 @@ +Spark Connect group 1 has the IR it needs: nested literal values, a declared schema on file reads and a data-carrying local relation, each with a plan codec feature (PR-12 part 2a). From 61bf266028a4d09c58b1ea41215ee60b848743ea Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 23:04:07 +0000 Subject: [PATCH 05/10] Revert "changelog: PR fragment" This reverts commit a9a188e57469dafd65bf28dda1d864c5d86c18f9. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- changelog.d/.added.md | 1 - 1 file changed, 1 deletion(-) delete mode 100644 changelog.d/.added.md diff --git a/changelog.d/.added.md b/changelog.d/.added.md deleted file mode 100644 index f80c8a185..000000000 --- a/changelog.d/.added.md +++ /dev/null @@ -1 +0,0 @@ -Spark Connect group 1 has the IR it needs: nested literal values, a declared schema on file reads and a data-carrying local relation, each with a plan codec feature (PR-12 part 2a). From 665243668c0493cdebac97de2e41f2b9b01e57c9 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sat, 14 Nov 2026 23:58:09 +0000 Subject: [PATCH 06/10] plan, exec, optimizer, session: fix the part 2a review findings (P3-T012a) B1: nested literal types (Array, KeyValueMap, Struct) are written to the canonical encoding without loss (nullability, struct fields past the 25th, comments), so literals Literal::eq tells apart no longer share bytes. The tags are new; no existing tag, golden or ENCODING_VERSION changes. N-1: the optimizer and exec literal sites name the nested variants. N-2: a struct<> literal is a one-row array. N-3/N-4: recorded as golden-pending in the code. N-5: the property generator makes nested literals and declared file schemas. N-6: note that ReadSource::Files is now #[non_exhaustive] (a breaking change under unstable-ir): build it with ReadSource::files / with_schema. N-7: LocalRelation::from_decoded documents its use. Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-exec/src/planner/eval/mod.rs | 2 +- crates/krishiv-exec/src/planner/pushdown.rs | 5 +- crates/krishiv-optimizer/src/fold.rs | 14 +++-- crates/krishiv-plan/src/encode.rs | 46 ++++++++++++++- crates/krishiv-plan/src/literal.rs | 5 ++ crates/krishiv-plan/src/local.rs | 5 ++ .../krishiv-plan/src/tests/nested_literals.rs | 47 +++++++++++++++ crates/krishiv-plan/src/tests/props.rs | 57 +++++++++++++++++-- crates/krishiv-session/src/local_relation.rs | 4 ++ 9 files changed, 171 insertions(+), 14 deletions(-) diff --git a/crates/krishiv-exec/src/planner/eval/mod.rs b/crates/krishiv-exec/src/planner/eval/mod.rs index 19c0dd870..fcadfa283 100644 --- a/crates/krishiv-exec/src/planner/eval/mod.rs +++ b/crates/krishiv-exec/src/planner/eval/mod.rs @@ -1150,7 +1150,7 @@ pub(crate) fn literal_array(l: &Literal) -> Result { .map(literal_array) .collect::>>()?; Arc::new( - StructArray::try_new(arrow_fields, columns, None) + StructArray::try_new_with_length(arrow_fields, columns, None, 1) .map_err(|e| internal(&e.to_string()))?, ) } diff --git a/crates/krishiv-exec/src/planner/pushdown.rs b/crates/krishiv-exec/src/planner/pushdown.rs index 8adca3a3c..ac00c7237 100644 --- a/crates/krishiv-exec/src/planner/pushdown.rs +++ b/crates/krishiv-exec/src/planner/pushdown.rs @@ -244,7 +244,10 @@ fn value_or_null(e: &Expr) -> Option> { Literal::Date(v) => ScalarValue::Date32(*v), Literal::Timestamp(v) | Literal::TimestampNtz(v) => ScalarValue::TimestampMicros(*v), // No scalar form for the file filters, so never pushed down: nested - // values (`Literal::{Array, KeyValueMap, Struct}`), maps, intervals. + // values, maps, intervals. + Literal::Array { .. } | Literal::KeyValueMap { .. } | Literal::Struct { .. } => { + return None; + } _ => return None, })) } diff --git a/crates/krishiv-optimizer/src/fold.rs b/crates/krishiv-optimizer/src/fold.rs index 1dbbe6c14..1979ee930 100644 --- a/crates/krishiv-optimizer/src/fold.rs +++ b/crates/krishiv-optimizer/src/fold.rs @@ -245,8 +245,10 @@ pub(crate) fn compare_values(x: &Literal, y: &Literal) -> Option { |v: i128, from: u8| v.checked_mul(10_i128.checked_pow(u32::from(scale - from))?); up(*a, s.scale())?.cmp(&up(*b, t.scale())?) } - // Not foldable: nested values (`Literal::{Array, KeyValueMap, - // Struct}`) are not ordered here, nor are mixed types. + // Not foldable: nested values are not ordered here, nor are mixed types. + (Literal::Array { .. } | Literal::KeyValueMap { .. } | Literal::Struct { .. }, _) => { + return None; + } _ => return None, }) } @@ -406,9 +408,11 @@ fn to_array(l: &Literal) -> Option { Literal::String(v) => Arc::new(StringArray::from(vec![v.as_str()])), Literal::Date(v) => Arc::new(Date32Array::from(vec![*v])), Literal::TimestampNtz(v) => Arc::new(TimestampMicrosecondArray::from(vec![*v])), - // Not foldable: `Literal::{Array, KeyValueMap, Struct}` have no - // one-row form here, so a cast or call over a nested value stays as - // written. + // Not foldable: nested values have no one-row form here, so a cast or + // call over one stays as written. + Literal::Array { .. } | Literal::KeyValueMap { .. } | Literal::Struct { .. } => { + return None; + } _ => return None, }) } diff --git a/crates/krishiv-plan/src/encode.rs b/crates/krishiv-plan/src/encode.rs index a0d48e07b..e8929af4f 100644 --- a/crates/krishiv-plan/src/encode.rs +++ b/crates/krishiv-plan/src/encode.rs @@ -718,6 +718,46 @@ impl<'k, 'a> Encoder<'k, 'a> { self.str(&t.simple_string()); } + /// A nested literal's type without loss: `simple_string` drops + /// nullability, comments and fields past the 25th, which `Literal::eq` + /// keeps. Bounded by `DataType::MAX_NESTING_DEPTH`. + fn lossless_type(&mut self, t: &krishiv_types::DataType) { + use krishiv_types::DataType; + match t { + DataType::Array(a) => { + self.u8(1); + self.bool(a.contains_null()); + self.lossless_type(a.element_type()); + } + DataType::Map(m) => { + self.u8(2); + self.bool(m.value_contains_null()); + self.lossless_type(m.key_type()); + self.lossless_type(m.value_type()); + } + DataType::Struct(st) => { + self.u8(3); + self.len(st.fields().len()); + for f in st.fields() { + self.str(f.name()); + self.bool(f.nullable()); + // The map is ordered by key, so its debug text is stable. + let md = f.metadata(); + self.str(&if md.is_empty() { + String::new() + } else { + format!("{:?}", md.as_map()) + }); + self.lossless_type(f.data_type()); + } + } + leaf => { + self.u8(0); + self.data_type(leaf); + } + } + } + fn exprs(&mut self, v: &'a [Arc]) { self.len(v.len()); v.iter().for_each(|e| self.expr(e)); @@ -1580,7 +1620,7 @@ impl<'k, 'a> Encoder<'k, 'a> { elements, } => { self.u8(20); - self.data_type(&krishiv_types::DataType::Array(Arc::clone(data_type))); + self.lossless_type(&krishiv_types::DataType::Array(Arc::clone(data_type))); self.len(elements.len()); for e in elements { self.literal(e); @@ -1588,7 +1628,7 @@ impl<'k, 'a> Encoder<'k, 'a> { } Literal::KeyValueMap { data_type, entries } => { self.u8(21); - self.data_type(&krishiv_types::DataType::Map(Arc::clone(data_type))); + self.lossless_type(&krishiv_types::DataType::Map(Arc::clone(data_type))); self.len(entries.len()); for (k, v) in entries { self.literal(k); @@ -1597,7 +1637,7 @@ impl<'k, 'a> Encoder<'k, 'a> { } Literal::Struct { data_type, fields } => { self.u8(22); - self.data_type(&krishiv_types::DataType::Struct((**data_type).clone())); + self.lossless_type(&krishiv_types::DataType::Struct((**data_type).clone())); self.len(fields.len()); for f in fields { self.literal(f); diff --git a/crates/krishiv-plan/src/literal.rs b/crates/krishiv-plan/src/literal.rs index 46398688b..3659033f2 100644 --- a/crates/krishiv-plan/src/literal.rs +++ b/crates/krishiv-plan/src/literal.rs @@ -290,6 +290,11 @@ fn check_member( impl Literal { /// The type of this literal; `None` for an untyped `NULL`. + /// + /// Golden-pending (P3-T012a): a typed null of `VOID` is also `None`, so + /// an `ARRAY`/`MAP<_, VOID>` literal holding a null member is + /// refused as an untyped null; whether Spark's `typedLit` reaches that + /// is settled with the literal goldens in part 2b. #[must_use] pub fn data_type(&self) -> Option { Some(match self { diff --git a/crates/krishiv-plan/src/local.rs b/crates/krishiv-plan/src/local.rs index add449705..1f52720d2 100644 --- a/crates/krishiv-plan/src/local.rs +++ b/crates/krishiv-plan/src/local.rs @@ -41,6 +41,11 @@ impl LocalRelation { /// A relation named by its 32-byte content hash, holding no lease (the /// plan codec's decode: the entry is found by hash where the plan runs). + /// + /// For decoding and for internal plans that must not keep the entry + /// alive; a session builds one from a lease with [`Self::from_lease`]. + /// The hash is the caller's claim: a wrong one fails loudly at the + /// relation provider, which finds no entry (or none with that content). #[must_use] pub fn from_decoded( schema: SchemaRef, diff --git a/crates/krishiv-plan/src/tests/nested_literals.rs b/crates/krishiv-plan/src/tests/nested_literals.rs index 64840e46d..5314a5c39 100644 --- a/crates/krishiv-plan/src/tests/nested_literals.rs +++ b/crates/krishiv-plan/src/tests/nested_literals.rs @@ -187,3 +187,50 @@ fn nested_literals_print_as_constructor_calls() { .unwrap(); assert_eq!(map.to_string(), "MAP('a', 1)"); } + +/// Every pair `Literal::eq` tells apart must encode differently: the +/// canonical type keeps nullability and every struct field. +#[test] +fn the_hash_keeps_nullability_and_every_field_of_the_type() { + let arr = |contains_null: bool| { + Literal::array( + ArrayType::new(DataType::Integer, contains_null), + vec![Literal::Integer(1)], + ) + .unwrap() + }; + assert_ne!(arr(true), arr(false)); + assert_ne!(hash_of(arr(true)), hash_of(arr(false))); + + let map = |value_null: bool| { + Literal::key_value_map( + MapType::new(DataType::String, DataType::Integer, value_null), + vec![(Literal::String("a".to_owned()), Literal::Integer(1))], + ) + .unwrap() + }; + assert_ne!(map(true), map(false)); + assert_ne!(hash_of(map(true)), hash_of(map(false))); + + let st = |nullable: bool| { + Literal::struct_value( + StructType::new([StructField::new("a", DataType::Integer, nullable)]), + vec![Literal::Integer(1)], + ) + .unwrap() + }; + assert_ne!(st(true), st(false)); + assert_ne!(hash_of(st(true)), hash_of(st(false))); + + // Names that differ only after the 25th field (`simple_string` stops there). + let wide = |last: &str| { + let fields: Vec<_> = (0..25) + .map(|i| StructField::new(format!("f{i}"), DataType::Integer, true)) + .chain([StructField::new(last, DataType::Integer, true)]) + .collect(); + Literal::struct_value(StructType::new(fields), vec![Literal::Integer(1); 26]).unwrap() + }; + assert_ne!(wide("x"), wide("y")); + assert_ne!(hash_of(wide("x")), hash_of(wide("y"))); + assert_eq!(hash_of(wide("x")), hash_of(wide("x"))); +} diff --git a/crates/krishiv-plan/src/tests/props.rs b/crates/krishiv-plan/src/tests/props.rs index 26e79f19c..74405bb9b 100644 --- a/crates/krishiv-plan/src/tests/props.rs +++ b/crates/krishiv-plan/src/tests/props.rs @@ -113,7 +113,7 @@ impl Gen<'_> { DayTimeField::Minute, DayTimeField::Second, ]; - match self.pick(16) { + match self.pick(19) { 0 => Literal::Null(self.flag().then(|| self.data_type())), 1 => Literal::Boolean(v.is_multiple_of(2)), 2 => Literal::Byte(i8::from_le_bytes([v])), @@ -139,6 +139,38 @@ impl Gen<'_> { end: start.max(ym[usize::from(self.pick(2))]), } } + 16 => { + let contains_null = self.flag(); + Literal::array( + krishiv_types::ArrayType::new(DataType::Integer, contains_null), + vec![Literal::Integer(i32::from(v)); usize::from(v % 3)], + ) + .unwrap() + } + 17 => { + let value_null = self.flag(); + Literal::key_value_map( + krishiv_types::MapType::new(DataType::String, DataType::Integer, value_null), + vec![( + Literal::String("k".to_owned()), + Literal::Integer(i32::from(v)), + )], + ) + .unwrap() + } + 18 => { + let nullable = self.flag(); + let name = self.name(); + Literal::struct_value( + krishiv_types::StructType::new([krishiv_types::StructField::new( + name, + DataType::Integer, + nullable, + )]), + vec![Literal::Integer(i32::from(v))], + ) + .unwrap() + } _ => { let start = dt[usize::from(self.pick(4))]; Literal::DayTimeInterval { @@ -378,7 +410,13 @@ impl Gen<'_> { format: self.name(), paths: self.many(1, 2, Gen::option), options, - schema: None, + schema: self.flag().then(|| { + krishiv_types::StructType::new([krishiv_types::StructField::new( + "c", + DataType::Integer, + self.flag(), + )]) + }), } } } @@ -1106,8 +1144,19 @@ fn m_source(s: &mut ReadSource, w: &mut Walk) -> bool { format, paths, options, - .. - } => m_string(format, w) || m_values(paths, w) || m_options(options, w), + schema, + } => { + let flipped = schema.as_ref().map_or_else( + || { + Some(krishiv_types::StructType::new([ + krishiv_types::StructField::new("c", DataType::Integer, true), + ])) + }, + |_| None, + ); + site!(w, *schema = flipped); + m_string(format, w) || m_values(paths, w) || m_options(options, w) + } // Nothing inside to mutate: the hash and counts are the whole value. ReadSource::Local(_) => false, } diff --git a/crates/krishiv-session/src/local_relation.rs b/crates/krishiv-session/src/local_relation.rs index 7b0ac57d0..3b9da0bfc 100644 --- a/crates/krishiv-session/src/local_relation.rs +++ b/crates/krishiv-session/src/local_relation.rs @@ -245,6 +245,10 @@ fn null_in(path: &str) -> KrishivError { /// Fails when `array` holds a null where `field` (and the fields under it) is non-nullable. A struct /// child's nulls count only where the struct itself is valid; a list's elements are checked over /// the whole values array (a null parent owns an empty range in the arrays Arrow writes). +/// +/// Golden-pending (P3-T012a): a writer that leaves data under a null list slot, or a null grandchild +/// under a null struct two levels up, can raise a violation Spark would not; carrying the parent +/// validity mask down is settled with the nullability goldens. fn check_non_null(array: &ArrayRef, field: &Field, path: &str) -> Result<(), KrishivError> { if !field.is_nullable() && array.logical_null_count() > 0 { return Err(null_in(path)); From 5ed9090cbe9422352ab104857d58aff8d9be77cd Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sun, 15 Nov 2026 00:08:16 +0000 Subject: [PATCH 07/10] changelog: fragment for #772 Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- changelog.d/772.added.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changelog.d/772.added.md diff --git a/changelog.d/772.added.md b/changelog.d/772.added.md new file mode 100644 index 000000000..3169e601b --- /dev/null +++ b/changelog.d/772.added.md @@ -0,0 +1 @@ +Spark Connect group 1 has the IR it needs: nested literal values, a declared schema on file reads and a data-carrying local relation, each with a plan codec feature (PR-12 part 2a, #772). From c15e33a8fe4bf6d76864b68979a5491414697657 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sun, 15 Nov 2026 01:07:06 +0000 Subject: [PATCH 08/10] plan, session: fix three review findings on local relations (P3-T012a) A declared file schema is encoded losslessly (tag 4), so reads differing only in nested nullability no longer share canonical bytes. An empty IPC stream validates the schema conversion like a zero-row batch. Utf8 and Binary columns move to the view carriers by an Arrow cast before the logical cast. Signed-off-by: gopal tathe Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- Cargo.lock | 1 + crates/krishiv-plan/src/encode.rs | 2 +- crates/krishiv-plan/src/tests/files.rs | 34 ++++++- crates/krishiv-session/Cargo.toml | 1 + crates/krishiv-session/src/local_relation.rs | 96 ++++++++++++++++++- .../src/local_relation/tests.rs | 52 ++++++++++ 6 files changed, 183 insertions(+), 3 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 33b63805b..0b9abbd87 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2760,6 +2760,7 @@ version = "0.2.0" dependencies = [ "arrow-array", "arrow-buffer", + "arrow-cast", "arrow-schema", "arrow-select", "aws-lc-rs", diff --git a/crates/krishiv-plan/src/encode.rs b/crates/krishiv-plan/src/encode.rs index e8929af4f..cc7d0df85 100644 --- a/crates/krishiv-plan/src/encode.rs +++ b/crates/krishiv-plan/src/encode.rs @@ -1182,7 +1182,7 @@ impl<'k, 'a> Encoder<'k, 'a> { paths.iter().for_each(|p| self.option_value(p)); self.options(options); if let Some(schema) = schema { - self.str(&krishiv_types::DataType::Struct(schema.clone()).sql()); + self.lossless_type(&krishiv_types::DataType::Struct(schema.clone())); } } } diff --git a/crates/krishiv-plan/src/tests/files.rs b/crates/krishiv-plan/src/tests/files.rs index 80698c074..f27d61e5f 100644 --- a/crates/krishiv-plan/src/tests/files.rs +++ b/crates/krishiv-plan/src/tests/files.rs @@ -673,5 +673,37 @@ fn a_declared_schema_prints_and_encodes_apart() { ] ); assert!(b.len() > a.len()); - assert!(b.windows(3).any(|w| w == b"BIG")); + assert!(b.windows(3).any(|w| w == b"big")); +} + +/// The declared schema is encoded without loss: reads that differ only in an +/// array's `containsNull`, a map's `valueContainsNull` or a nested field's +/// nullability do not share canonical bytes. +#[test] +fn a_declared_schema_encodes_nested_nullability() { + use krishiv_types::{ArrayType, MapType, StructField, StructType}; + let read = |ty: DataType| { + let schema = StructType::new([StructField::new("c", ty, true)]); + let p = LogicalPlan::new(PlanNode::Read( + ReadSource::files( + "parquet", + vec![OptionValue::from_user("path", "/d")], + BTreeMap::new(), + ) + .with_schema(schema), + )); + Plan::try_new(p).unwrap().canonical_encoding() + }; + let array = |n| ArrayType::new(DataType::Integer, n).into(); + assert_ne!(read(array(true)), read(array(false))); + let map = |n| MapType::new(DataType::Integer, DataType::Integer, n).into(); + assert_ne!(read(map(true)), read(map(false))); + let st = |n| { + DataType::Struct(StructType::new([StructField::new( + "x", + DataType::Integer, + n, + )])) + }; + assert_ne!(read(st(true)), read(st(false))); } diff --git a/crates/krishiv-session/Cargo.toml b/crates/krishiv-session/Cargo.toml index 4cb007c28..196da5216 100644 --- a/crates/krishiv-session/Cargo.toml +++ b/crates/krishiv-session/Cargo.toml @@ -11,6 +11,7 @@ publish = true [dependencies] arrow-array.workspace = true arrow-buffer.workspace = true +arrow-cast.workspace = true arrow-schema.workspace = true thiserror = { workspace = true, optional = true } # `executor`: file footers are read on a thread of their own (catalog.rs). diff --git a/crates/krishiv-session/src/local_relation.rs b/crates/krishiv-session/src/local_relation.rs index 3b9da0bfc..1ad6b226b 100644 --- a/crates/krishiv-session/src/local_relation.rs +++ b/crates/krishiv-session/src/local_relation.rs @@ -78,14 +78,23 @@ impl Session { } .into()); } + let (schema, records) = carried(&schema, records, &target)?; let renamed = renamed(&schema, &target); let same = renamed .fields() .iter() .zip(target.fields()) .all(|(a, b)| same_shape(a.data_type(), b.data_type())); - let converted = if same || records.is_empty() { + let converted = if same { records + } else if records.is_empty() { + // No batches: the conversion is still checked, as for a zero-row batch. + let empty = RecordBatch::new_empty(Arc::clone(&renamed)); + self.cast_batches(&renamed, &target, vec![empty]) + .await? + .into_iter() + .filter(|b| b.num_rows() > 0) + .collect() } else { self.cast_batches(&renamed, &target, records).await? }; @@ -150,6 +159,91 @@ impl Session { } } +/// The view carrier of `t` for encoding-only differences: `Utf8`/`LargeUtf8` as `Utf8View` and +/// `Binary`/`LargeBinary` as `BinaryView`, at any depth. +fn viewed(t: &ArrowType) -> ArrowType { + let field = |f: &Arc| Arc::new(f.as_ref().clone().with_data_type(viewed(f.data_type()))); + match t { + ArrowType::Utf8 | ArrowType::LargeUtf8 => ArrowType::Utf8View, + ArrowType::Binary | ArrowType::LargeBinary => ArrowType::BinaryView, + ArrowType::List(f) => ArrowType::List(field(f)), + ArrowType::LargeList(f) => ArrowType::LargeList(field(f)), + ArrowType::FixedSizeList(f, n) => ArrowType::FixedSizeList(field(f), *n), + ArrowType::Map(f, sorted) => ArrowType::Map(field(f), *sorted), + ArrowType::Struct(fields) => ArrowType::Struct( + fields + .iter() + .map(|f| f.as_ref().clone().with_data_type(viewed(f.data_type()))) + .collect(), + ), + other => other.clone(), + } +} + +/// `records` (of `schema`) with the encoding-only carriers of string and binary columns moved to +/// the views the engine runs on, by an Arrow cast of the buffers: a logical STRING-to-STRING cast +/// would be folded away by the optimizer, so this is done here and explicitly. Only columns whose +/// target is the same logical kind are touched; the rest keep their cast path. +fn carried( + schema: &SchemaRef, + records: Vec, + target: &SchemaRef, +) -> Result<(SchemaRef, Vec), KrishivError> { + let wanted: Vec = schema + .fields() + .iter() + .zip(target.fields()) + .map(|(have, want)| { + let v = viewed(have.data_type()); + if &v != have.data_type() && v == viewed(want.data_type()) { + v + } else { + have.data_type().clone() + } + }) + .collect(); + if wanted + .iter() + .zip(schema.fields()) + .all(|(w, f)| w == f.data_type()) + { + return Ok((Arc::clone(schema), records)); + } + let out = Arc::new(Schema::new_with_metadata( + schema + .fields() + .iter() + .zip(&wanted) + .map(|(f, w)| f.as_ref().clone().with_data_type(w.clone())) + .collect::>(), + schema.metadata().clone(), + )); + let batches = records + .iter() + .map(|b| { + let columns = b + .columns() + .iter() + .zip(&wanted) + .map(|(c, w)| { + if c.data_type() == w { + Ok(Arc::clone(c)) + } else { + arrow_cast::cast(c, w).map_err(|e| internal(e.to_string())) + } + }) + .collect::, _>>()?; + RecordBatch::try_new_with_options( + Arc::clone(&out), + columns, + &RecordBatchOptions::new().with_row_count(Some(b.num_rows())), + ) + .map_err(|e| internal(e.to_string())) + }) + .collect::, _>>()?; + Ok((out, batches)) +} + /// The Arrow schema of a user-declared struct type. fn arrow_schema_of(schema: &StructType) -> Schema { Schema::new( diff --git a/crates/krishiv-session/src/local_relation/tests.rs b/crates/krishiv-session/src/local_relation/tests.rs index 182341a44..95bb5a4a9 100644 --- a/crates/krishiv-session/src/local_relation/tests.rs +++ b/crates/krishiv-session/src/local_relation/tests.rs @@ -324,3 +324,55 @@ fn conform_relabels_names_without_copying() { out.column(0).to_data().buffers()[0].as_ptr() ); } + +#[test] +fn an_empty_stream_validates_the_conversion_like_a_zero_row_batch() { + let s = Session::new().unwrap(); + let schema = Arc::new(Schema::new(vec![ints("x", true)])); + let want = || { + user(&[( + "x", + krishiv_types::ArrayType::new(DataType::Integer, true).into(), + true, + )]) + }; + let empty = ipc(&schema, Vec::new()); + let zero_col: ArrayRef = Arc::new(Int32Array::from(Vec::::new())); + let zero = ipc(&schema, vec![batch(&schema, vec![zero_col])]); + let a = block_on(s.local_relation(Some(&zero), Some(want()))).unwrap_err(); + let b = block_on(s.local_relation(Some(&empty), Some(want()))).unwrap_err(); + assert_eq!(a.error_class(), b.error_class()); +} + +#[test] +fn utf8_and_binary_columns_meet_a_string_and_binary_schema() { + use arrow_array::BinaryArray; + let s = Session::new().unwrap(); + let schema = Arc::new(Schema::new(vec![ + Field::new("s", A::Utf8, true), + Field::new("b", A::Binary, true), + ])); + let cols: Vec = vec![ + Arc::new(StringArray::from(vec![ + Some("a"), + None, + Some("a longer string value"), + ])), + Arc::new(BinaryArray::from_iter(vec![ + Some(b"x".as_slice()), + None, + Some(b"yz".as_slice()), + ])), + ]; + let bytes = ipc(&schema, vec![batch(&schema, cols)]); + let want = user(&[("s", DataType::String, true), ("b", DataType::Binary, true)]); + let l = block_on(s.local_relation(Some(&bytes), Some(want))).unwrap(); + let out = block_on(dataframe(&s, &l).collect()).unwrap(); + assert_eq!(rows(&out), 3); + let strs = out[0].column(0).as_string_view(); + assert_eq!(strs.value(0), "a"); + assert!(strs.is_null(1)); + assert_eq!(strs.value(2), "a longer string value"); + let bins = out[0].column(1).as_binary_view(); + assert_eq!(bins.value(2), b"yz"); +} From 61016e169e3bd3914fa6ab56720196b9f781c1a8 Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sun, 15 Nov 2026 06:57:57 +0000 Subject: [PATCH 09/10] plan: tests for the Debug and data type of nested literals (diff-coverage of the T012a IR) Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- .../krishiv-plan/src/tests/nested_literals.rs | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/crates/krishiv-plan/src/tests/nested_literals.rs b/crates/krishiv-plan/src/tests/nested_literals.rs index 5314a5c39..ed04bba99 100644 --- a/crates/krishiv-plan/src/tests/nested_literals.rs +++ b/crates/krishiv-plan/src/tests/nested_literals.rs @@ -234,3 +234,33 @@ fn the_hash_keeps_nullability_and_every_field_of_the_type() { assert_ne!(hash_of(wide("x")), hash_of(wide("y"))); assert_eq!(hash_of(wide("x")), hash_of(wide("x"))); } + +#[test] +fn debug_and_data_type_of_nested_literals() { + let arr = ints(&[1]); + let map = Literal::key_value_map( + MapType::new(DataType::String, DataType::Integer, true), + vec![(Literal::String("a".to_owned()), Literal::Integer(1))], + ) + .unwrap(); + let st_type = StructType::new([StructField::new("a", DataType::Integer, false)]); + let st = Literal::struct_value(st_type.clone(), vec![Literal::Integer(1)]).unwrap(); + assert!(format!("{arr:?}").starts_with("Array")); + assert!(format!("{map:?}").starts_with("KeyValueMap")); + assert!(format!("{st:?}").starts_with("Struct")); + assert_eq!( + arr.data_type(), + Some(DataType::Array( + ArrayType::new(DataType::Integer, true).into() + )) + ); + assert!(matches!(map.data_type(), Some(DataType::Map(_)))); + assert_eq!(st.data_type(), Some(DataType::Struct(st_type))); + let interval = Literal::CalendarInterval { + months: 1, + days: 2, + micros: 3, + }; + assert!(format!("{interval:?}").starts_with("CalendarInterval")); + assert_eq!(interval.data_type(), Some(DataType::CalendarInterval)); +} From 0855dc4507319f0f539adaa1737ac073a1149e2f Mon Sep 17 00:00:00 2001 From: gopal tathe <44522021+techgopal@users.noreply.github.com> Date: Sun, 15 Nov 2026 08:37:19 +0000 Subject: [PATCH 10/10] plan: refuse an options-only map as a nested member, name the element type of an empty array in identity text, stronger nested literal and local relation tests (review of #772) Signed-off-by: gopal tathe <44522021+techgopal@users.noreply.github.com> --- crates/krishiv-plan/src/display.rs | 13 ++++++- crates/krishiv-plan/src/literal.rs | 7 ++++ crates/krishiv-plan/src/tests/local.rs | 11 ++++++ .../krishiv-plan/src/tests/nested_literals.rs | 37 ++++++++++++++++++- 4 files changed, 65 insertions(+), 3 deletions(-) diff --git a/crates/krishiv-plan/src/display.rs b/crates/krishiv-plan/src/display.rs index 4d7435c24..565ce22ad 100644 --- a/crates/krishiv-plan/src/display.rs +++ b/crates/krishiv-plan/src/display.rs @@ -1314,8 +1314,17 @@ impl Display for Literal { ), // Nested values print as the constructor call that builds them; // Spark's own text is golden-pending (design P3-T012a note, 5). - Self::Array { elements, .. } => { - f.write_str("ARRAY(")?; + Self::Array { + data_type, + elements, + } => { + // Two empty arrays of other element types are different values: identity text + // names the type (the Spark-facing text stays `ARRAY()`). + if elements.is_empty() && IDENTITY.with(std::cell::Cell::get) { + write!(f, "ARRAY<{}>(", data_type.element_type().sql())?; + } else { + f.write_str("ARRAY(")?; + } for (i, e) in elements.iter().enumerate() { if i > 0 { f.write_str(", ")?; diff --git a/crates/krishiv-plan/src/literal.rs b/crates/krishiv-plan/src/literal.rs index 3659033f2..a5bee381c 100644 --- a/crates/krishiv-plan/src/literal.rs +++ b/crates/krishiv-plan/src/literal.rs @@ -275,6 +275,13 @@ fn check_member( "a null {what} in a nested literal whose type forbids nulls" ))); } + // The options-only map literal is not a value: the codec refuses it as a member and the + // executor builds it with other child types than a declared `MAP`. + if matches!(value, Literal::Map(_)) { + return Err(invalid(format!( + "an options-only map as a nested literal {what}; use a key/value map literal" + ))); + } match value.data_type() { Some(t) if t == *expected => Ok(()), Some(t) => Err(invalid(format!( diff --git a/crates/krishiv-plan/src/tests/local.rs b/crates/krishiv-plan/src/tests/local.rs index 1de438959..2d80d9e60 100644 --- a/crates/krishiv-plan/src/tests/local.rs +++ b/crates/krishiv-plan/src/tests/local.rs @@ -187,3 +187,14 @@ fn display_names_the_columns_and_rows() { // Spark-format text is golden-pending; the row count is Krishiv's own. insta::assert_snapshot!(analyzed.plan.root().to_string(), @"LocalRelation [a#0, b#1] (2 rows)"); } + +#[test] +fn different_content_of_the_same_shape_is_not_equal() { + let store = RelationStore::new(); + let (a, b) = (relation(&store, &[1, 2]), relation(&store, &[1, 3])); + assert_ne!(a, b); + let (pa, pb) = (plan(read(&a)), plan(read(&b))); + assert!(!pa.structurally_eq(&pb)); + assert_ne!(pa.canonical_encoding(), pb.canonical_encoding()); + assert_ne!(normalized_plan_hash(&pa), normalized_plan_hash(&pb)); +} diff --git a/crates/krishiv-plan/src/tests/nested_literals.rs b/crates/krishiv-plan/src/tests/nested_literals.rs index ed04bba99..ad1257041 100644 --- a/crates/krishiv-plan/src/tests/nested_literals.rs +++ b/crates/krishiv-plan/src/tests/nested_literals.rs @@ -85,7 +85,18 @@ fn nulls_inside_carry_their_type() { ArrayType::new(DataType::Integer, true), vec![Literal::Integer(1), Literal::Null(Some(DataType::Integer))], ); - assert!(ok.is_ok()); + let Literal::Array { + elements, + data_type, + } = ok.unwrap() + else { + panic!("not an array"); + }; + assert_eq!( + elements, + vec![Literal::Integer(1), Literal::Null(Some(DataType::Integer))] + ); + assert_eq!(*data_type, ArrayType::new(DataType::Integer, true)); // An untyped null, a null of another type, and a null where the type // forbids one are refused. for (contains_null, element) in [ @@ -264,3 +275,27 @@ fn debug_and_data_type_of_nested_literals() { assert!(format!("{interval:?}").starts_with("CalendarInterval")); assert_eq!(interval.data_type(), Some(DataType::CalendarInterval)); } + +#[test] +fn an_options_only_map_is_not_a_nested_member() { + let e = Literal::array( + ArrayType::new( + DataType::Map(MapType::new(DataType::String, DataType::String, false).into()), + false, + ), + vec![Literal::Map(vec![("k".to_owned(), "v".to_owned())])], + ) + .unwrap_err(); + assert_eq!(class(&e), "INTERNAL_ERROR"); +} + +#[test] +fn identity_text_of_an_empty_array_names_its_element_type() { + let empty = |t: DataType| Expr::from(Literal::array(ArrayType::new(t, true), vec![]).unwrap()); + let (a, b) = (empty(DataType::Integer), empty(DataType::Long)); + assert_eq!(a.to_string(), b.to_string()); + assert_ne!( + crate::display::identity_string(&a), + crate::display::identity_string(&b) + ); +}