Samyak2 commented on code in PR #7919:
URL: https://github.com/apache/arrow-rs/pull/7919#discussion_r2203142467
##########
parquet-variant-compute/src/variant_get.rs:
##########
@@ -0,0 +1,265 @@
+use std::sync::Arc;
+
+use arrow::{
+ array::{
+ Array, ArrayRef, ArrowPrimitiveType, BinaryArray, PrimitiveArray,
PrimitiveBuilder,
+ StructArray,
+ },
+ compute::CastOptions,
+ datatypes::UInt64Type,
+ error::Result,
+};
+use arrow_schema::{ArrowError, DataType, Field};
+use parquet_variant::Variant;
+
+use crate::utils::variant_from_struct_array;
+
+/// Returns an array with the specified path extracted from the variant values.
+pub fn variant_get(input: &ArrayRef, options: GetOptions) -> Result<ArrayRef> {
+ let (struct_array, metadata_array, value_array) =
variant_from_struct_array(input)?;
+
+ // TODO: can we use OffsetBuffer and NullBuffer here instead?
+ // I couldn't find a way to set individual indices here, so went
with vecs.
+ let mut offsets = vec![0; struct_array.len()];
+ let mut nulls = if let Some(struct_nulls) = struct_array.nulls() {
+ struct_nulls.iter().collect()
+ } else {
+ vec![true; struct_array.len()]
+ };
+
+ for path in options
+ .path
+ .0
+ .iter()
+ .take(options.path.0.len().saturating_sub(1))
+ {
+ match path {
+ VariantPathElement::Field { name } => {
+ go_to_object_field(
+ struct_array,
+ metadata_array,
+ value_array,
+ name,
+ &mut offsets,
+ &mut nulls,
+ )?;
+ }
+ VariantPathElement::Index { offset } => {
+ go_to_array_index(
+ struct_array,
+ metadata_array,
+ value_array,
+ *offset,
+ &mut offsets,
+ &mut nulls,
+ )?;
+ }
+ }
+ }
+
+ let as_type = options.as_type.ok_or_else(|| {
Review Comment:
That makes sense. I'll focus on returning a variant for this PR
--
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: github-unsubscr...@arrow.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
