@@ -12,17 +12,27 @@ use parquet_variant_compute::GetOptions;
1212use parquet_variant_compute:: VariantArray as ArrowVariantArray ;
1313use vortex_array:: ArrayRef ;
1414use vortex_array:: ExecutionCtx ;
15+ use vortex_array:: IntoArray ;
16+ use vortex_array:: arrays:: VariantArray ;
1517use vortex_array:: arrays:: scalar_fn:: ExactScalarFn ;
1618use vortex_array:: arrays:: scalar_fn:: ScalarFnArrayView ;
19+ use vortex_array:: arrow:: FromArrowArray ;
20+ use vortex_array:: dtype:: DType ;
1721use vortex_array:: dtype:: FieldName ;
22+ use vortex_array:: dtype:: Nullability ;
1823use vortex_array:: kernel:: ExecuteParentKernel ;
1924use vortex_array:: scalar_fn:: fns:: variant_get:: VariantGet ;
25+ use vortex_array:: validity:: Validity ;
26+ use vortex_buffer:: BitBuffer ;
2027use vortex_error:: VortexResult ;
2128use vortex_error:: vortex_err;
2229
2330use crate :: ParquetVariant ;
2431use crate :: array:: ParquetVariantArray ;
2532
33+ #[ cfg( test) ]
34+ mod tests;
35+
2636#[ derive( Debug ) ]
2737pub ( 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