plan, memory, plan-codec, session, exec: nested literals, a declared file schema and the LocalRelation carrier (P3-T012a part 2a, PR-12)

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).

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),
+                )
             }
         }))
     }

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(),
     }
 }

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

@@ -1090,10 +1090,88 @@ pub(crate) fn literal_array(l: &Literal) -> Result<ArrayRef> {
                 .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::<Result<Vec<_>>>()?;
+            Arc::new(
+                StructArray::try_new_with_length(arrow_fields, columns, None, 1)
+                    .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<Item = &'a Literal>,
+    data_type: &ArrowType,
+) -> Result<ArrayRef> {
+    let arrays = values.map(literal_array).collect::<Result<Vec<_>>>()?;
+    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>| {
+            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::<i32>().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<Option<i32>> = row
+            .as_primitive::<arrow_array::types::Int32Type>()
+            .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::<i32>().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]

crates/krishiv-exec/src/planner/pushdown.rs

@@ -243,6 +243,11 @@ fn value_or_null(e: &Expr) -> Option<Option<ScalarValue>> {
         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, maps, intervals.
+        Literal::Array { .. } | Literal::KeyValueMap { .. } | Literal::Struct { .. } => {
+            return None;
+        }
         _ => return None,
     }))
 }

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.

crates/krishiv-exec/src/planner/table.rs

@@ -253,7 +253,8 @@ impl TableProviderResolver for MemoryTables {
     fn provider(&self, scan: &Scan) -> Result<Arc<dyn TableProvider>> {
         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());
         };

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());