mirror of
https://github.com/GreptimeTeam/greptimedb.git
synced 2026-05-17 21:40:37 +00:00
@@ -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 }
|
||||
|
||||
Reference in New Issue
Block a user