feat: index for DataType::List (#562)

Adds type inference support for lists and adds just the `index`
function.

Further additions: #494

crates/sparrow-api/src/kaskada/v1alpha/schema_traits.rs

@@ -24,7 +24,7 @@ impl DataType {
  let value = &fields[1];
  Self {
  kind: Some(data_type::Kind::Map(Box::new(data_type::Map {
- name: name.to_string(),
+ name: name.to_owned(),
  ordered,
  key_name: key.name.clone(),
  key_type: Some(Box::new(
@@ -44,6 +44,22 @@ impl DataType {
  }
  }
 
+ pub fn new_list(name: &str, field: schema::Field) -> Self {
+ Self {
+ kind: Some(data_type::Kind::List(Box::new(data_type::List {
+ name: name.to_owned(),
+ item_type: Some(Box::new(
+ field
+ .data_type
+ .as_ref()
+ .expect("data type to exist")
+ .clone(),
+ )),
+ nullable: field.nullable,
+ }))),
+ }
+ }
+
  pub fn new_primitive(primitive: data_type::PrimitiveType) -> Self {
  Self {
  kind: Some(data_type::Kind::Primitive(primitive as i32)),
@@ -221,6 +237,21 @@ impl TryFrom<&arrow::datatypes::DataType> for DataType {
 
  Ok(DataType::new_map(s.name(), *is_ordered, vec![key, value]))
  }
+ arrow::datatypes::DataType::List(field) => {
+ let name = field.name();
+ let field = schema::Field {
+ name: name.to_owned(),
+ data_type: Some(field.data_type().try_into().map_err(
+ |err: ConversionError<arrow::datatypes::DataType>| {
+ err.with_prepend_field("list item".to_owned())
+ },
+ )?),
+ nullable: field.is_nullable(),
+ };
+
+ Ok(DataType::new_list(name, field))
+ }
+
  unsupported => Err(ConversionError::new_unsupported(unsupported.clone())),
  }
  }
@@ -337,12 +368,16 @@ impl TryFrom<&DataType> for arrow::datatypes::DataType {
  Some(data_type::Kind::Struct(schema)) => Ok(arrow::datatypes::DataType::Struct(
  fields_to_arrow(&schema.fields)?.into(),
  )),
- Some(data_type::Kind::List(item_type)) => {
- let item_type = arrow::datatypes::DataType::try_from(item_type.as_ref())
- .map_err(|e| e.with_prepend_field("list item".to_owned()))?;
- let item_type = arrow::datatypes::Field::new("item", item_type, true);
- Ok(arrow::datatypes::DataType::List(Arc::new(item_type)))
- }
+ Some(data_type::Kind::List(list)) => match list.item_type.as_ref() {
+ Some(item_type) => {
+ let item_type = arrow::datatypes::DataType::try_from(item_type.as_ref())
+ .map_err(|e| e.with_prepend_field("list item".to_owned()))?;
+ let item_type =
+ arrow::datatypes::Field::new(list.name.clone(), item_type, list.nullable);
+ Ok(arrow::datatypes::DataType::List(Arc::new(item_type)))
+ }
+ None => Err(ConversionError::new_unsupported(value.clone())),
+ },
  Some(data_type::Kind::Map(map)) => {
  match (map.key_type.as_ref(), map.value_type.as_ref()) {
  (Some(key), Some(value)) => {

crates/sparrow-arrow/src/downcast.rs

@@ -2,8 +2,7 @@
 
 use anyhow::Context;
 use arrow::array::{
- Array, BooleanArray, GenericStringArray, ListArray, OffsetSizeTrait, PrimitiveArray,
- StructArray,
+ Array, BooleanArray, GenericStringArray, OffsetSizeTrait, PrimitiveArray, StructArray,
 };
 use arrow::datatypes::ArrowPrimitiveType;
 use arrow_array::MapArray;
@@ -24,13 +23,6 @@ pub fn downcast_primitive_array<T: ArrowPrimitiveType>(
  })
 }
 
-pub fn downcast_list_array(array: &dyn Array) -> anyhow::Result<&ListArray> {
- array
- .as_any()
- .downcast_ref::<ListArray>()
- .with_context(|| format!("Unable to downcast {:?} to ListArray", array.data_type()))
-}
-
 /// Downcast an array into a string array.
 pub fn downcast_string_array<T>(array: &dyn Array) -> anyhow::Result<&GenericStringArray<T>>
 where

crates/sparrow-compiler/src/ast_to_dfg.rs

@@ -674,6 +674,13 @@ fn cast_if_needed(
  {
  Ok(value)
  }
+ // Ensures that list types with the same inner types are compatible, regardless of the (arbitary) field naming.
+ (FenlType::Concrete(DataType::List(s)), FenlType::Concrete(DataType::List(s2)))
+ if list_types_are_equal(s, s2) =>
+ {
+ Ok(value)
+ }
+
  (FenlType::Concrete(DataType::Null), FenlType::Window) => Ok(value),
  (
  FenlType::Concrete(DataType::Struct(actual_fields)),
@@ -733,6 +740,14 @@ fn map_types_are_equal(a: &FieldRef, b: &FieldRef) -> bool {
  }
 }
 
+// When constructing the concrete list during inference, we use arbitary names for the inner data
+// field since we don't have access to the user's naming patterns there.
+// By comparing the list types based on just the inner type, we can ensure that the types are
+// still treated as equal.
+fn list_types_are_equal(a: &FieldRef, b: &FieldRef) -> bool {
+ a.data_type() == b.data_type()
+}
+
 pub(crate) fn is_any_new(dfg: &mut Dfg, arguments: &[Located<AstDfgRef>]) -> anyhow::Result<Id> {
  let mut argument_is_new = arguments.iter().map(|a| a.is_new()).unique();
  let mut result = argument_is_new

crates/sparrow-compiler/src/dfg.rs

@@ -108,7 +108,6 @@ impl Default for Dfg {
 impl Dfg {
  pub(super) fn add_literal(&mut self, literal: impl Into<ScalarValue>) -> anyhow::Result<Id> {
  let literal = literal.into();
- // TODO: FRAZ - do I need to support large string literal here?
  if let ScalarValue::Utf8(Some(literal)) = literal {
  self.add_string_literal(&literal)
  } else {

crates/sparrow-compiler/src/functions/collection.rs

@@ -7,4 +7,9 @@ pub(super) fn register(registry: &mut Registry) {
  .register("get<K: key, V: any>(key: K, map: map<K, V>) -> V")
  .with_implementation(Implementation::Instruction(InstOp::Get))
  .set_internal();
+
+ registry
+ .register("index<T: any>(i: i64, list: list<T>) -> T")
+ .with_implementation(Implementation::Instruction(InstOp::Index))
+ .set_internal();
 }

crates/sparrow-compiler/src/functions/pushdown.rs

@@ -111,6 +111,7 @@ impl Pushdown {
  | DataType::Interval(_)
  | DataType::Utf8
  | DataType::LargeUtf8
+ | DataType::List(..)
  | DataType::Map(..) => {
  let mut subst = subst.clone();
  subst.insert(

crates/sparrow-compiler/src/plan/operation_to_plan.rs

@@ -21,10 +21,16 @@ use crate::DataContext;
 /// DataType protobuf representing a list of u64.
 #[static_init::dynamic]
 static LIST_U64_DATA_TYPE: DataType = DataType {
- kind: Some(data_type::Kind::List(Box::new(DataType {
- kind: Some(data_type::Kind::Primitive(
- data_type::PrimitiveType::U64 as i32,
- )),
+ kind: Some(data_type::Kind::List(Box::new(data_type::List {
+ // Note: The fields here must match the default fields used when creating
+ // types during type inference, otherwise schema validation will fail.
+ name: "item".to_owned(),
+ item_type: Some(Box::new(DataType {
+ kind: Some(data_type::Kind::Primitive(
+ data_type::PrimitiveType::U64 as i32,
+ )),
+ })),
+ nullable: true,
  }))),
 };
 

crates/sparrow-compiler/src/types/inference.rs

@@ -210,8 +210,30 @@ pub fn validate_instantiation(
  }
  }
  }
- FenlType::Collection(Collection::List, _) => {
- todo!("list unsupported")
+ FenlType::Collection(Collection::List, type_vars) => {
+ debug_assert!(type_vars.len() == 1);
+ let item_type = match argument_type {
+ FenlType::Concrete(DataType::List(f)) => {
+ FenlType::Concrete(f.data_type().clone())
+ }
+ other => anyhow::bail!("expected list, saw {:?}", other),
+ };
+
+ match types_for_variable.entry(type_vars[0].clone()) {
+ Entry::Occupied(occupied) => {
+ anyhow::ensure!(
+ occupied.get() == &item_type
+ || matches!(occupied.get(), FenlType::Error)
+ || matches!(argument_type, FenlType::Error),
+ "Failed type validation: expected {} but was {}",
+ occupied.get(),
+ item_type
+ );
+ }
+ Entry::Vacant(vacant) => {
+ vacant.insert(item_type.clone());
+ }
+ }
  }
  FenlType::Error => {
  // Assume the argument matches (since we already reported what
@@ -359,7 +381,20 @@ fn instantiate_type(fenl_type: &FenlType, solutions: &HashMap<TypeVariable, Fenl
  let s = Arc::new(Field::new("entries", DataType::Struct(fields), false));
  FenlType::Concrete(DataType::Map(s, false))
  }
- FenlType::Collection(Collection::List, _) => todo!("unsupported"),
+ FenlType::Collection(Collection::List, type_vars) => {
+ debug_assert!(type_vars.len() == 1);
+ let concrete_type = solutions
+ .get(&type_vars[0])
+ .cloned()
+ .unwrap_or(FenlType::Concrete(DataType::Null));
+ let field = match concrete_type {
+ // TODO: Should the fields be nullable?
+ FenlType::Concrete(t) => Arc::new(Field::new("item", t, false)),
+ other => panic!("expected concrete type, got {:?}", other),
+ };
+
+ FenlType::Concrete(DataType::List(field))
+ }
  FenlType::Concrete(_) => fenl_type.clone(),
  FenlType::Window => fenl_type.clone(),
  FenlType::Json => fenl_type.clone(),

crates/sparrow-compiler/tests/snapshots/compiler_golden_tests__lookup_diff_grouping.snap

@@ -227,8 +227,11 @@ operations:
  result_type:
  kind:
  List:
- kind:
- Primitive: 10
+ name: item
+ item_type:
+ kind:
+ Primitive: 10
+ nullable: true
  output: true
  operator:
  Input:
@@ -251,8 +254,11 @@ operations:
  result_type:
  kind:
  List:
- kind:
- Primitive: 10
+ name: item
+ item_type:
+ kind:
+ Primitive: 10
+ nullable: true
  output: true
  operator:
  Input:

crates/sparrow-compiler/tests/snapshots/compiler_golden_tests__lookup_same_grouping.snap

@@ -217,8 +217,11 @@ operations:
  result_type:
  kind:
  List:
- kind:
- Primitive: 10
+ name: item
+ item_type:
+ kind:
+ Primitive: 10
+ nullable: true
  output: true
  operator:
  Input:
@@ -241,8 +244,11 @@ operations:
  result_type:
  kind:
  List:
- kind:
- Primitive: 10
+ name: item
+ item_type:
+ kind:
+ Primitive: 10
+ nullable: true
  output: true
  operator:
  Input:

crates/sparrow-compiler/tests/snapshots/compiler_golden_tests__lookup_same_grouping_with_slicing.snap

@@ -219,8 +219,11 @@ operations:
  result_type:
  kind:
  List:
- kind:
- Primitive: 10
+ name: item
+ item_type:
+ kind:
+ Primitive: 10
+ nullable: true
  output: true
  operator:
  Input:
@@ -243,8 +246,11 @@ operations:
  result_type:
  kind:
  List:
- kind:
- Primitive: 10
+ name: item
+ item_type:
+ kind:
+ Primitive: 10
+ nullable: true
  output: true
  operator:
  Input:

crates/sparrow-instructions/src/evaluators.rs

@@ -16,6 +16,7 @@ mod equality;
 mod field_ref;
 mod general;
 mod json_field;
+mod list;
 mod logical;
 mod macros;
 mod map;
@@ -31,6 +32,7 @@ use equality::*;
 use field_ref::*;
 use general::*;
 use json_field::*;
+use list::*;
 use logical::*;
 use map::*;
 use math::*;
@@ -204,6 +206,7 @@ fn create_simple_evaluator(
  }
  InstOp::Floor => FloorEvaluator::try_new(info),
  InstOp::Get => GetEvaluator::try_new(info),
+ InstOp::Index => IndexEvaluator::try_new(info),
  InstOp::Gt => match (info.args[0].is_literal(), info.args[1].is_literal()) {
  (_, true) => {
  create_ordered_evaluator!(&info.args[0].data_type, GtScalarEvaluator, info)

crates/sparrow-instructions/src/evaluators/list.rs

@@ -0,0 +1,2 @@
+mod index;
+pub(super) use index::*;