Skip to content

Commit 90e660d

Browse files
committed
address comments
1 parent a126b92 commit 90e660d

File tree

2 files changed

+57
-32
lines changed

2 files changed

+57
-32
lines changed

parquet-variant-compute/src/variant_get.rs

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -319,6 +319,7 @@ mod test {
319319
use arrow::compute::CastOptions;
320320
use arrow::datatypes::DataType::{Int16, Int32, Int64};
321321
use arrow::datatypes::i256;
322+
use arrow::util::display::FormatOptions;
322323
use arrow_schema::DataType::{Boolean, Float32, Float64, Int8};
323324
use arrow_schema::{DataType, Field, FieldRef, Fields, IntervalUnit, TimeUnit};
324325
use chrono::DateTime;
@@ -1043,11 +1044,25 @@ mod test {
10431044
Int32Array::from(vec![Some(32), Some(64), Some(48)])
10441045
});
10451046

1047+
// We append null values if type miss match happens in safe mode
1048+
perfectly_shredded_to_arrow_primitive_test!(
1049+
get_variant_perfectly_shredded_null_with_type_missmatch_in_safe_mode,
1050+
DataType::Null,
1051+
perfectly_shredded_null_variant_array_with_int,
1052+
arrow::array::NullArray::new(3)
1053+
);
1054+
1055+
// We'll return an error if type miss match happens in strict mode
10461056
#[test]
1047-
fn get_variant_perfectly_shredded_null_as_null_with_type_missmatch() {
1057+
fn get_variant_perfectly_shredded_null_as_null_with_type_missmatch_in_strict_mode() {
10481058
let array = perfectly_shredded_null_variant_array_with_int();
10491059
let field = Field::new("typed_value", DataType::Null, true);
1050-
let options = GetOptions::new().with_as_type(Some(FieldRef::from(field)));
1060+
let options = GetOptions::new()
1061+
.with_as_type(Some(FieldRef::from(field)))
1062+
.with_cast_options(CastOptions {
1063+
safe: false,
1064+
format_options: FormatOptions::default(),
1065+
});
10511066

10521067
let result = variant_get(&array, options);
10531068

parquet-variant-compute/src/variant_to_arrow.rs

Lines changed: 40 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ use std::sync::Arc;
3737
/// `VariantToArrowRowBuilder` (below) and `VariantToShreddedPrimitiveVariantRowBuilder` (in
3838
/// `shred_variant.rs`).
3939
pub(crate) enum PrimitiveVariantToArrowRowBuilder<'a> {
40-
Null(VariantToNullArrowRowBuilder),
40+
Null(VariantToNullArrowRowBuilder<'a>),
4141
Boolean(VariantToBooleanArrowRowBuilder<'a>),
4242
Int8(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Int8Type>),
4343
Int16(VariantToPrimitiveArrowRowBuilder<'a, datatypes::Int16Type>),
@@ -232,7 +232,7 @@ pub(crate) fn make_primitive_variant_to_arrow_row_builder<'a>(
232232

233233
let builder =
234234
match data_type {
235-
DataType::Null => Null(VariantToNullArrowRowBuilder::new(capacity)),
235+
DataType::Null => Null(VariantToNullArrowRowBuilder::new(cast_options)),
236236
DataType::Boolean => {
237237
Boolean(VariantToBooleanArrowRowBuilder::new(cast_options, capacity))
238238
}
@@ -696,8 +696,43 @@ impl VariantToBinaryVariantArrowRowBuilder {
696696
}
697697
}
698698

699-
pub(crate) struct VariantToNullArrowRowBuilder {
700-
capacity: usize,
699+
pub(crate) struct VariantToNullArrowRowBuilder<'a> {
700+
cast_option: &'a CastOptions<'a>,
701+
item_count: usize,
702+
}
703+
704+
705+
impl<'a> VariantToNullArrowRowBuilder<'a> {
706+
fn new(cast_option: &'a CastOptions<'a>) -> Self {
707+
Self {
708+
cast_option,
709+
item_count: 0,
710+
}
711+
}
712+
713+
fn append_value(&mut self, value: &Variant<'_, '_>) -> Result<bool> {
714+
if value.as_null().is_some() {
715+
self.item_count += 1;
716+
Ok(true)
717+
} else if self.cast_option.safe {
718+
self.append_null()?;
719+
Ok(false)
720+
} else {
721+
Err(ArrowError::CastError(format!(
722+
"Failed to extract Null from variant {:?} at path VariantPath([])",
723+
value
724+
)))
725+
}
726+
}
727+
728+
fn append_null(&mut self) -> Result<()> {
729+
self.item_count += 1;
730+
Ok(())
731+
}
732+
733+
fn finish(self) -> Result<ArrayRef> {
734+
Ok(Arc::new(NullArray::new(self.item_count)))
735+
}
701736
}
702737

703738

@@ -755,29 +790,4 @@ mod tests {
755790
}
756791
}
757792
}
758-
}
759-
impl VariantToNullArrowRowBuilder {
760-
fn new(capacity: usize) -> Self {
761-
Self { capacity }
762-
}
763-
764-
fn append_value(&mut self, value: &Variant<'_, '_>) -> Result<bool> {
765-
if value.as_null().is_some() {
766-
Ok(true)
767-
} else {
768-
// Null type only accepts nulls
769-
Err(ArrowError::CastError(format!(
770-
"Failed to extract Null from variant {:?} at path VariantPath([])",
771-
value
772-
)))
773-
}
774-
}
775-
776-
fn append_null(&mut self) -> Result<()> {
777-
Ok(())
778-
}
779-
780-
fn finish(self) -> Result<ArrayRef> {
781-
Ok(Arc::new(NullArray::new(self.capacity)))
782-
}
783-
}
793+
}

0 commit comments

Comments
 (0)