From 8d5745e66b8ae0aca37f59a103bc12436c6b558d Mon Sep 17 00:00:00 2001 From: shuiyisong <113876041+shuiyisong@users.noreply.github.com> Date: Sat, 25 Jul 2026 16:33:49 +0800 Subject: [PATCH] fix(datatypes): replicate nested list and struct vectors (#8638) * fix(datatypes): replicate nested list and struct vectors Signed-off-by: shuiyisong * chore: fix issues Signed-off-by: shuiyisong --------- Signed-off-by: shuiyisong --- src/datatypes/src/schema/column_schema.rs | 19 ++++ src/datatypes/src/vectors/operations.rs | 20 +++- .../src/vectors/operations/replicate.rs | 46 +++++++- src/datatypes/src/vectors/struct_vector.rs | 101 +++++++++++++++--- 4 files changed, 167 insertions(+), 19 deletions(-) diff --git a/src/datatypes/src/schema/column_schema.rs b/src/datatypes/src/schema/column_schema.rs index 2479f4fc41..d3ecd90116 100644 --- a/src/datatypes/src/schema/column_schema.rs +++ b/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); diff --git a/src/datatypes/src/vectors/operations.rs b/src/datatypes/src/vectors/operations.rs index 4eb74b0307..1dc9e21a0d 100644 --- a/src/datatypes/src/vectors/operations.rs +++ b/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 { + filter::filter_non_constant!(self, ListVector, filter) + } + + fn cast(&self, to_type: &ConcreteDataType) -> Result { + cast::cast_non_constant!(self, to_type) + } + + fn take(&self, indices: &UInt32Vector) -> Result { + take::take_indices!(self, ListVector, indices) + } +} impl VectorOp for Decimal128Vector { fn replicate(&self, offsets: &[usize]) -> VectorRef { diff --git a/src/datatypes/src/vectors/operations/replicate.rs b/src/datatypes/src/vectors/operations/replicate.rs index 9e0591efa3..dee0b3855b 100644 --- a/src/datatypes/src/vectors/operations/replicate.rs +++ b/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: &C, offsets: &[usize]) -> VectorRef { assert_eq!(offsets.len(), c.len()); @@ -36,6 +37,26 @@ pub(crate) fn replicate_scalar(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"])); diff --git a/src/datatypes/src/vectors/struct_vector.rs b/src/datatypes/src/vectors/struct_vector.rs index 79abef255b..e6284351b4 100644 --- a/src/datatypes/src/vectors/struct_vector.rs +++ b/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::>(); - 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 { @@ -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::::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