diff --git a/src/datatypes/src/vectors/json/array.rs b/src/datatypes/src/vectors/json/array.rs index 9def5735bb..f2a5c426a9 100644 --- a/src/datatypes/src/vectors/json/array.rs +++ b/src/datatypes/src/vectors/json/array.rs @@ -18,9 +18,9 @@ use std::sync::Arc; use arrow::compute; use arrow::util::display::{ArrayFormatter, FormatOptions}; use arrow_array::cast::AsArray; -use arrow_array::{Array, ArrayRef, GenericListArray, StructArray, new_null_array}; -use arrow_schema::DataType; -use snafu::{OptionExt, ResultExt}; +use arrow_array::{Array, ArrayRef, GenericListArray, ListArray, StructArray, new_null_array}; +use arrow_schema::{DataType, FieldRef}; +use snafu::{OptionExt, ResultExt, ensure}; use crate::arrow_array::StringArray; use crate::error::{AlignJsonArraySnafu, ArrowComputeSnafu, Result}; @@ -81,24 +81,7 @@ impl JsonArray<'_> { } (DataType::List(expect_item), DataType::List(array_item)) => { let list_array = array_columns[j].as_list::(); - let item_aligned = - match (expect_item.data_type(), array_item.data_type()) { - (DataType::Struct(_), DataType::Struct(_)) => { - JsonArray::from(list_array.values()) - .try_align(expect_item.data_type())? - } - _ => JsonArray::from(list_array.values()) - .try_cast(expect_item.data_type())?, - }; - Arc::new( - GenericListArray::::try_new( - expect_item.clone(), - list_array.offsets().clone(), - item_aligned, - list_array.nulls().cloned(), - ) - .context(ArrowComputeSnafu)?, - ) + try_align_list(list_array, expect_item, array_item)? } _ => JsonArray::from(&array_columns[j]).try_cast(expect_type)?, }; @@ -124,8 +107,9 @@ impl JsonArray<'_> { aligned.push(new_null_array(field.data_type(), struct_array.len())); } } - if j < array_fields.len() { - return AlignJsonArraySnafu { + ensure!( + j < array_fields.len(), + AlignJsonArraySnafu { reason: format!( "extra fields are found: [{}]", array_fields[j..] @@ -135,8 +119,7 @@ impl JsonArray<'_> { .join(", ") ), } - .fail(); - } + ); let json_array = StructArray::try_new( expect_fields.clone(), @@ -157,6 +140,7 @@ impl JsonArray<'_> { return compute::cast(&self.inner, to_type).context(ArrowComputeSnafu); } + // TODO(LFC): Cast according to `to_type` instead of formatting to String here. let formatter = ArrayFormatter::try_new(&self.inner, &FormatOptions::default()) .context(ArrowComputeSnafu)?; let values = (0..self.inner.len()) @@ -170,6 +154,32 @@ impl JsonArray<'_> { } } +fn try_align_list( + list_array: &ListArray, + expect_item: &FieldRef, + array_item: &FieldRef, +) -> Result { + let item_aligned = match (expect_item.data_type(), array_item.data_type()) { + (DataType::Struct(_), DataType::Struct(_)) => { + JsonArray::from(list_array.values()).try_align(expect_item.data_type())? + } + (DataType::List(expect_item), DataType::List(array_item)) => { + let list_array = list_array.values().as_list::(); + try_align_list(list_array, expect_item, array_item)? + } + _ => JsonArray::from(list_array.values()).try_cast(expect_item.data_type())?, + }; + Ok(Arc::new( + GenericListArray::::try_new( + expect_item.clone(), + list_array.offsets().clone(), + item_aligned, + list_array.nulls().cloned(), + ) + .context(ArrowComputeSnafu)?, + )) +} + impl<'a> From<&'a ArrayRef> for JsonArray<'a> { fn from(inner: &'a ArrayRef) -> Self { Self { inner }