scovich commented on code in PR #8366:
URL: https://github.com/apache/arrow-rs/pull/8366#discussion_r2355205126
##########
parquet-variant-compute/src/variant_to_arrow.rs:
##########
@@ -26,28 +26,34 @@ use crate::{VariantArray, VariantValueArrayBuilder};
use std::sync::Arc;
-/// Builder for converting variant values into strongly typed Arrow arrays.
-///
-/// Useful for variant_get kernels that need to extract specific paths from
variant values, possibly
-/// with casting of leaf values to specific types.
-pub(crate) enum VariantToArrowRowBuilder<'a> {
- // Direct builders (no path extraction)
+/// Builder for converting variant values to primitive Arrow arrays. It is
used by both
+/// `VariantToArrowRowBuilder` (below) and
`VariantToShreddedPrimitiveVariantRowBuilder` (in
+/// `shred_variant.rs`).
+pub(crate) enum PrimitiveVariantToArrowRowBuilder<'a> {
Review Comment:
I don't love splitting this out, but it seemed better than copying the whole
thing for both parent builders that use the same logic? Very open to ideas for
a better approach.
(the new definition of `VariantToArrowRowBuilder` is at L95 below)
##########
parquet-variant-compute/src/variant_to_arrow.rs:
##########
@@ -86,62 +88,129 @@ impl<'a> VariantToArrowRowBuilder<'a> {
Float16(b) => b.finish(),
Float32(b) => b.finish(),
Float64(b) => b.finish(),
+ }
+ }
+}
+
+impl<'a> VariantToArrowRowBuilder<'a> {
+ pub fn append_null(&mut self) -> Result<()> {
+ use VariantToArrowRowBuilder::*;
+ match self {
+ Primitive(b) => b.append_null(),
+ BinaryVariant(b) => b.append_null(),
+ WithPath(path_builder) => path_builder.append_null(),
+ }
+ }
+
+ pub fn append_value(&mut self, value: Variant<'_, '_>) -> Result<bool> {
+ use VariantToArrowRowBuilder::*;
+ match self {
+ Primitive(b) => b.append_value(&value),
+ BinaryVariant(b) => b.append_value(value),
+ WithPath(path_builder) => path_builder.append_value(value),
+ }
+ }
+
+ pub fn finish(self) -> Result<ArrayRef> {
+ use VariantToArrowRowBuilder::*;
+ match self {
+ Primitive(b) => b.finish(),
BinaryVariant(b) => b.finish(),
WithPath(path_builder) => path_builder.finish(),
}
}
}
-pub(crate) fn make_variant_to_arrow_row_builder<'a>(
- metadata: &BinaryViewArray,
- path: VariantPath<'a>,
- data_type: Option<&'a DataType>,
+/// Creates a primitive row builder, returning Err if the requested data type
is not primitive.
+pub(crate) fn make_primitive_variant_to_arrow_row_builder<'a>(
Review Comment:
A new helper for creating the new primitive row builder enum
(the original helper moved to L175 below)
##########
parquet-variant-compute/src/variant_to_arrow.rs:
##########
@@ -166,9 +235,9 @@ impl<'a> VariantPathRowBuilder<'a> {
self.builder.append_null()
}
- fn append_value(&mut self, value: &Variant<'_, '_>) -> Result<bool> {
+ fn append_value(&mut self, value: Variant<'_, '_>) -> Result<bool> {
Review Comment:
Now that the primitive row builder is separate a separate API from the
normal row builder, we can make the latter take `&Variant` while this can go
back to taking `Variant`.
(again below)
##########
parquet-variant-compute/src/shred_variant.rs:
##########
@@ -0,0 +1,683 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Module for shredding VariantArray with a given schema.
+
+use crate::variant_array::{ShreddedVariantFieldArray, StructArrayBuilder};
+use crate::variant_to_arrow::{
+ make_primitive_variant_to_arrow_row_builder,
PrimitiveVariantToArrowRowBuilder,
+};
+use crate::{VariantArray, VariantValueArrayBuilder};
+use arrow::array::Array as _;
+use arrow::array::{ArrayRef, BinaryViewArray, NullBufferBuilder};
+use arrow::buffer::NullBuffer;
+use arrow::compute::CastOptions;
+use arrow::datatypes::{DataType, Fields};
+use arrow::error::{ArrowError, Result};
+use parquet_variant::{ObjectBuilder, ReadOnlyMetadataBuilder, Variant,
EMPTY_VARIANT_METADATA};
+
+use indexmap::IndexMap;
+use std::sync::Arc;
+
+pub fn shred_variant(array: &VariantArray, as_type: &DataType) ->
Result<VariantArray> {
+ if array.typed_value_field().is_some() {
+ return Err(ArrowError::InvalidArgumentError(
+ "Input is already shredded".to_string(),
+ ));
+ }
+
+ if array.value_field().is_none() {
+ // all-null case
+ return Ok(VariantArray::from_parts(
+ array.metadata_field().clone(),
+ None,
+ None,
+ None,
+ ));
+ };
+
+ let cast_options = CastOptions::default();
+ let mut builder = make_variant_to_shredded_variant_arrow_row_builder(
+ as_type,
+ &cast_options,
+ array.len(),
+ true,
+ )?;
+ for i in 0..array.len() {
+ if array.is_null(i) {
+ builder.append_null()?;
+ } else {
+ builder.append_value(array.value(i))?;
+ }
+ }
+ let (value, typed_value, nulls) = builder.finish()?;
+ Ok(VariantArray::from_parts(
+ array.metadata_field().clone(),
+ Some(value),
+ Some(typed_value),
+ nulls,
+ ))
+}
+
+pub(crate) fn make_variant_to_shredded_variant_arrow_row_builder<'a>(
+ data_type: &'a DataType,
+ cast_options: &'a CastOptions,
+ len: usize,
Review Comment:
```suggestion
capacity: usize,
```
##########
parquet-variant-compute/src/shred_variant.rs:
##########
@@ -0,0 +1,683 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Module for shredding VariantArray with a given schema.
+
+use crate::variant_array::{ShreddedVariantFieldArray, StructArrayBuilder};
+use crate::variant_to_arrow::{
+ make_primitive_variant_to_arrow_row_builder,
PrimitiveVariantToArrowRowBuilder,
+};
+use crate::{VariantArray, VariantValueArrayBuilder};
+use arrow::array::Array as _;
+use arrow::array::{ArrayRef, BinaryViewArray, NullBufferBuilder};
+use arrow::buffer::NullBuffer;
+use arrow::compute::CastOptions;
+use arrow::datatypes::{DataType, Fields};
+use arrow::error::{ArrowError, Result};
+use parquet_variant::{ObjectBuilder, ReadOnlyMetadataBuilder, Variant,
EMPTY_VARIANT_METADATA};
+
+use indexmap::IndexMap;
+use std::sync::Arc;
+
+pub fn shred_variant(array: &VariantArray, as_type: &DataType) ->
Result<VariantArray> {
Review Comment:
needs a doc comment...
##########
parquet-variant-compute/src/shred_variant.rs:
##########
@@ -0,0 +1,683 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Module for shredding VariantArray with a given schema.
+
+use crate::variant_array::{ShreddedVariantFieldArray, StructArrayBuilder};
+use crate::variant_to_arrow::{
+ make_primitive_variant_to_arrow_row_builder,
PrimitiveVariantToArrowRowBuilder,
+};
+use crate::{VariantArray, VariantValueArrayBuilder};
+use arrow::array::Array as _;
+use arrow::array::{ArrayRef, BinaryViewArray, NullBufferBuilder};
+use arrow::buffer::NullBuffer;
+use arrow::compute::CastOptions;
+use arrow::datatypes::{DataType, Fields};
+use arrow::error::{ArrowError, Result};
+use parquet_variant::{ObjectBuilder, ReadOnlyMetadataBuilder, Variant,
EMPTY_VARIANT_METADATA};
+
+use indexmap::IndexMap;
+use std::sync::Arc;
+
+pub fn shred_variant(array: &VariantArray, as_type: &DataType) ->
Result<VariantArray> {
+ if array.typed_value_field().is_some() {
+ return Err(ArrowError::InvalidArgumentError(
+ "Input is already shredded".to_string(),
+ ));
+ }
+
+ if array.value_field().is_none() {
+ // all-null case
+ return Ok(VariantArray::from_parts(
+ array.metadata_field().clone(),
+ None,
+ None,
+ None,
+ ));
+ };
+
+ let cast_options = CastOptions::default();
+ let mut builder = make_variant_to_shredded_variant_arrow_row_builder(
+ as_type,
+ &cast_options,
+ array.len(),
+ true,
+ )?;
+ for i in 0..array.len() {
+ if array.is_null(i) {
+ builder.append_null()?;
+ } else {
+ builder.append_value(array.value(i))?;
+ }
+ }
+ let (value, typed_value, nulls) = builder.finish()?;
+ Ok(VariantArray::from_parts(
+ array.metadata_field().clone(),
+ Some(value),
+ Some(typed_value),
+ nulls,
+ ))
+}
+
+pub(crate) fn make_variant_to_shredded_variant_arrow_row_builder<'a>(
+ data_type: &'a DataType,
+ cast_options: &'a CastOptions,
+ len: usize,
+ top_level: bool,
+) -> Result<VariantToShreddedVariantRowBuilder<'a>> {
+ let builder = match data_type {
+ DataType::Struct(fields) => {
+ let typed_value_builder =
VariantToShreddedObjectVariantRowBuilder::try_new(
+ fields,
+ cast_options,
+ len,
+ top_level,
+ )?;
+ VariantToShreddedVariantRowBuilder::Object(typed_value_builder)
+ }
+ DataType::List(_)
+ | DataType::LargeList(_)
+ | DataType::ListView(_)
+ | DataType::LargeListView(_)
+ | DataType::FixedSizeList(..) => {
+ // TODO: Special handling for shredded variant arrays
+ return Err(ArrowError::NotYetImplemented(
+ "shred_variant not yet implemented for lists".to_string(),
+ ));
+ }
+ _ => {
+ let builder =
+ make_primitive_variant_to_arrow_row_builder(data_type,
cast_options, len)?;
+ let typed_value_builder =
+ VariantToShreddedPrimitiveVariantRowBuilder::new(builder, len,
top_level);
+ VariantToShreddedVariantRowBuilder::Primitive(typed_value_builder)
+ }
+ };
+ Ok(builder)
+}
+
+pub(crate) enum VariantToShreddedVariantRowBuilder<'a> {
+ Primitive(VariantToShreddedPrimitiveVariantRowBuilder<'a>),
+ Object(VariantToShreddedObjectVariantRowBuilder<'a>),
+}
+impl<'a> VariantToShreddedVariantRowBuilder<'a> {
+ pub fn append_null(&mut self) -> Result<()> {
+ use VariantToShreddedVariantRowBuilder::*;
+ match self {
+ Primitive(b) => b.append_null(),
+ Object(b) => b.append_null(),
+ }
+ }
+
+ pub fn append_value(&mut self, value: Variant<'_, '_>) -> Result<bool> {
+ use VariantToShreddedVariantRowBuilder::*;
+ match self {
+ Primitive(b) => b.append_value(value),
+ Object(b) => b.append_value(value),
+ }
+ }
+
+ pub fn finish(self) -> Result<(BinaryViewArray, ArrayRef,
Option<NullBuffer>)> {
+ use VariantToShreddedVariantRowBuilder::*;
+ match self {
+ Primitive(b) => b.finish(),
+ Object(b) => b.finish(),
+ }
+ }
+}
+
+/// A top-level variant shredder -- appending NULL produces typed_value=NULL
and value=Variant::Null
+pub(crate) struct VariantToShreddedPrimitiveVariantRowBuilder<'a> {
+ value_builder: VariantValueArrayBuilder,
+ typed_value_builder: PrimitiveVariantToArrowRowBuilder<'a>,
+ nulls: NullBufferBuilder,
+ top_level: bool,
+}
+
+impl<'a> VariantToShreddedPrimitiveVariantRowBuilder<'a> {
+ pub(crate) fn new(
+ typed_value_builder: PrimitiveVariantToArrowRowBuilder<'a>,
+ len: usize,
+ top_level: bool,
+ ) -> Self {
+ Self {
+ value_builder: VariantValueArrayBuilder::new(len),
+ typed_value_builder,
+ nulls: NullBufferBuilder::new(len),
+ top_level,
+ }
+ }
+ fn append_null(&mut self) -> Result<()> {
+ if self.top_level {
+ self.nulls.append_null();
+ self.value_builder.append_null();
+ } else {
+ self.nulls.append_non_null();
+ self.value_builder.append_value(Variant::Null);
+ }
+ self.typed_value_builder.append_null()
+ }
+ fn append_value(&mut self, value: Variant<'_, '_>) -> Result<bool> {
+ self.nulls.append_non_null();
+ if self.typed_value_builder.append_value(&value)? {
+ self.value_builder.append_null();
+ } else {
+ self.value_builder.append_value(value);
+ }
+ Ok(true)
+ }
+ fn finish(mut self) -> Result<(BinaryViewArray, ArrayRef,
Option<NullBuffer>)> {
+ Ok((
+ self.value_builder.build()?,
+ self.typed_value_builder.finish()?,
+ self.nulls.finish(),
+ ))
+ }
+}
+
+pub(crate) struct VariantToShreddedObjectVariantRowBuilder<'a> {
+ value_builder: VariantValueArrayBuilder,
+ typed_value_builders: IndexMap<&'a str,
VariantToShreddedVariantRowBuilder<'a>>,
+ typed_value_nulls: NullBufferBuilder,
+ nulls: NullBufferBuilder,
+ top_level: bool,
+}
+
+impl<'a> VariantToShreddedObjectVariantRowBuilder<'a> {
+ fn try_new(
+ fields: &'a Fields,
+ cast_options: &'a CastOptions,
+ len: usize,
+ top_level: bool,
+ ) -> Result<Self> {
+ let typed_value_builders = fields.iter().map(|field| {
+ let builder = make_variant_to_shredded_variant_arrow_row_builder(
+ field.data_type(),
+ cast_options,
+ len,
+ top_level,
+ )?;
+ Ok((field.name().as_str(), builder))
+ });
+ Ok(Self {
+ value_builder: VariantValueArrayBuilder::new(len),
+ typed_value_builders: typed_value_builders.collect::<Result<_>>()?,
+ typed_value_nulls: NullBufferBuilder::new(len),
+ nulls: NullBufferBuilder::new(len),
+ top_level,
+ })
+ }
+
+ fn append_null(&mut self) -> Result<()> {
+ if self.top_level {
+ self.nulls.append_null();
+ self.value_builder.append_null();
+ } else {
+ self.nulls.append_non_null();
+ self.value_builder.append_value(Variant::Null);
+ }
+ self.typed_value_nulls.append_null();
+ for (_, typed_value_builder) in &mut self.typed_value_builders {
+ typed_value_builder.append_null()?;
+ }
+ Ok(())
+ }
+ fn append_value(&mut self, value: Variant<'_, '_>) -> Result<bool> {
+ let Variant::Object(ref obj) = value else {
+ // Not an object => fall back
+ self.nulls.append_non_null();
+ self.value_builder.append_value(value);
+ self.typed_value_nulls.append_null();
+ for (_, typed_value_builder) in &mut self.typed_value_builders {
+ typed_value_builder.append_null()?;
+ }
+ return Ok(false);
+ };
+
+ // Route the object's fields by name as either shredded or unshredded
+ let metadata =
value.metadata().cloned().unwrap_or(EMPTY_VARIANT_METADATA);
+ let mut metadata_builder = ReadOnlyMetadataBuilder::new(metadata);
+ let state = self.value_builder.parent_state(&mut metadata_builder);
+ let mut object_builder = ObjectBuilder::new(state, false);
+ let mut seen = std::collections::HashSet::new();
+ let mut partially_shredded = false;
+ for (field_name, value) in obj.iter() {
+ match self.typed_value_builders.get_mut(field_name) {
+ Some(typed_value_builder) => {
+ typed_value_builder.append_value(value)?;
+ seen.insert(field_name);
+ }
+ None => {
+ object_builder.insert_bytes(field_name, value);
+ partially_shredded = true;
+ }
+ }
+ }
+
+ // Handle missing fields
+ for (field_name, typed_value_builder) in &mut
self.typed_value_builders {
+ if !seen.contains(field_name) {
+ typed_value_builder.append_null()?;
+ }
+ }
+
+ // Only emit the value if it captured any unshredded object fields
+ if partially_shredded {
+ object_builder.finish();
+ } else {
+ drop(object_builder);
+ self.value_builder.append_null();
+ }
+
+ self.typed_value_nulls.append_non_null();
+ self.nulls.append_non_null();
+ Ok(true)
+ }
+ fn finish(mut self) -> Result<(BinaryViewArray, ArrayRef,
Option<NullBuffer>)> {
+ let mut builder = StructArrayBuilder::new();
+ for (field_name, typed_value_builder) in self.typed_value_builders {
+ let (value, typed_value, nulls) = typed_value_builder.finish()?;
+ let array =
+ ShreddedVariantFieldArray::from_parts(Some(value),
Some(typed_value), nulls);
+ builder = builder.with_field(field_name, Arc::new(array), false);
+ }
+ if let Some(nulls) = self.typed_value_nulls.finish() {
+ builder = builder.with_nulls(nulls);
+ }
+ Ok((
+ self.value_builder.build()?,
+ Arc::new(builder.build()),
+ self.nulls.finish(),
+ ))
+ }
+}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use crate::VariantArrayBuilder;
+ use arrow::array::{Float64Array, Int64Array};
+ use arrow::datatypes::{DataType, Field, Fields};
+ use parquet_variant::{Variant, VariantBuilderExt};
+ use std::sync::Arc;
+
+ fn create_test_variant_array(values: Vec<Option<Variant<'_, '_>>>) ->
VariantArray {
+ let mut builder = VariantArrayBuilder::new(values.len());
+ for value in values {
+ match value {
+ Some(v) => builder.append_variant(v),
+ None => builder.append_null(),
+ }
+ }
+ builder.build()
+ }
+
+ // Input Validation Tests (3 tests - cannot consolidate)
+
Review Comment:
```suggestion
```
##########
parquet-variant-compute/src/shred_variant.rs:
##########
@@ -0,0 +1,683 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Module for shredding VariantArray with a given schema.
+
+use crate::variant_array::{ShreddedVariantFieldArray, StructArrayBuilder};
+use crate::variant_to_arrow::{
+ make_primitive_variant_to_arrow_row_builder,
PrimitiveVariantToArrowRowBuilder,
+};
+use crate::{VariantArray, VariantValueArrayBuilder};
+use arrow::array::Array as _;
+use arrow::array::{ArrayRef, BinaryViewArray, NullBufferBuilder};
Review Comment:
```suggestion
use arrow::array::{Array as _, ArrayRef, BinaryViewArray, NullBufferBuilder};
```
##########
parquet-variant-compute/src/shred_variant.rs:
##########
@@ -0,0 +1,683 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Module for shredding VariantArray with a given schema.
+
+use crate::variant_array::{ShreddedVariantFieldArray, StructArrayBuilder};
+use crate::variant_to_arrow::{
+ make_primitive_variant_to_arrow_row_builder,
PrimitiveVariantToArrowRowBuilder,
+};
+use crate::{VariantArray, VariantValueArrayBuilder};
+use arrow::array::Array as _;
+use arrow::array::{ArrayRef, BinaryViewArray, NullBufferBuilder};
+use arrow::buffer::NullBuffer;
+use arrow::compute::CastOptions;
+use arrow::datatypes::{DataType, Fields};
+use arrow::error::{ArrowError, Result};
+use parquet_variant::{ObjectBuilder, ReadOnlyMetadataBuilder, Variant,
EMPTY_VARIANT_METADATA};
+
+use indexmap::IndexMap;
+use std::sync::Arc;
+
+pub fn shred_variant(array: &VariantArray, as_type: &DataType) ->
Result<VariantArray> {
+ if array.typed_value_field().is_some() {
+ return Err(ArrowError::InvalidArgumentError(
+ "Input is already shredded".to_string(),
+ ));
+ }
+
+ if array.value_field().is_none() {
+ // all-null case
+ return Ok(VariantArray::from_parts(
+ array.metadata_field().clone(),
+ None,
+ None,
+ None,
+ ));
+ };
+
+ let cast_options = CastOptions::default();
+ let mut builder = make_variant_to_shredded_variant_arrow_row_builder(
+ as_type,
+ &cast_options,
+ array.len(),
+ true,
+ )?;
+ for i in 0..array.len() {
+ if array.is_null(i) {
+ builder.append_null()?;
+ } else {
+ builder.append_value(array.value(i))?;
+ }
+ }
+ let (value, typed_value, nulls) = builder.finish()?;
+ Ok(VariantArray::from_parts(
+ array.metadata_field().clone(),
+ Some(value),
+ Some(typed_value),
+ nulls,
+ ))
+}
+
+pub(crate) fn make_variant_to_shredded_variant_arrow_row_builder<'a>(
+ data_type: &'a DataType,
+ cast_options: &'a CastOptions,
+ len: usize,
+ top_level: bool,
+) -> Result<VariantToShreddedVariantRowBuilder<'a>> {
+ let builder = match data_type {
+ DataType::Struct(fields) => {
+ let typed_value_builder =
VariantToShreddedObjectVariantRowBuilder::try_new(
+ fields,
+ cast_options,
+ len,
+ top_level,
+ )?;
+ VariantToShreddedVariantRowBuilder::Object(typed_value_builder)
+ }
+ DataType::List(_)
+ | DataType::LargeList(_)
+ | DataType::ListView(_)
+ | DataType::LargeListView(_)
+ | DataType::FixedSizeList(..) => {
+ // TODO: Special handling for shredded variant arrays
+ return Err(ArrowError::NotYetImplemented(
+ "shred_variant not yet implemented for lists".to_string(),
+ ));
+ }
+ _ => {
+ let builder =
+ make_primitive_variant_to_arrow_row_builder(data_type,
cast_options, len)?;
+ let typed_value_builder =
+ VariantToShreddedPrimitiveVariantRowBuilder::new(builder, len,
top_level);
+ VariantToShreddedVariantRowBuilder::Primitive(typed_value_builder)
+ }
+ };
+ Ok(builder)
+}
+
+pub(crate) enum VariantToShreddedVariantRowBuilder<'a> {
+ Primitive(VariantToShreddedPrimitiveVariantRowBuilder<'a>),
+ Object(VariantToShreddedObjectVariantRowBuilder<'a>),
+}
+impl<'a> VariantToShreddedVariantRowBuilder<'a> {
+ pub fn append_null(&mut self) -> Result<()> {
+ use VariantToShreddedVariantRowBuilder::*;
+ match self {
+ Primitive(b) => b.append_null(),
+ Object(b) => b.append_null(),
+ }
+ }
+
+ pub fn append_value(&mut self, value: Variant<'_, '_>) -> Result<bool> {
+ use VariantToShreddedVariantRowBuilder::*;
+ match self {
+ Primitive(b) => b.append_value(value),
+ Object(b) => b.append_value(value),
+ }
+ }
+
+ pub fn finish(self) -> Result<(BinaryViewArray, ArrayRef,
Option<NullBuffer>)> {
+ use VariantToShreddedVariantRowBuilder::*;
+ match self {
+ Primitive(b) => b.finish(),
+ Object(b) => b.finish(),
+ }
+ }
+}
+
+/// A top-level variant shredder -- appending NULL produces typed_value=NULL
and value=Variant::Null
+pub(crate) struct VariantToShreddedPrimitiveVariantRowBuilder<'a> {
+ value_builder: VariantValueArrayBuilder,
+ typed_value_builder: PrimitiveVariantToArrowRowBuilder<'a>,
+ nulls: NullBufferBuilder,
+ top_level: bool,
+}
+
+impl<'a> VariantToShreddedPrimitiveVariantRowBuilder<'a> {
+ pub(crate) fn new(
+ typed_value_builder: PrimitiveVariantToArrowRowBuilder<'a>,
+ len: usize,
+ top_level: bool,
+ ) -> Self {
+ Self {
+ value_builder: VariantValueArrayBuilder::new(len),
+ typed_value_builder,
+ nulls: NullBufferBuilder::new(len),
+ top_level,
+ }
+ }
+ fn append_null(&mut self) -> Result<()> {
+ if self.top_level {
+ self.nulls.append_null();
+ self.value_builder.append_null();
+ } else {
+ self.nulls.append_non_null();
+ self.value_builder.append_value(Variant::Null);
+ }
+ self.typed_value_builder.append_null()
+ }
+ fn append_value(&mut self, value: Variant<'_, '_>) -> Result<bool> {
+ self.nulls.append_non_null();
+ if self.typed_value_builder.append_value(&value)? {
+ self.value_builder.append_null();
+ } else {
+ self.value_builder.append_value(value);
+ }
+ Ok(true)
+ }
+ fn finish(mut self) -> Result<(BinaryViewArray, ArrayRef,
Option<NullBuffer>)> {
+ Ok((
+ self.value_builder.build()?,
+ self.typed_value_builder.finish()?,
+ self.nulls.finish(),
+ ))
+ }
+}
+
+pub(crate) struct VariantToShreddedObjectVariantRowBuilder<'a> {
+ value_builder: VariantValueArrayBuilder,
+ typed_value_builders: IndexMap<&'a str,
VariantToShreddedVariantRowBuilder<'a>>,
+ typed_value_nulls: NullBufferBuilder,
+ nulls: NullBufferBuilder,
+ top_level: bool,
+}
+
+impl<'a> VariantToShreddedObjectVariantRowBuilder<'a> {
+ fn try_new(
+ fields: &'a Fields,
+ cast_options: &'a CastOptions,
+ len: usize,
+ top_level: bool,
+ ) -> Result<Self> {
+ let typed_value_builders = fields.iter().map(|field| {
+ let builder = make_variant_to_shredded_variant_arrow_row_builder(
+ field.data_type(),
+ cast_options,
+ len,
+ top_level,
+ )?;
+ Ok((field.name().as_str(), builder))
+ });
+ Ok(Self {
+ value_builder: VariantValueArrayBuilder::new(len),
+ typed_value_builders: typed_value_builders.collect::<Result<_>>()?,
+ typed_value_nulls: NullBufferBuilder::new(len),
+ nulls: NullBufferBuilder::new(len),
+ top_level,
+ })
+ }
+
+ fn append_null(&mut self) -> Result<()> {
+ if self.top_level {
+ self.nulls.append_null();
+ self.value_builder.append_null();
+ } else {
+ self.nulls.append_non_null();
+ self.value_builder.append_value(Variant::Null);
+ }
Review Comment:
Duplicated above...
--
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
