liamzwbao commented on code in PR #8235:
URL: https://github.com/apache/arrow-rs/pull/8235#discussion_r2305694210
##########
parquet-variant-compute/src/cast_to_variant.rs:
##########
@@ -412,140 +301,230 @@ pub fn cast_to_variant(input: &dyn Array) ->
Result<VariantArray, ArrowError> {
)));
}
},
- DataType::Dictionary(_, _) => {
- convert_dictionary_encoded(input, &mut builder)?;
+ dt => {
+ return Err(ArrowError::CastError(format!(
+ "Unsupported data type for casting to Variant: {dt:?}",
+ )));
}
+ };
+ Ok(builder.build())
+}
- DataType::Map(field, _) => match field.data_type() {
- DataType::Struct(_) => {
- let map_array = input.as_map();
- let keys = cast(map_array.keys(), &DataType::Utf8)?;
- let key_strings = keys.as_string::<i32>();
- let values = cast_to_variant(map_array.values())?;
- let offsets = map_array.offsets();
-
- let mut start_offset = offsets[0];
- for end_offset in offsets.iter().skip(1) {
- if start_offset >= *end_offset {
- builder.append_null();
- continue;
- }
-
- let length = (end_offset - start_offset) as usize;
+// TODO do we need a cast_with_options to allow specifying conversion behavior,
+// e.g. how to handle overflows, whether to convert to Variant::Null or return
+// an error, etc. ?
- let mut variant_builder = VariantBuilder::new();
- let mut object_builder = variant_builder.new_object();
+/// Convert timestamp arrays to native datetimes
+fn convert_timestamp(
+ time_unit: &TimeUnit,
+ time_zone: &Option<Arc<str>>,
+ input: &dyn Array,
+ builder: &mut VariantArrayBuilder,
+) {
+ let native_datetimes: Vec<Option<NaiveDateTime>> = match time_unit {
+ arrow_schema::TimeUnit::Second => {
+ let ts_array = input
+ .as_any()
+ .downcast_ref::<TimestampSecondArray>()
+ .expect("Array is not TimestampSecondArray");
- for i in start_offset..*end_offset {
- let value = values.value(i as usize);
- object_builder.insert(key_strings.value(i as usize),
value);
- }
- object_builder.finish();
- let (metadata, value) = variant_builder.finish();
- let variant = Variant::try_new(&metadata, &value)?;
+ ts_array
+ .iter()
+ .map(|x| x.map(|y| timestamp_s_to_datetime(y).unwrap()))
+ .collect()
+ }
+ arrow_schema::TimeUnit::Millisecond => {
+ let ts_array = input
+ .as_any()
+ .downcast_ref::<TimestampMillisecondArray>()
+ .expect("Array is not TimestampMillisecondArray");
- builder.append_variant(variant);
+ ts_array
+ .iter()
+ .map(|x| x.map(|y| timestamp_ms_to_datetime(y).unwrap()))
+ .collect()
+ }
+ arrow_schema::TimeUnit::Microsecond => {
+ let ts_array = input
+ .as_any()
+ .downcast_ref::<TimestampMicrosecondArray>()
+ .expect("Array is not TimestampMicrosecondArray");
+ ts_array
+ .iter()
+ .map(|x| x.map(|y| timestamp_us_to_datetime(y).unwrap()))
+ .collect()
+ }
+ arrow_schema::TimeUnit::Nanosecond => {
+ let ts_array = input
+ .as_any()
+ .downcast_ref::<TimestampNanosecondArray>()
+ .expect("Array is not TimestampNanosecondArray");
+ ts_array
+ .iter()
+ .map(|x| x.map(|y| timestamp_ns_to_datetime(y).unwrap()))
+ .collect()
+ }
+ };
- start_offset += length as i32;
+ for x in native_datetimes {
+ match x {
+ Some(ndt) => {
+ if time_zone.is_none() {
+ builder.append_variant(ndt.into());
+ } else {
+ let utc_dt: DateTime<Utc> = Utc.from_utc_datetime(&ndt);
+ builder.append_variant(utc_dt.into());
}
}
- _ => {
- return Err(ArrowError::CastError(format!(
- "Unsupported map field type for casting to Variant:
{field:?}",
- )));
+ None => {
+ builder.append_null();
}
- },
- DataType::List(_) => {
- let list_array = input.as_list::<i32>();
- let values = list_array.values();
- let offsets = list_array.offsets();
-
- let first_offset = offsets.first().expect("There should be an
offset");
- let length = offsets.last().expect("There should be an offset") -
first_offset;
- let sliced_values = values.slice(*first_offset as usize, length as
usize);
-
- let values_variant_array =
cast_to_variant(sliced_values.as_ref())?;
- let new_offsets = OffsetBuffer::new(ScalarBuffer::from_iter(
- offsets.iter().map(|o| o - first_offset),
- ));
+ }
+ }
+}
- for i in 0..list_array.len() {
- if list_array.is_null(i) {
- builder.append_null();
- continue;
- }
+/// Generic function to convert list arrays (both List and LargeList) to
variant arrays
+fn convert_list<O: OffsetSizeTrait>(
+ input: &dyn Array,
+ builder: &mut VariantArrayBuilder,
+) -> Result<(), ArrowError> {
Review Comment:
The logic for `List` and `LargeList` is the same, refactor them into this
generic helper function
##########
parquet-variant-compute/src/cast_to_variant.rs:
##########
@@ -412,140 +301,230 @@ pub fn cast_to_variant(input: &dyn Array) ->
Result<VariantArray, ArrowError> {
)));
}
},
- DataType::Dictionary(_, _) => {
- convert_dictionary_encoded(input, &mut builder)?;
+ dt => {
+ return Err(ArrowError::CastError(format!(
+ "Unsupported data type for casting to Variant: {dt:?}",
+ )));
}
+ };
+ Ok(builder.build())
+}
- DataType::Map(field, _) => match field.data_type() {
- DataType::Struct(_) => {
- let map_array = input.as_map();
- let keys = cast(map_array.keys(), &DataType::Utf8)?;
- let key_strings = keys.as_string::<i32>();
- let values = cast_to_variant(map_array.values())?;
- let offsets = map_array.offsets();
-
- let mut start_offset = offsets[0];
- for end_offset in offsets.iter().skip(1) {
- if start_offset >= *end_offset {
- builder.append_null();
- continue;
- }
-
- let length = (end_offset - start_offset) as usize;
+// TODO do we need a cast_with_options to allow specifying conversion behavior,
+// e.g. how to handle overflows, whether to convert to Variant::Null or return
+// an error, etc. ?
- let mut variant_builder = VariantBuilder::new();
- let mut object_builder = variant_builder.new_object();
+/// Convert timestamp arrays to native datetimes
+fn convert_timestamp(
+ time_unit: &TimeUnit,
+ time_zone: &Option<Arc<str>>,
+ input: &dyn Array,
+ builder: &mut VariantArrayBuilder,
+) {
+ let native_datetimes: Vec<Option<NaiveDateTime>> = match time_unit {
+ arrow_schema::TimeUnit::Second => {
+ let ts_array = input
+ .as_any()
+ .downcast_ref::<TimestampSecondArray>()
+ .expect("Array is not TimestampSecondArray");
- for i in start_offset..*end_offset {
- let value = values.value(i as usize);
- object_builder.insert(key_strings.value(i as usize),
value);
- }
- object_builder.finish();
- let (metadata, value) = variant_builder.finish();
- let variant = Variant::try_new(&metadata, &value)?;
+ ts_array
+ .iter()
+ .map(|x| x.map(|y| timestamp_s_to_datetime(y).unwrap()))
+ .collect()
+ }
+ arrow_schema::TimeUnit::Millisecond => {
+ let ts_array = input
+ .as_any()
+ .downcast_ref::<TimestampMillisecondArray>()
+ .expect("Array is not TimestampMillisecondArray");
- builder.append_variant(variant);
+ ts_array
+ .iter()
+ .map(|x| x.map(|y| timestamp_ms_to_datetime(y).unwrap()))
+ .collect()
+ }
+ arrow_schema::TimeUnit::Microsecond => {
+ let ts_array = input
+ .as_any()
+ .downcast_ref::<TimestampMicrosecondArray>()
+ .expect("Array is not TimestampMicrosecondArray");
+ ts_array
+ .iter()
+ .map(|x| x.map(|y| timestamp_us_to_datetime(y).unwrap()))
+ .collect()
+ }
+ arrow_schema::TimeUnit::Nanosecond => {
+ let ts_array = input
+ .as_any()
+ .downcast_ref::<TimestampNanosecondArray>()
+ .expect("Array is not TimestampNanosecondArray");
+ ts_array
+ .iter()
+ .map(|x| x.map(|y| timestamp_ns_to_datetime(y).unwrap()))
+ .collect()
+ }
+ };
- start_offset += length as i32;
+ for x in native_datetimes {
+ match x {
+ Some(ndt) => {
+ if time_zone.is_none() {
+ builder.append_variant(ndt.into());
+ } else {
+ let utc_dt: DateTime<Utc> = Utc.from_utc_datetime(&ndt);
+ builder.append_variant(utc_dt.into());
}
}
- _ => {
- return Err(ArrowError::CastError(format!(
- "Unsupported map field type for casting to Variant:
{field:?}",
- )));
+ None => {
+ builder.append_null();
}
- },
- DataType::List(_) => {
- let list_array = input.as_list::<i32>();
- let values = list_array.values();
- let offsets = list_array.offsets();
-
- let first_offset = offsets.first().expect("There should be an
offset");
- let length = offsets.last().expect("There should be an offset") -
first_offset;
- let sliced_values = values.slice(*first_offset as usize, length as
usize);
-
- let values_variant_array =
cast_to_variant(sliced_values.as_ref())?;
- let new_offsets = OffsetBuffer::new(ScalarBuffer::from_iter(
- offsets.iter().map(|o| o - first_offset),
- ));
+ }
+ }
+}
- for i in 0..list_array.len() {
- if list_array.is_null(i) {
- builder.append_null();
- continue;
- }
+/// Generic function to convert list arrays (both List and LargeList) to
variant arrays
+fn convert_list<O: OffsetSizeTrait>(
+ input: &dyn Array,
+ builder: &mut VariantArrayBuilder,
+) -> Result<(), ArrowError> {
+ let list_array = input.as_list::<O>();
+ let values = list_array.values();
+ let offsets = list_array.offsets();
- let start = new_offsets[i] as usize;
- let end = new_offsets[i + 1] as usize;
+ let first_offset = *offsets.first().expect("There should be an offset");
+ let length = *offsets.last().expect("There should be an offset") -
first_offset;
+ let sliced_values = values.slice(first_offset.as_usize(),
length.as_usize());
- // Start building the inner VariantList
- let mut variant_builder = VariantBuilder::new();
- let mut list_builder = variant_builder.new_list();
+ let values_variant_array = cast_to_variant(sliced_values.as_ref())?;
+ let new_offsets = OffsetBuffer::new(ScalarBuffer::from_iter(
+ offsets.iter().map(|o| *o - first_offset),
+ ));
- // Add all values from the slice
- for j in start..end {
- list_builder.append_value(values_variant_array.value(j));
- }
+ for i in 0..list_array.len() {
+ if list_array.is_null(i) {
+ builder.append_null();
+ continue;
+ }
- list_builder.finish();
+ let start = new_offsets[i].as_usize();
+ let end = new_offsets[i + 1].as_usize();
- let (metadata, value) = variant_builder.finish();
- let variant = Variant::new(&metadata, &value);
- let variant_list = variant.as_list().expect("Variant should be
list");
- builder.append_variant(Variant::List(variant_list.clone()))
- }
+ // Start building the inner VariantList
+ let mut variant_builder = VariantBuilder::new();
+ let mut list_builder = variant_builder.new_list();
+
+ // Add all values from the slice
+ for j in start..end {
+ list_builder.append_value(values_variant_array.value(j));
}
- DataType::LargeList(_) => {
- let large_list_array = input.as_list::<i64>();
- let values = large_list_array.values();
- let offsets = large_list_array.offsets();
-
- let first_offset = offsets.first().expect("There should be an
offset");
- let length = offsets.last().expect("There should be an offset") -
first_offset;
- let sliced_values = values.slice(*first_offset as usize, length as
usize);
-
- let values_variant_array =
cast_to_variant(sliced_values.as_ref())?;
- let new_offsets = OffsetBuffer::new(ScalarBuffer::from_iter(
- offsets.iter().map(|o| o - first_offset),
- ));
- for i in 0..large_list_array.len() {
- if large_list_array.is_null(i) {
- builder.append_null();
- continue;
- }
+ list_builder.finish();
- let start = new_offsets[i] as usize; // What if the system is
32bit and offset is > usize::MAX?
- let end = new_offsets[i + 1] as usize;
+ let (metadata, value) = variant_builder.finish();
+ let variant = Variant::new(&metadata, &value);
+ builder.append_variant(variant)
+ }
- // Start building the inner VariantList
- let mut variant_builder = VariantBuilder::new();
- let mut list_builder = variant_builder.new_list();
+ Ok(())
+}
+
+fn convert_struct(input: &dyn Array, builder: &mut VariantArrayBuilder) ->
Result<(), ArrowError> {
+ let struct_array = input.as_struct();
+
+ // Pre-convert all field arrays once for better performance
+ // This avoids converting the same field array multiple times
+ // Alternative approach: Use slicing per row: field_array.slice(i, 1)
+ // However, pre-conversion is more efficient for typical use cases
+ let field_variant_arrays: Result<Vec<_>, _> = struct_array
+ .columns()
+ .iter()
+ .map(|field_array| cast_to_variant(field_array.as_ref()))
+ .collect();
+ let field_variant_arrays = field_variant_arrays?;
+
+ // Cache column names to avoid repeated calls
+ let column_names = struct_array.column_names();
+
+ for i in 0..struct_array.len() {
+ if struct_array.is_null(i) {
+ builder.append_null();
+ continue;
+ }
+
+ // Create a VariantBuilder for this struct instance
+ let mut variant_builder = VariantBuilder::new();
+ let mut object_builder = variant_builder.new_object();
+
+ // Iterate through all fields in the struct
+ for (field_idx, field_name) in column_names.iter().enumerate() {
+ // Use pre-converted field variant arrays for better performance
+ // Check nulls directly from the pre-converted arrays instead of
accessing column again
+ if !field_variant_arrays[field_idx].is_null(i) {
+ let field_variant = field_variant_arrays[field_idx].value(i);
+ object_builder.insert(field_name, field_variant);
+ }
+ // Note: we skip null fields rather than inserting Variant::Null
+ // to match Arrow struct semantics where null fields are omitted
+ }
+
+ object_builder.finish();
+ let (metadata, value) = variant_builder.finish();
+ let variant = Variant::try_new(&metadata, &value)?;
+ builder.append_variant(variant);
+ }
+
+ Ok(())
+}
- // Add all values from the slice
- for j in start..end {
- list_builder.append_value(values_variant_array.value(j));
+fn convert_map(
Review Comment:
moved from the `Map` branch without any changes
##########
parquet-variant-compute/src/cast_to_variant.rs:
##########
@@ -2243,7 +2131,7 @@ mod tests {
}
#[test]
- fn test_cast_map_with_non_string_keys_to_variant_object() {
+ fn test_cast_to_variant_map_with_non_string_keys() {
Review Comment:
rename to match the naming convention in this file
##########
parquet-variant-compute/src/cast_to_variant.rs:
##########
@@ -40,80 +40,12 @@ use arrow::temporal_conversions::{
timestamp_ms_to_datetime, timestamp_ns_to_datetime,
timestamp_s_to_datetime,
timestamp_us_to_datetime,
};
-use arrow_schema::{ArrowError, DataType, TimeUnit, UnionFields};
+use arrow_schema::{ArrowError, DataType, FieldRef, TimeUnit, UnionFields};
use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, TimeZone, Utc};
use parquet_variant::{
Variant, VariantBuilder, VariantDecimal16, VariantDecimal4,
VariantDecimal8,
};
-fn convert_timestamp(
Review Comment:
Move this helper function below `cast_to_variant` so that the public
function appears first, making this file easier to read and reason about.
##########
parquet-variant-compute/src/cast_to_variant.rs:
##########
@@ -412,140 +301,230 @@ pub fn cast_to_variant(input: &dyn Array) ->
Result<VariantArray, ArrowError> {
)));
}
},
- DataType::Dictionary(_, _) => {
- convert_dictionary_encoded(input, &mut builder)?;
+ dt => {
+ return Err(ArrowError::CastError(format!(
+ "Unsupported data type for casting to Variant: {dt:?}",
+ )));
}
+ };
+ Ok(builder.build())
+}
- DataType::Map(field, _) => match field.data_type() {
- DataType::Struct(_) => {
- let map_array = input.as_map();
- let keys = cast(map_array.keys(), &DataType::Utf8)?;
- let key_strings = keys.as_string::<i32>();
- let values = cast_to_variant(map_array.values())?;
- let offsets = map_array.offsets();
-
- let mut start_offset = offsets[0];
- for end_offset in offsets.iter().skip(1) {
- if start_offset >= *end_offset {
- builder.append_null();
- continue;
- }
-
- let length = (end_offset - start_offset) as usize;
+// TODO do we need a cast_with_options to allow specifying conversion behavior,
+// e.g. how to handle overflows, whether to convert to Variant::Null or return
+// an error, etc. ?
- let mut variant_builder = VariantBuilder::new();
- let mut object_builder = variant_builder.new_object();
+/// Convert timestamp arrays to native datetimes
+fn convert_timestamp(
+ time_unit: &TimeUnit,
+ time_zone: &Option<Arc<str>>,
+ input: &dyn Array,
+ builder: &mut VariantArrayBuilder,
+) {
+ let native_datetimes: Vec<Option<NaiveDateTime>> = match time_unit {
+ arrow_schema::TimeUnit::Second => {
+ let ts_array = input
+ .as_any()
+ .downcast_ref::<TimestampSecondArray>()
+ .expect("Array is not TimestampSecondArray");
- for i in start_offset..*end_offset {
- let value = values.value(i as usize);
- object_builder.insert(key_strings.value(i as usize),
value);
- }
- object_builder.finish();
- let (metadata, value) = variant_builder.finish();
- let variant = Variant::try_new(&metadata, &value)?;
+ ts_array
+ .iter()
+ .map(|x| x.map(|y| timestamp_s_to_datetime(y).unwrap()))
+ .collect()
+ }
+ arrow_schema::TimeUnit::Millisecond => {
+ let ts_array = input
+ .as_any()
+ .downcast_ref::<TimestampMillisecondArray>()
+ .expect("Array is not TimestampMillisecondArray");
- builder.append_variant(variant);
+ ts_array
+ .iter()
+ .map(|x| x.map(|y| timestamp_ms_to_datetime(y).unwrap()))
+ .collect()
+ }
+ arrow_schema::TimeUnit::Microsecond => {
+ let ts_array = input
+ .as_any()
+ .downcast_ref::<TimestampMicrosecondArray>()
+ .expect("Array is not TimestampMicrosecondArray");
+ ts_array
+ .iter()
+ .map(|x| x.map(|y| timestamp_us_to_datetime(y).unwrap()))
+ .collect()
+ }
+ arrow_schema::TimeUnit::Nanosecond => {
+ let ts_array = input
+ .as_any()
+ .downcast_ref::<TimestampNanosecondArray>()
+ .expect("Array is not TimestampNanosecondArray");
+ ts_array
+ .iter()
+ .map(|x| x.map(|y| timestamp_ns_to_datetime(y).unwrap()))
+ .collect()
+ }
+ };
- start_offset += length as i32;
+ for x in native_datetimes {
+ match x {
+ Some(ndt) => {
+ if time_zone.is_none() {
+ builder.append_variant(ndt.into());
+ } else {
+ let utc_dt: DateTime<Utc> = Utc.from_utc_datetime(&ndt);
+ builder.append_variant(utc_dt.into());
}
}
- _ => {
- return Err(ArrowError::CastError(format!(
- "Unsupported map field type for casting to Variant:
{field:?}",
- )));
+ None => {
+ builder.append_null();
}
- },
- DataType::List(_) => {
- let list_array = input.as_list::<i32>();
- let values = list_array.values();
- let offsets = list_array.offsets();
-
- let first_offset = offsets.first().expect("There should be an
offset");
- let length = offsets.last().expect("There should be an offset") -
first_offset;
- let sliced_values = values.slice(*first_offset as usize, length as
usize);
-
- let values_variant_array =
cast_to_variant(sliced_values.as_ref())?;
- let new_offsets = OffsetBuffer::new(ScalarBuffer::from_iter(
- offsets.iter().map(|o| o - first_offset),
- ));
+ }
+ }
+}
- for i in 0..list_array.len() {
- if list_array.is_null(i) {
- builder.append_null();
- continue;
- }
+/// Generic function to convert list arrays (both List and LargeList) to
variant arrays
+fn convert_list<O: OffsetSizeTrait>(
+ input: &dyn Array,
+ builder: &mut VariantArrayBuilder,
+) -> Result<(), ArrowError> {
+ let list_array = input.as_list::<O>();
+ let values = list_array.values();
+ let offsets = list_array.offsets();
- let start = new_offsets[i] as usize;
- let end = new_offsets[i + 1] as usize;
+ let first_offset = *offsets.first().expect("There should be an offset");
+ let length = *offsets.last().expect("There should be an offset") -
first_offset;
+ let sliced_values = values.slice(first_offset.as_usize(),
length.as_usize());
- // Start building the inner VariantList
- let mut variant_builder = VariantBuilder::new();
- let mut list_builder = variant_builder.new_list();
+ let values_variant_array = cast_to_variant(sliced_values.as_ref())?;
+ let new_offsets = OffsetBuffer::new(ScalarBuffer::from_iter(
+ offsets.iter().map(|o| *o - first_offset),
+ ));
- // Add all values from the slice
- for j in start..end {
- list_builder.append_value(values_variant_array.value(j));
- }
+ for i in 0..list_array.len() {
+ if list_array.is_null(i) {
+ builder.append_null();
+ continue;
+ }
- list_builder.finish();
+ let start = new_offsets[i].as_usize();
+ let end = new_offsets[i + 1].as_usize();
- let (metadata, value) = variant_builder.finish();
- let variant = Variant::new(&metadata, &value);
- let variant_list = variant.as_list().expect("Variant should be
list");
- builder.append_variant(Variant::List(variant_list.clone()))
- }
+ // Start building the inner VariantList
+ let mut variant_builder = VariantBuilder::new();
+ let mut list_builder = variant_builder.new_list();
+
+ // Add all values from the slice
+ for j in start..end {
+ list_builder.append_value(values_variant_array.value(j));
}
- DataType::LargeList(_) => {
- let large_list_array = input.as_list::<i64>();
- let values = large_list_array.values();
- let offsets = large_list_array.offsets();
-
- let first_offset = offsets.first().expect("There should be an
offset");
- let length = offsets.last().expect("There should be an offset") -
first_offset;
- let sliced_values = values.slice(*first_offset as usize, length as
usize);
-
- let values_variant_array =
cast_to_variant(sliced_values.as_ref())?;
- let new_offsets = OffsetBuffer::new(ScalarBuffer::from_iter(
- offsets.iter().map(|o| o - first_offset),
- ));
- for i in 0..large_list_array.len() {
- if large_list_array.is_null(i) {
- builder.append_null();
- continue;
- }
+ list_builder.finish();
- let start = new_offsets[i] as usize; // What if the system is
32bit and offset is > usize::MAX?
- let end = new_offsets[i + 1] as usize;
+ let (metadata, value) = variant_builder.finish();
+ let variant = Variant::new(&metadata, &value);
+ builder.append_variant(variant)
+ }
- // Start building the inner VariantList
- let mut variant_builder = VariantBuilder::new();
- let mut list_builder = variant_builder.new_list();
+ Ok(())
+}
+
+fn convert_struct(input: &dyn Array, builder: &mut VariantArrayBuilder) ->
Result<(), ArrowError> {
Review Comment:
moved from the `Struct` branch without any changes
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
