[Variant] Support list paths in ShreddedSchemaBuilder

parquet-variant-compute/src/shred_variant.rs

@@ -528,8 +528,9 @@ impl IntoShreddingField for (DataType, bool) {
 /// should be shredded and with what types. Fields are nullable by default; pass
 /// a `(data_type, nullable)` pair or a `FieldRef` to control nullability.
 ///
-/// Note: this builder currently only supports struct fields. List support
-/// will be added in the future.
+/// Numeric indexes represent the shared element schema of a list. The index
+/// value itself is ignored, so `items[0].id` and `items[1].name` describe fields
+/// on the same list element struct.
 ///
 /// # Example
 ///
@@ -556,6 +557,8 @@ impl IntoShreddingField for (DataType, bool) {
 ///         VariantPath::from_iter([VariantPathElement::from("metrics.cpu")]),
 ///         &DataType::Float64,
 ///     )?
+///     // indexes describe the shared schema for every element of a list
+///     .with_path("items[0].id", &DataType::Int64)?
 ///     .build();
 ///    Ok(())
 /// }
@@ -614,6 +617,8 @@ enum VariantSchemaNode {
     Leaf(ShreddingField),
     /// An inner struct node with nested fields
     Struct(BTreeMap<String, VariantSchemaNode>),
+    /// An inner list node with a shared element schema
+    List(Box<VariantSchemaNode>),
 }
 
 impl Default for VariantSchemaNode {
@@ -654,8 +659,18 @@ impl VariantSchemaNode {
                     .insert_path_elements(tail, field);
             }
             VariantPathElement::Index { .. } => {
-                // List support to be added later; reject for now
-                unreachable!("List paths are not supported yet");
+                let element = match self {
+                    Self::List(element) => element,
+                    _ => {
+                        *self = Self::List(Box::default());
+                        match self {
+                            Self::List(element) => element,
+                            _ => unreachable!(),
+                        }
+                    }
+                };
+
+                element.insert_path_elements(tail, field);
             }
         }
     }
@@ -677,6 +692,7 @@ impl VariantSchemaNode {
                     Some(DataType::Struct(Fields::from(child_fields)))
                 }
             }
+            Self::List(element) => element.to_shredding_field("item").map(DataType::List),
         }
     }
 
@@ -687,7 +703,7 @@ impl VariantSchemaNode {
                 field.data_type.clone(),
                 field.nullable,
             ))),
-            Self::Struct(_) => self
+            Self::Struct(_) | Self::List(_) => self
                 .to_shredding_type()
                 .map(|data_type| Arc::new(Field::new(name, data_type, true))),
         }
@@ -1847,15 +1863,12 @@ mod tests {
         ]);
 
         // Target schema is List<Struct<id:int64,name:utf8>>
-        let object_fields = Fields::from(vec![
-            Field::new("id", DataType::Int64, true),
-            Field::new("name", DataType::Utf8, true),
-        ]);
-        let list_schema = DataType::List(Arc::new(Field::new(
-            "item",
-            DataType::Struct(object_fields),
-            true,
-        )));
+        let list_schema = ShreddedSchemaBuilder::default()
+            .with_path("[0].id", &DataType::Int64)
+            .unwrap()
+            .with_path("[0].name", &DataType::Utf8)
+            .unwrap()
+            .build();
         let result = shred_variant(&input, &list_schema).unwrap();
         assert_eq!(result.len(), 3);
 
@@ -2898,6 +2911,49 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_variant_schema_builder_list() -> Result<()> {
+        let shredding_type = ShreddedSchemaBuilder::default()
+            .with_path("items[0].id", &DataType::Int64)?
+            .with_path("items[42].name", &DataType::Utf8)?
+            .build();
+
+        assert_eq!(
+            shredding_type,
+            DataType::Struct(Fields::from(vec![Field::new(
+                "items",
+                DataType::new_list(
+                    DataType::Struct(Fields::from(vec![
+                        Field::new("id", DataType::Int64, true),
+                        Field::new("name", DataType::Utf8, true),
+                    ])),
+                    true,
+                ),
+                true,
+            )]))
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    fn test_variant_schema_builder_nested_lists() -> Result<()> {
+        let shredding_type = ShreddedSchemaBuilder::default()
+            .with_path("matrix[0][1]", (&DataType::Float64, false))?
+            .build();
+
+        assert_eq!(
+            shredding_type,
+            DataType::Struct(Fields::from(vec![Field::new(
+                "matrix",
+                DataType::new_list(DataType::new_list(DataType::Float64, false), true),
+                true,
+            )]))
+        );
+
+        Ok(())
+    }
+
     #[test]
     fn test_variant_schema_builder_with_path_variant_path_arg() -> Result<()> {
         let path = VariantPath::from_iter([VariantPathElement::from("a.b")]);