klion26 commented on code in PR #8516:
URL: https://github.com/apache/arrow-rs/pull/8516#discussion_r2398906273
##########
parquet-variant-compute/src/variant_to_arrow.rs:
##########
@@ -293,55 +324,90 @@ fn get_type_name<T: ArrowPrimitiveType>() -> &'static str
{
"arrow_array::types::Float32Type" => "Float32",
"arrow_array::types::Float64Type" => "Float64",
"arrow_array::types::Float16Type" => "Float16",
+ "arrow_array::types::TimestampMicrosecondType" =>
"Timestamp(Microsecond)",
+ "arrow_array::types::TimestampNanosecondType" =>
"Timestamp(Nanosecond)",
+ "arrow_array::types::Date32Type" => "Date32",
_ => "Unknown",
}
}
-/// Builder for converting variant values to primitive values
-pub(crate) struct VariantToPrimitiveArrowRowBuilder<'a, T:
PrimitiveFromVariant> {
- builder: arrow::array::PrimitiveBuilder<T>,
- cast_options: &'a CastOptions<'a>,
-}
-
-impl<'a, T: PrimitiveFromVariant> VariantToPrimitiveArrowRowBuilder<'a, T> {
- fn new(cast_options: &'a CastOptions<'a>, capacity: usize) -> Self {
- Self {
- builder: PrimitiveBuilder::<T>::with_capacity(capacity),
- cast_options,
+macro_rules! define_variant_to_primitive_builder {
+ (struct $name:ident<$lifetime:lifetime $(, $generic:ident: $bound:path )?>
+ |$array_param:ident $(, $field:ident: $field_type:ty)?| ->
$builder_name:ident $(< $array_type:ty >)? { $init_expr: expr },
+ |$value: ident| $value_transform:expr,
+ type_name: $type_name:expr) => {
+ pub(crate) struct $name<$lifetime $(, $generic : $bound )?>
+ {
+ builder: $builder_name $(<$array_type>)?,
+ cast_options: &$lifetime CastOptions<$lifetime>,
}
- }
-}
-impl<'a, T: PrimitiveFromVariant> VariantToPrimitiveArrowRowBuilder<'a, T> {
- fn append_null(&mut self) -> Result<()> {
- self.builder.append_null();
- Ok(())
- }
+ impl<$lifetime $(, $generic: $bound+ )?> $name<$lifetime $(, $generic
)?> {
+ fn new(
+ cast_options: &$lifetime CastOptions<$lifetime>,
+ $array_param: usize
+ // add this so that $init_expr can use it
+ $(, $field: $field_type)?) -> Self {
+ Self {
+ builder: $init_expr,
+ cast_options,
+ }
+ }
- fn append_value(&mut self, value: &Variant<'_, '_>) -> Result<bool> {
- if let Some(v) = T::from_variant(value) {
- self.builder.append_value(v);
- Ok(true)
- } else {
- if !self.cast_options.safe {
- // Unsafe casting: return error on conversion failure
- return Err(ArrowError::CastError(format!(
- "Failed to extract primitive of type {} from variant {:?}
at path VariantPath([])",
- get_type_name::<T>(),
- value
- )));
+ fn append_null(&mut self) -> Result<()> {
+ self.builder.append_null();
+ Ok(())
}
- // Safe casting: append null on conversion failure
- self.builder.append_null();
- Ok(false)
- }
- }
- fn finish(mut self) -> Result<ArrayRef> {
- Ok(Arc::new(self.builder.finish()))
+ fn append_value(&mut self, $value: &Variant<'_, '_>) ->
Result<bool> {
+ if let Some(v) = $value_transform {
+ self.builder.append_value(v);
+ Ok(true)
+ } else {
+ if !self.cast_options.safe {
+ // Unsafe casting: return error on conversion failure
+ return Err(ArrowError::CastError(format!(
+ "Failed to extract primitive of type {} from
variant {:?} at path VariantPath([])",
+ $type_name,
+ $value
+ )));
+ }
+ // Safe casting: append null on conversion failure
+ self.builder.append_null();
+ Ok(false)
+ }
+ }
+
+ fn finish(mut self) -> Result<ArrayRef> {
+ Ok(Arc::new(self.builder.finish()))
+ }
+ }
}
}
+define_variant_to_primitive_builder!(
+ struct VariantToBooleanArrowRowBuilder<'a>
+ |capacity| -> BooleanBuilder { BooleanBuilder::with_capacity(capacity) },
+ |value| value.as_boolean(),
+ type_name: "Boolean"
+);
+
+define_variant_to_primitive_builder!(
+ struct VariantToPrimitiveArrowRowBuilder<'a, T:PrimitiveFromVariant>
+ |capacity| -> PrimitiveBuilder<T> {
PrimitiveBuilder::<T>::with_capacity(capacity) },
+ |value| T::from_variant(value),
+ type_name: get_type_name::<T>()
+);
+
+define_variant_to_primitive_builder!(
Review Comment:
The type of `T` needs to be `ArrowTimestampType` because the
`with_timezone_opt` function only exists for `ArrowTimestampType`
##########
parquet-variant-compute/src/variant_to_arrow.rs:
##########
@@ -293,21 +337,77 @@ fn get_type_name<T: ArrowPrimitiveType>() -> &'static str
{
"arrow_array::types::Float32Type" => "Float32",
"arrow_array::types::Float64Type" => "Float64",
"arrow_array::types::Float16Type" => "Float16",
+ "arrow_array::types::TimestampMicrosecondType" =>
"Timestamp(Microsecond)",
+ "arrow_array::types::TimestampNanosecondType" =>
"Timestamp(Nanosecond)",
_ => "Unknown",
Review Comment:
Created an issue(#8538) to track the cleanup logic
--
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]
