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.

Sign up for GitHub

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

feat: add support for converting LargeUtf8/LargeString in try_cast_to #71

Merged
merged 2 commits into from
Oct 25, 2024
Merged

feat: add support for converting LargeUtf8/LargeString in try_cast_to #71

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9ae73a2
feat: add support for converting LargeUtf8/LargeString in try_cast_to…
gengteng Oct 15, 2024
5a833a7
test: use assert_batches_eq
gengteng Oct 23, 2024
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
13 changes: 7 additions & 6 deletions datafusion-federation/src/schema_cast/lists_cast.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use arrow_json::ReaderBuilder;
use datafusion::arrow::array::{GenericStringArray, OffsetSizeTrait};
use datafusion::arrow::{
array::{
Array, ArrayRef, BooleanArray, BooleanBuilder, FixedSizeListBuilder, Float32Array,
Expand Down Expand Up @@ -193,13 +194,13 @@ macro_rules! cast_string_to_fixed_size_list_array {
}};
}

pub(crate) fn cast_string_to_list(
pub(crate) fn cast_string_to_list<StringOffsetSize: OffsetSizeTrait>(
array: &dyn Array,
list_item_field: &FieldRef,
) -> Result<ArrayRef, ArrowError> {
let string_array = array
.as_any()
.downcast_ref::<StringArray>()
.downcast_ref::<GenericStringArray<StringOffsetSize>>()
.ok_or_else(|| {
ArrowError::CastError(
"Failed to decode value: unable to downcast to StringArray".to_string(),
Expand Down Expand Up @@ -297,13 +298,13 @@ pub(crate) fn cast_string_to_list(
}
}

pub(crate) fn cast_string_to_large_list(
pub(crate) fn cast_string_to_large_list<StringOffsetSize: OffsetSizeTrait>(
array: &dyn Array,
list_item_field: &FieldRef,
) -> Result<ArrayRef, ArrowError> {
let string_array = array
.as_any()
.downcast_ref::<StringArray>()
.downcast_ref::<GenericStringArray<StringOffsetSize>>()
.ok_or_else(|| {
ArrowError::CastError(
"Failed to decode value: unable to downcast to StringArray".to_string(),
Expand Down Expand Up @@ -401,14 +402,14 @@ pub(crate) fn cast_string_to_large_list(
}
}

pub(crate) fn cast_string_to_fixed_size_list(
pub(crate) fn cast_string_to_fixed_size_list<StringOffsetSize: OffsetSizeTrait>(
array: &dyn Array,
list_item_field: &FieldRef,
value_length: i32,
) -> Result<ArrayRef, ArrowError> {
let string_array = array
.as_any()
.downcast_ref::<StringArray>()
.downcast_ref::<GenericStringArray<StringOffsetSize>>()
.ok_or_else(|| {
ArrowError::CastError(
"Failed to decode value: unable to downcast to StringArray".to_string(),
Expand Down
75 changes: 68 additions & 7 deletions datafusion-federation/src/schema_cast/record_convert.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -68,25 +68,47 @@ pub fn try_cast_to(record_batch: RecordBatch, expected_schema: SchemaRef) -> Res

match (record_batch_col.data_type(), expected_field.data_type()) {
(DataType::Utf8, DataType::List(item_type)) => {
cast_string_to_list(&Arc::clone(record_batch_col), item_type)
cast_string_to_list::<i32>(&Arc::clone(record_batch_col), item_type)
.map_err(|e| Error::UnableToConvertRecordBatch { source: e })
}
(DataType::Utf8, DataType::LargeList(item_type)) => {
cast_string_to_large_list(&Arc::clone(record_batch_col), item_type)
cast_string_to_large_list::<i32>(&Arc::clone(record_batch_col), item_type)
.map_err(|e| Error::UnableToConvertRecordBatch { source: e })
}
(DataType::Utf8, DataType::FixedSizeList(item_type, value_length)) => {
cast_string_to_fixed_size_list(
cast_string_to_fixed_size_list::<i32>(
&Arc::clone(record_batch_col),
item_type,
*value_length,
)
.map_err(|e| Error::UnableToConvertRecordBatch { source: e })
}
(DataType::Utf8, DataType::Struct(_)) => {
cast_string_to_struct(&Arc::clone(record_batch_col), expected_field.clone())
(DataType::Utf8, DataType::Struct(_)) => cast_string_to_struct::<i32>(
&Arc::clone(record_batch_col),
expected_field.clone(),
)
.map_err(|e| Error::UnableToConvertRecordBatch { source: e }),
(DataType::LargeUtf8, DataType::List(item_type)) => {
cast_string_to_list::<i64>(&Arc::clone(record_batch_col), item_type)
.map_err(|e| Error::UnableToConvertRecordBatch { source: e })
}
(DataType::LargeUtf8, DataType::LargeList(item_type)) => {
cast_string_to_large_list::<i64>(&Arc::clone(record_batch_col), item_type)
.map_err(|e| Error::UnableToConvertRecordBatch { source: e })
}
(DataType::LargeUtf8, DataType::FixedSizeList(item_type, value_length)) => {
cast_string_to_fixed_size_list::<i64>(
&Arc::clone(record_batch_col),
item_type,
*value_length,
)
.map_err(|e| Error::UnableToConvertRecordBatch { source: e })
}
(DataType::LargeUtf8, DataType::Struct(_)) => cast_string_to_struct::<i64>(
&Arc::clone(record_batch_col),
expected_field.clone(),
)
.map_err(|e| Error::UnableToConvertRecordBatch { source: e }),
(
DataType::Interval(IntervalUnit::MonthDayNano),
DataType::Interval(IntervalUnit::YearMonth),
Expand All @@ -109,13 +131,13 @@ pub fn try_cast_to(record_batch: RecordBatch, expected_schema: SchemaRef) -> Res

#[cfg(test)]
mod test {
use super::*;
use datafusion::arrow::array::LargeStringArray;
use datafusion::arrow::{
array::{Int32Array, StringArray},
datatypes::{DataType, Field, Schema, TimeUnit},
};

use super::*;

fn schema() -> SchemaRef {
Arc::new(Schema::new(vec![
Field::new("a", DataType::Int32, false),
Expand Down Expand Up @@ -153,4 +175,43 @@ mod test {
let result = try_cast_to(batch_input(), to_schema()).expect("converted");
assert_eq!(3, result.num_rows());
}

fn large_string_schema() -> SchemaRef {
gengteng marked this conversation as resolved.
Show resolved Hide resolved
Arc::new(Schema::new(vec![
Field::new("a", DataType::Int32, false),
Field::new("b", DataType::LargeUtf8, false),
Field::new("c", DataType::LargeUtf8, false),
]))
}

fn large_string_to_schema() -> SchemaRef {
Arc::new(Schema::new(vec![
Field::new("a", DataType::Int64, false),
Field::new("b", DataType::LargeUtf8, false),
Field::new("c", DataType::Timestamp(TimeUnit::Microsecond, None), false),
]))
}

fn large_string_batch_input() -> RecordBatch {
RecordBatch::try_new(
large_string_schema(),
vec![
Arc::new(Int32Array::from(vec![1, 2, 3])),
Arc::new(LargeStringArray::from(vec!["foo", "bar", "baz"])),
Arc::new(LargeStringArray::from(vec![
"2024-01-13 03:18:09.000000",
"2024-01-13 03:18:09",
"2024-01-13 03:18:09.000",
])),
],
)
.expect("record batch should not panic")
}

#[test]
fn test_large_string_to_timestamp_conversion() {
let result =
try_cast_to(large_string_batch_input(), large_string_to_schema()).expect("converted");
assert_eq!(3, result.num_rows());
gengteng marked this conversation as resolved.
Show resolved Hide resolved
}
}
7 changes: 4 additions & 3 deletions datafusion-federation/src/schema_cast/struct_cast.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,20 +1,21 @@
use arrow_json::ReaderBuilder;
use datafusion::arrow::array::{GenericStringArray, OffsetSizeTrait};
use datafusion::arrow::{
array::{Array, ArrayRef, StringArray},
array::{Array, ArrayRef},
datatypes::Field,
error::ArrowError,
};
use std::sync::Arc;

pub type Result<T, E = crate::schema_cast::record_convert::Error> = std::result::Result<T, E>;

pub(crate) fn cast_string_to_struct(
pub(crate) fn cast_string_to_struct<StringOffsetSize: OffsetSizeTrait>(
array: &dyn Array,
struct_field: Arc<Field>,
) -> Result<ArrayRef, ArrowError> {
let string_array = array
.as_any()
.downcast_ref::<StringArray>()
.downcast_ref::<GenericStringArray<StringOffsetSize>>()
.ok_or_else(|| ArrowError::CastError("Failed to downcast to StringArray".to_string()))?;

let mut decoder = ReaderBuilder::new_with_field(struct_field)
Expand Down