fix(datatypes): replicate nested list and struct vectors (#8638)

* fix(datatypes): replicate nested list and struct vectors

Signed-off-by: shuiyisong <xixing.sys@gmail.com>

* chore: fix issues

Signed-off-by: shuiyisong <xixing.sys@gmail.com>

---------

Signed-off-by: shuiyisong <xixing.sys@gmail.com>

src/datatypes/src/schema/column_schema.rs

@@ -1230,6 +1230,7 @@ mod tests {
     use arrow::datatypes::{DataType as ArrowDataType, TimeUnit};
 
     use super::*;
+    use crate::types::{StructField, StructType};
     use crate::value::Value;
     use crate::vectors::Int32Vector;
 
@@ -1343,6 +1344,24 @@ mod tests {
         assert!(v.only_null());
     }
 
+    #[test]
+    fn test_column_schema_create_default_null_struct_with_list() {
+        let list_type =
+            ConcreteDataType::list_datatype(Arc::new(ConcreteDataType::int32_datatype()));
+        let struct_type =
+            ConcreteDataType::struct_datatype(StructType::new(Arc::new(vec![StructField::new(
+                "values".to_string(),
+                list_type,
+                true,
+            )])));
+        let column_schema = ColumnSchema::new("test", struct_type, true);
+
+        let v = column_schema.create_default_vector(5).unwrap().unwrap();
+
+        assert_eq!(5, v.len());
+        assert!(v.only_null());
+    }
+
     #[test]
     fn test_column_schema_no_default() {
         let column_schema = ColumnSchema::new("test", ConcreteDataType::int32_datatype(), false);

src/datatypes/src/vectors/operations.rs

@@ -89,7 +89,25 @@ macro_rules! impl_scalar_vector_op {
     )+};
 }
 
-impl_scalar_vector_op!(BinaryVector, BooleanVector, ListVector, StringVector);
+impl_scalar_vector_op!(BinaryVector, BooleanVector, StringVector);
+
+impl VectorOp for ListVector {
+    fn replicate(&self, offsets: &[usize]) -> VectorRef {
+        replicate::replicate_list(self, offsets)
+    }
+
+    fn filter(&self, filter: &BooleanVector) -> Result<VectorRef> {
+        filter::filter_non_constant!(self, ListVector, filter)
+    }
+
+    fn cast(&self, to_type: &ConcreteDataType) -> Result<VectorRef> {
+        cast::cast_non_constant!(self, to_type)
+    }
+
+    fn take(&self, indices: &UInt32Vector) -> Result<VectorRef> {
+        take::take_indices!(self, ListVector, indices)
+    }
+}
 
 impl VectorOp for Decimal128Vector {
     fn replicate(&self, offsets: &[usize]) -> VectorRef {

src/datatypes/src/vectors/operations/replicate.rs

@@ -16,6 +16,7 @@ use crate::prelude::*;
 pub(crate) use crate::vectors::decimal::replicate_decimal128;
 pub(crate) use crate::vectors::null::replicate_null;
 pub(crate) use crate::vectors::primitive::replicate_primitive;
+use crate::vectors::{ListVector, ListVectorBuilder};
 
 pub(crate) fn replicate_scalar<C: ScalarVector>(c: &C, offsets: &[usize]) -> VectorRef {
     assert_eq!(offsets.len(), c.len());
@@ -36,6 +37,26 @@ pub(crate) fn replicate_scalar<C: ScalarVector>(c: &C, offsets: &[usize]) -> Vec
     builder.to_vector()
 }
 
+pub(crate) fn replicate_list(c: &ListVector, offsets: &[usize]) -> VectorRef {
+    assert_eq!(offsets.len(), c.len());
+
+    if offsets.is_empty() {
+        return c.slice(0, 0);
+    }
+    let mut builder =
+        ListVectorBuilder::with_type_capacity(c.item_type(), *offsets.last().unwrap());
+
+    let mut previous_offset = 0;
+    for (i, offset) in offsets.iter().enumerate() {
+        let data = c.get_data(i);
+        for _ in previous_offset..*offset {
+            builder.push(data.clone());
+        }
+        previous_offset = *offset;
+    }
+    builder.to_vector()
+}
+
 #[cfg(test)]
 mod tests {
     use std::sync::Arc;
@@ -45,8 +66,11 @@ mod tests {
     use paste::paste;
 
     use super::*;
+    use crate::value::{ListValue, ListValueRef};
     use crate::vectors::constant::ConstantVector;
-    use crate::vectors::{Decimal128Vector, Int32Vector, NullVector, StringVector, VectorOp};
+    use crate::vectors::{
+        Decimal128Vector, Int32Vector, ListVectorBuilder, NullVector, StringVector, VectorOp,
+    };
 
     #[test]
     fn test_replicate_primitive() {
@@ -93,6 +117,26 @@ mod tests {
         assert_eq!(expect, v);
     }
 
+    #[test]
+    fn test_replicate_list() {
+        let item_type = Arc::new(ConcreteDataType::int32_datatype());
+        let first = ListValue::new(vec![Value::Int32(1), Value::Int32(2)], item_type.clone());
+        let second = ListValue::new(vec![Value::Int32(3)], item_type.clone());
+        let mut builder = ListVectorBuilder::with_type_capacity(item_type, 2);
+        builder.push(Some(ListValueRef::Ref { val: &first }));
+        builder.push(Some(ListValueRef::Ref { val: &second }));
+        let v = builder.finish();
+
+        let v = v.replicate(&[1, 3]);
+        let mut expect_builder =
+            ListVectorBuilder::with_type_capacity(Arc::new(ConcreteDataType::int32_datatype()), 3);
+        expect_builder.push(Some(ListValueRef::Ref { val: &first }));
+        expect_builder.push(Some(ListValueRef::Ref { val: &second }));
+        expect_builder.push(Some(ListValueRef::Ref { val: &second }));
+        let expect = expect_builder.to_vector();
+        assert_eq!(expect, v);
+    }
+
     #[test]
     fn test_replicate_constant() {
         let v = Arc::new(StringVector::from_slice(&["hello"]));

src/datatypes/src/vectors/struct_vector.rs

@@ -15,7 +15,7 @@
 use std::any::Any;
 use std::sync::Arc;
 
-use arrow::array::NullBufferBuilder;
+use arrow::array::{MutableArrayData, NullBufferBuilder};
 use arrow::compute::TakeOptions;
 use arrow::datatypes::DataType as ArrowDataType;
 use arrow_array::{Array, ArrayRef, StructArray};
@@ -144,22 +144,28 @@ impl Vector for StructVector {
 
 impl VectorOp for StructVector {
     fn replicate(&self, offsets: &[usize]) -> VectorRef {
-        let column_arrays = self
-            .array
-            .columns()
-            .iter()
-            .map(|col| {
-                let vector = Helper::try_into_vector(col)
-                    .expect("Failed to replicate struct vector columns");
-                vector.replicate(offsets).to_arrow_array()
-            })
-            .collect::<Vec<_>>();
-        let replicated_array = StructArray::new(
-            self.array.fields().clone(),
-            column_arrays,
-            self.array.nulls().cloned(),
-        );
-        Arc::new(StructVector::try_new(self.fields.clone(), replicated_array).unwrap())
+        assert_eq!(offsets.len(), self.len());
+        assert!(offsets.is_sorted(), "offsets must be non-decreasing");
+
+        let Some(&output_len) = offsets.last() else {
+            return self.slice(0, 0);
+        };
+
+        let source = self.array.to_data();
+        let mut output = MutableArrayData::new(vec![&source], false, output_len);
+        let mut previous_offset = 0;
+
+        for (index, &offset) in offsets.iter().enumerate() {
+            for _ in previous_offset..offset {
+                output.extend(0, index, index + 1);
+            }
+            previous_offset = offset;
+        }
+
+        Arc::new(StructVector {
+            array: StructArray::from(output.freeze()),
+            fields: self.fields.clone(),
+        })
     }
 
     fn cast(&self, _to_type: &ConcreteDataType) -> Result<VectorRef> {
@@ -460,7 +466,13 @@ impl ScalarVectorBuilder for StructVectorBuilder {
 
 #[cfg(test)]
 mod tests {
+    use arrow::array::{DictionaryArray, Int8Array, StringArray};
+    use arrow::buffer::NullBuffer;
+    use arrow::datatypes::Int8Type;
+
     use super::*;
+    use crate::json::JsonSettings;
+    use crate::schema::{ColumnDefaultConstraint, ColumnSchema};
     use crate::types::StructField;
     use crate::value::ListValue;
     use crate::value::tests::*;
@@ -542,6 +554,61 @@ mod tests {
         assert_eq!(vector.get(1), Value::Null);
     }
 
+    #[test]
+    fn test_replicate_preserves_json2_identity() {
+        let json = JsonSettings::default()
+            .encode(serde_json::json!({"answer": 42}))
+            .unwrap();
+        let fields = StructType::new(Arc::new(vec![StructField::new(
+            "payload",
+            json.data_type(),
+            true,
+        )]));
+        let data_type = ConcreteDataType::struct_datatype(fields.clone());
+        let value = Value::Struct(StructValue::new(vec![json], fields));
+        let schema = ColumnSchema::new("nested", data_type.clone(), true)
+            .with_default_constraint(Some(ColumnDefaultConstraint::Value(value)))
+            .unwrap();
+
+        let replicated = schema.create_default_vector(2).unwrap().unwrap();
+
+        assert_eq!(replicated.data_type(), data_type);
+        assert_eq!(replicated.len(), 2);
+    }
+
+    #[test]
+    fn test_replicate_preserves_dictionary_and_nulls() {
+        let fields = StructType::new(Arc::new(vec![StructField::new(
+            "label",
+            ConcreteDataType::dictionary_datatype(
+                ConcreteDataType::int8_datatype(),
+                ConcreteDataType::string_datatype(),
+            ),
+            true,
+        )]));
+        let dictionary = DictionaryArray::<Int8Type>::new(
+            Int8Array::from(vec![Some(0), Some(1)]),
+            Arc::new(StringArray::from(vec!["a", "b"])),
+        );
+        let array = StructArray::new(
+            fields.as_arrow_fields(),
+            vec![Arc::new(dictionary)],
+            Some(NullBuffer::from(vec![true, false])),
+        );
+        let vector = StructVector::try_new(fields.clone(), array).unwrap();
+
+        let replicated = vector.replicate(&[2, 3]);
+
+        assert_eq!(
+            replicated.data_type(),
+            ConcreteDataType::struct_datatype(fields)
+        );
+        assert_eq!(replicated.len(), 3);
+        assert_eq!(replicated.null_count(), 1);
+        assert_eq!(replicated.get(0), replicated.get(1));
+        assert!(replicated.is_null(2));
+    }
+
     #[test]
     fn test_deep_nested_struct_list() {
         // level 1: struct