Skip to content

Commit 38345ea

Browse files
committed
stuff
Signed-off-by: Adam Gutglick <adam@spiraldb.com>
1 parent 44f1e3b commit 38345ea

4 files changed

Lines changed: 395 additions & 168 deletions

File tree

encodings/parquet-variant/src/lib.rs

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,6 @@ mod kernel;
2929
mod operations;
3030
mod validity;
3131
mod variant_get;
32-
#[cfg(test)]
33-
mod variant_get_tests;
3432
mod vtable;
3533

3634
pub use array::ParquetVariantArray;

encodings/parquet-variant/src/variant_get.rs renamed to encodings/parquet-variant/src/variant_get/mod.rs

Lines changed: 44 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,17 +12,27 @@ use parquet_variant_compute::GetOptions;
1212
use parquet_variant_compute::VariantArray as ArrowVariantArray;
1313
use vortex_array::ArrayRef;
1414
use vortex_array::ExecutionCtx;
15+
use vortex_array::IntoArray;
16+
use vortex_array::arrays::VariantArray;
1517
use vortex_array::arrays::scalar_fn::ExactScalarFn;
1618
use vortex_array::arrays::scalar_fn::ScalarFnArrayView;
19+
use vortex_array::arrow::FromArrowArray;
20+
use vortex_array::dtype::DType;
1721
use vortex_array::dtype::FieldName;
22+
use vortex_array::dtype::Nullability;
1823
use vortex_array::kernel::ExecuteParentKernel;
1924
use vortex_array::scalar_fn::fns::variant_get::VariantGet;
25+
use vortex_array::validity::Validity;
26+
use vortex_buffer::BitBuffer;
2027
use vortex_error::VortexResult;
2128
use vortex_error::vortex_err;
2229

2330
use crate::ParquetVariant;
2431
use crate::array::ParquetVariantArray;
2532

33+
#[cfg(test)]
34+
mod tests;
35+
2636
#[derive(Debug)]
2737
pub(crate) struct VariantGetExecuteParent;
2838

@@ -61,7 +71,7 @@ fn variant_get_impl(
6171
let arrow_result = parquet_variant_compute::variant_get(&inner, options)
6272
.map_err(|e| vortex_err!("variant_get failed: {e}"))?;
6373

64-
// Convert back to Vortex
74+
// Convert back to Vortex.
6575
let result_variant = ArrowVariantArray::try_new(
6676
arrow_result
6777
.as_any()
@@ -70,5 +80,37 @@ fn variant_get_impl(
7080
)
7181
.map_err(|e| vortex_err!("failed to create VariantArray from result: {e}"))?;
7282

73-
ParquetVariantArray::from_arrow_variant(&result_variant)
83+
// Ensure the result is always nullable (matching variant_get's return_dtype).
84+
// Arrow may return a non-nullable result when no nulls are present.
85+
let validity = result_variant
86+
.nulls()
87+
.map(|nulls| {
88+
if nulls.null_count() == nulls.len() {
89+
Validity::AllInvalid
90+
} else {
91+
Validity::from(BitBuffer::from(nulls.inner().clone()))
92+
}
93+
})
94+
.unwrap_or(Validity::AllValid);
95+
96+
let metadata = ArrayRef::from_arrow(
97+
result_variant.metadata_field() as &dyn arrow_array::Array,
98+
false,
99+
)?;
100+
let value = result_variant
101+
.value_field()
102+
.map(|v| ArrayRef::from_arrow(v as &dyn arrow_array::Array, true))
103+
.transpose()?;
104+
let typed_value = result_variant
105+
.typed_value_field()
106+
.map(|tv| ArrayRef::from_arrow(tv.as_ref(), true))
107+
.transpose()?;
108+
109+
let pv = ParquetVariantArray::try_new(validity, metadata, value, typed_value)?;
110+
debug_assert_eq!(
111+
pv.dtype,
112+
DType::Variant(Nullability::Nullable),
113+
"variant_get result must be nullable"
114+
);
115+
Ok(VariantArray::new(pv.into_array()).into_array())
74116
}

0 commit comments

Comments
 (0)