resolve PR comments

Signed-off-by: luofucong <luofc@foxmail.com>
This commit is contained in:
luofucong
2026-04-22 21:07:08 +08:00
parent 2fc6c2da7b
commit b4dc6c2b8e

View File

@@ -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::<i32>();
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::<i32>::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<ArrayRef> {
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::<i32>();
try_align_list(list_array, expect_item, array_item)?
}
_ => JsonArray::from(list_array.values()).try_cast(expect_item.data_type())?,
};
Ok(Arc::new(
GenericListArray::<i32>::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 }