Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Preserve field name when casting List #13468

Merged
merged 4 commits into from
Nov 25, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 64 additions & 8 deletions datafusion/common/src/scalar/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,9 @@ use crate::cast::{
use crate::error::{DataFusionError, Result, _exec_err, _internal_err, _not_impl_err};
use crate::hash_utils::create_hashes;
use crate::utils::{
array_into_fixed_size_list_array, array_into_large_list_array, array_into_list_array,
array_into_fixed_size_list_array_with_field_name, array_into_large_list_array,
array_into_large_list_array_with_field_name, array_into_list_array,
array_into_list_array_with_field_name,
};
use arrow::compute::kernels::numeric::*;
use arrow::util::display::{array_value_to_string, ArrayFormatter, FormatOptions};
Expand Down Expand Up @@ -2663,27 +2665,36 @@ impl ScalarValue {
let list_array = array.as_list::<i32>();
let nested_array = list_array.value(index);
// Produces a single element `ListArray` with the value at `index`.
let arr =
Arc::new(array_into_list_array(nested_array, field.is_nullable()));
let arr = Arc::new(array_into_list_array_with_field_name(
nested_array,
field.is_nullable(),
field.name(),
));

ScalarValue::List(arr)
}
DataType::LargeList(_) => {
DataType::LargeList(field) => {
let list_array = as_large_list_array(array);
let nested_array = list_array.value(index);
// Produces a single element `LargeListArray` with the value at `index`.
let arr = Arc::new(array_into_large_list_array(nested_array));
let arr = Arc::new(array_into_large_list_array_with_field_name(
nested_array,
field.name(),
));

ScalarValue::LargeList(arr)
}
// TODO: There is no test for FixedSizeList now, add it later
DataType::FixedSizeList(_, _) => {
DataType::FixedSizeList(field, _) => {
let list_array = as_fixed_size_list_array(array)?;
let nested_array = list_array.value(index);
// Produces a single element `ListArray` with the value at `index`.
let list_size = nested_array.len();
let arr =
Arc::new(array_into_fixed_size_list_array(nested_array, list_size));
let arr = Arc::new(array_into_fixed_size_list_array_with_field_name(
nested_array,
list_size,
field.name(),
));

ScalarValue::FixedSizeList(arr)
}
Expand Down Expand Up @@ -5970,6 +5981,51 @@ mod tests {
ScalarValue::from("larger than 12 bytes string"),
DataType::Utf8View,
);
check_scalar_cast(
{
let element_field =
Arc::new(Field::new("element", DataType::Int32, true));

let mut builder =
ListBuilder::new(Int32Builder::new()).with_field(element_field);
builder.append_value([Some(1)]);
builder.append(true);

ScalarValue::List(Arc::new(builder.finish()))
},
DataType::List(Arc::new(Field::new("element", DataType::Int64, true))),
);
check_scalar_cast(
{
let element_field =
Arc::new(Field::new("element", DataType::Int32, true));

let mut builder = FixedSizeListBuilder::new(Int32Builder::new(), 1)
.with_field(element_field);
builder.values().append_value(1);
builder.append(true);

ScalarValue::FixedSizeList(Arc::new(builder.finish()))
},
DataType::FixedSizeList(
Arc::new(Field::new("element", DataType::Int64, true)),
1,
),
);
check_scalar_cast(
{
let element_field =
Arc::new(Field::new("element", DataType::Int32, true));

let mut builder =
LargeListBuilder::new(Int32Builder::new()).with_field(element_field);
builder.append_value([Some(1)]);
builder.append(true);

ScalarValue::LargeList(Arc::new(builder.finish()))
},
DataType::LargeList(Arc::new(Field::new("element", DataType::Int64, true))),
);
}

// mimics how casting work on scalar values by `casting` `scalar` to `desired_type`
Expand Down
41 changes: 41 additions & 0 deletions datafusion/common/src/utils/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -342,6 +342,20 @@ pub fn array_into_list_array(arr: ArrayRef, nullable: bool) -> ListArray {
)
}

pub fn array_into_list_array_with_field_name(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

So is the idea that we will backport this to 43.0.0 and then consolidate the other functions on main? That makes sense to me if so

arr: ArrayRef,
nullable: bool,
field_name: &str,
) -> ListArray {
let offsets = OffsetBuffer::from_lengths([arr.len()]);
ListArray::new(
Arc::new(Field::new(field_name, arr.data_type().to_owned(), nullable)),
offsets,
arr,
None,
)
}

/// Wrap an array into a single element `LargeListArray`.
/// For example `[1, 2, 3]` would be converted into `[[1, 2, 3]]`
pub fn array_into_large_list_array(arr: ArrayRef) -> LargeListArray {
Expand All @@ -354,6 +368,19 @@ pub fn array_into_large_list_array(arr: ArrayRef) -> LargeListArray {
)
}

pub fn array_into_large_list_array_with_field_name(
arr: ArrayRef,
field_name: &str,
) -> LargeListArray {
let offsets = OffsetBuffer::from_lengths([arr.len()]);
LargeListArray::new(
Arc::new(Field::new(field_name, arr.data_type().to_owned(), true)),
offsets,
arr,
None,
)
}

pub fn array_into_fixed_size_list_array(
arr: ArrayRef,
list_size: usize,
Expand All @@ -367,6 +394,20 @@ pub fn array_into_fixed_size_list_array(
)
}

pub fn array_into_fixed_size_list_array_with_field_name(
arr: ArrayRef,
list_size: usize,
field_name: &str,
) -> FixedSizeListArray {
let list_size = list_size as i32;
FixedSizeListArray::new(
Arc::new(Field::new(field_name, arr.data_type().to_owned(), true)),
list_size,
arr,
None,
)
}

/// Wrap arrays into a single element `ListArray`.
///
/// Example:
Expand Down