types: array, map and struct types (P2-T001b, part 1)

crates/krishiv-connect/src/translate/expression/tests.rs

@@ -298,7 +298,7 @@ fn cast_errors() {
         err(&cast(None, 0)).error_class(),
         ErrorClass::ConnectInvalidPlanInvalidOneOfFieldNotSet
     );
-    let e = err(&cast(Some(CastToType::TypeStr("array<int>".into())), 0));
+    let e = err(&cast(Some(CastToType::TypeStr("array<char(3)>".into())), 0));
     assert_eq!(e.error_class(), ErrorClass::KrishivUnsupportedFeature);
     let e = err(&cast(Some(CastToType::TypeStr("notatype".into())), 0));
     assert_eq!(e.error_class(), ErrorClass::UnsupportedDatatype);

crates/krishiv-connect/src/translate/types.rs

@@ -2,7 +2,7 @@
 //! Connect `DataType` to Krishiv's.
 
 use krishiv_common::error::KrishivError;
-use krishiv_types::{DataType, DecimalType};
+use krishiv_types::{ArrayType, DataType, DecimalType, MapType, StructField, StructType};
 
 use super::errors::{assertion, not_set, unsupported_field, unsupported_variant};
 use crate::proto::v4_2::{DataType as Proto, data_type::Kind};
@@ -12,14 +12,28 @@ const DEFAULT_COLLATION: &str = "UTF8_BINARY";
 
 /// Maps a Connect data type, following Spark's `DataTypeProtoConverter`:
 /// a decimal with no precision is `decimal(10,0)`, one with a precision but
-/// no scale has scale 0, and `string` with the default collation is `STRING`.
+/// no scale has scale 0, `string` with the default collation is `STRING`,
+/// and a struct field's metadata is Spark's `Metadata` JSON. An unset
+/// element, key, value or field type is an unset kind.
 ///
 /// # Errors
 /// `KRISHIV_UNSUPPORTED_FEATURE` for a type Krishiv does not have yet
-/// (`NULL`, `CHAR`, `VARCHAR`, intervals, nested types, ...) or a
-/// non-default collation, `CONNECT_INVALID_PLAN.*` for an unset kind or an
-/// invalid decimal.
+/// (`NULL`, `CHAR`, `VARCHAR`, intervals, ...) or a non-default collation,
+/// `CONNECT_INVALID_PLAN.*` for an unset kind or an invalid decimal,
+/// `INVALID_JSON_DATA_TYPE` for field metadata that is not a JSON object and
+/// `FAILED_TO_PARSE_TOO_COMPLEX` for a type nested deeper than
+/// `DataType::MAX_NESTING_DEPTH`.
 pub(crate) fn data_type(proto: &Proto) -> Result<DataType, KrishivError> {
+    data_type_at(proto, 0)
+}
+
+fn data_type_at(proto: &Proto, depth: usize) -> Result<DataType, KrishivError> {
+    let child = |t: Option<&Proto>| {
+        if depth >= DataType::MAX_NESTING_DEPTH {
+            return Err(krishiv_common::error::FailedToParseTooComplex {}.into());
+        }
+        data_type_at(t.unwrap_or(&Proto::default()), depth + 1)
+    };
     let Some(kind) = &proto.kind else {
         return Err(not_set("spark.connect.DataType", "KIND_NOT_SET"));
     };
@@ -46,6 +60,27 @@ pub(crate) fn data_type(proto: &Proto) -> Result<DataType, KrishivError> {
             }
             DataType::String
         }
+        Kind::Array(a) => ArrayType::new(child(a.element_type.as_deref())?, a.contains_null).into(),
+        Kind::Map(m) => MapType::new(
+            child(m.key_type.as_deref())?,
+            child(m.value_type.as_deref())?,
+            m.value_contains_null,
+        )
+        .into(),
+        Kind::Struct(s) => {
+            let fields = s
+                .fields
+                .iter()
+                .map(|f| {
+                    let field = StructField::new(&f.name, child(f.data_type.as_ref())?, f.nullable);
+                    Ok(match &f.metadata {
+                        Some(json) => field.with_metadata(krishiv_types::parse_metadata(json)?),
+                        None => field,
+                    })
+                })
+                .collect::<Result<Vec<_>, KrishivError>>()?;
+            StructType::new(fields).into()
+        }
         other => {
             return Err(unsupported_variant(
                 "data type",

crates/krishiv-connect/src/translate/types/tests.rs

@@ -103,9 +103,6 @@ fn unsupported_and_unset_kinds() {
         Kind::Null(dt::Null::default()),
         Kind::Char(dt::Char::default()),
         Kind::VarChar(dt::VarChar::default()),
-        Kind::Array(Box::default()),
-        Kind::Struct(dt::Struct::default()),
-        Kind::Map(Box::default()),
         Kind::Variant(dt::Variant::default()),
         Kind::CalendarInterval(dt::CalendarInterval::default()),
         Kind::YearMonthInterval(dt::YearMonthInterval::default()),
@@ -122,3 +119,81 @@ fn unsupported_and_unset_kinds() {
 }
 
 use crate::proto::v4_2::DataType as DataTypeProto;
+
+/// `array<struct<a:int not null comment 'c', b:map<string,array<double>>>>`
+/// exactly as pyspark-client 4.2 sends it
+/// (`pyspark_types_to_proto_types`): metadata as JSON with spaces, absent
+/// when empty, and the default flags (false) left out.
+#[test]
+fn nested_types_map_as_pyspark_sends_them() {
+    let boxed = |kind| Some(Box::new(proto_type(kind)));
+    let string = Kind::String(dt::String {
+        collation: "UTF8_BINARY".to_owned(),
+        ..dt::String::default()
+    });
+    let map = Kind::Map(Box::new(dt::Map {
+        key_type: boxed(string),
+        value_type: boxed(Kind::Array(Box::new(dt::Array {
+            element_type: boxed(Kind::Double(dt::Double::default())),
+            ..dt::Array::default()
+        }))),
+        ..dt::Map::default()
+    }));
+    let fields = vec![
+        dt::StructField {
+            name: "a".to_owned(),
+            data_type: Some(proto_type(Kind::Integer(dt::Integer::default()))),
+            nullable: false,
+            metadata: Some(r#"{"comment": "c"}"#.to_owned()),
+        },
+        dt::StructField {
+            name: "b".to_owned(),
+            data_type: Some(proto_type(map)),
+            nullable: true,
+            metadata: None,
+        },
+    ];
+    let proto = proto_type(Kind::Array(Box::new(dt::Array {
+        element_type: boxed(Kind::Struct(dt::Struct {
+            fields,
+            ..dt::Struct::default()
+        })),
+        ..dt::Array::default()
+    })));
+    let t = data_type(&proto).expect("maps");
+    assert_eq!(
+        t.json(),
+        r#"{"type":"array","elementType":{"type":"struct","fields":[{"name":"a","type":"integer","nullable":false,"metadata":{"comment":"c"}},{"name":"b","type":{"type":"map","keyType":"string","valueType":{"type":"array","elementType":"double","containsNull":false},"valueContainsNull":false},"nullable":true,"metadata":{}}]},"containsNull":false}"#
+    );
+    // An unset element type is an unset kind; bad metadata is not JSON.
+    let e = data_type(&proto_type(Kind::Array(Box::default()))).expect_err("unset");
+    assert_eq!(
+        e.error_class(),
+        ErrorClass::ConnectInvalidPlanInvalidOneOfFieldNotSet
+    );
+    let bad = proto_type(Kind::Struct(dt::Struct {
+        fields: vec![dt::StructField {
+            name: "a".to_owned(),
+            data_type: Some(proto_type(Kind::Integer(dt::Integer::default()))),
+            nullable: true,
+            metadata: Some("[1]".to_owned()),
+        }],
+        ..dt::Struct::default()
+    }));
+    assert_eq!(
+        data_type(&bad).expect_err("metadata").error_class(),
+        ErrorClass::InvalidJsonDataType
+    );
+    // Nesting is bounded.
+    let mut deep = proto_type(Kind::Integer(dt::Integer::default()));
+    for _ in 0..=DataType::MAX_NESTING_DEPTH {
+        deep = proto_type(Kind::Array(Box::new(dt::Array {
+            element_type: Some(Box::new(deep)),
+            ..dt::Array::default()
+        })));
+    }
+    assert_eq!(
+        data_type(&deep).expect_err("deep").error_class(),
+        ErrorClass::FailedToParseTooComplex
+    );
+}

crates/krishiv-connect/tests/goldens.rs

@@ -24,10 +24,6 @@ use prost_reflect::{DescriptorPool, DynamicMessage};
 const UNSUPPORTED: &[(&str, ErrorClass)] = &[
     ("alias_multi", ErrorClass::KrishivUnsupportedFeature),
     ("aggregate_pivot", ErrorClass::KrishivUnsupportedFeature),
-    (
-        "cast_unsupported_type",
-        ErrorClass::KrishivUnsupportedFeature,
-    ),
     ("distinct", ErrorClass::KrishivUnsupportedFeature),
     (
         "filter_expression_string",

crates/krishiv-connect/tests/snapshots/cast_unsupported_type.snap

@@ -1,4 +1,8 @@
 ---
 source: crates/krishiv-connect/tests/goldens.rs
 ---
-[KRISHIV_UNSUPPORTED_FEATURE] The feature nested data types ("ARRAY") is not supported by Krishiv yet. SQLSTATE: 0A000
+Project [CAST(a AS ARRAY<INT>)]
++- Read parquet [/data/t.parquet]
+-- plan ids --
+Project plan_id=1
+  Read plan_id=0

crates/krishiv-exec/src/planner/eval/cast.rs

@@ -9,6 +9,7 @@ use krishiv_kernels::cast::{CastMode as KernelMode, CastOptions, FixedOffset, ca
 use krishiv_plan::CastMode;
 use krishiv_types::DataType;
 
+use super::nested_cast;
 use super::strings::{from_view, to_view};
 use super::{Input, PhysExpr, Value};
 use crate::error::Result;
@@ -75,8 +76,33 @@ pub(super) fn parts(
     })
 }
 
+impl CastParts {
+    /// The cast of an element, value or field to `to`, in the same mode
+    /// and zone.
+    pub(super) fn child(&self, to: &ArrowType) -> Self {
+        let kernel_to = match to {
+            ArrowType::Utf8View => ArrowType::Utf8,
+            ArrowType::BinaryView => ArrowType::Binary,
+            other => other.clone(),
+        };
+        Self {
+            to: to.clone(),
+            kernel_to,
+            mode: self.mode,
+            zone: self.zone,
+        }
+    }
+}
+
 /// `a` through `p`.
 pub(super) fn apply(p: &CastParts, a: &ArrayRef) -> Result<ArrayRef> {
+    // Every value of a `void` array is NULL, of any type.
+    if a.data_type() == &ArrowType::Null {
+        return Ok(arrow_array::new_null_array(&p.to, a.len()));
+    }
+    if nested_cast::is_nested(a.data_type()) || nested_cast::is_nested(&p.to) {
+        return nested_cast::cast(p, a);
+    }
     let input = from_view(a)?;
     let opts = CastOptions {
         mode: p.mode,

crates/krishiv-exec/src/planner/eval/mod.rs

@@ -27,6 +27,7 @@ mod cast;
 mod cond;
 mod func;
 mod logic;
+mod nested_cast;
 mod strings;
 
 use std::fmt;