alamb commented on code in PR #7911:
URL: https://github.com/apache/arrow-rs/pull/7911#discussion_r2205100761
##########
parquet-variant-compute/src/variant_array_builder.rs:
##########
@@ -168,7 +166,169 @@ impl VariantArrayBuilder {
self.value_buffer.extend_from_slice(value);
}
- // TODO: Return a Variant builder that will write to the underlying
buffers (TODO)
+ /// Return a `VariantArrayVariantBuilder` that writes directly to the
+ /// buffers of this builder.
+ ///
+ /// You must call [`VariantArrayVariantBuilder::finish`] to complete the
builder
+ ///
+ /// # Example
+ /// ```
+ /// # use parquet_variant::{Variant, VariantBuilder, VariantBuilderExt};
+ /// # use parquet_variant_compute::{VariantArray, VariantArrayBuilder};
+ /// let mut array_builder = VariantArrayBuilder::new(10);
+ ///
+ /// // First row has a string
+ /// let mut variant_builder = array_builder.variant_builder();
+ /// variant_builder.append_value("Hello, World!");
+ /// // must call finish to write the variant to the buffers
+ /// variant_builder.finish();
+ ///
+ /// // Second row is an object
+ /// let mut variant_builder = array_builder.variant_builder();
+ /// let mut obj = variant_builder.new_object();
+ /// obj.insert("my_field", 42i64);
+ /// obj.finish().unwrap();
+ /// variant_builder.finish();
+ ///
+ /// // finalize the array
+ /// let variant_array: VariantArray = array_builder.build();
+ ///
+ /// // verify what we wrote is still there
+ /// assert_eq!(variant_array.value(0), Variant::from("Hello, World!"));
+ /// assert!(variant_array.value(1).as_object().is_some());
+ /// ```
+ pub fn variant_builder(&mut self) -> VariantArrayVariantBuilder {
+ // append directly into the metadata and value buffers
+ let metadata_buffer = std::mem::take(&mut self.metadata_buffer);
+ let value_buffer = std::mem::take(&mut self.value_buffer);
+ VariantArrayVariantBuilder::new(self, metadata_buffer, value_buffer)
+ }
+}
+
+/// A `VariantBuilder` that writes directly to the buffers of a
`VariantArrayBuilder`.
+///
+/// Note this struct implements [`VariantBuilderExt`], so it can be used
+/// as a drop-in replacement for [`VariantBuilder`] in most cases.
+///
+/// If [`Self::finish`] is not called, any changes will be rolled back
+///
+/// See [`VariantArrayBuilder::variant_builder`] for an example
+pub struct VariantArrayVariantBuilder<'a> {
+ /// was finish called?
+ finished: bool,
+ /// starting offset in the variant_builder's `metadata` buffer
+ metadata_offset: usize,
+ /// starting offset in the variant_builder's `value` buffer
+ value_offset: usize,
+ /// Parent array builder that this variant builder writes to. Buffers
+ /// have been moved into the variant builder, and must be returned on
+ /// drop
+ array_builder: &'a mut VariantArrayBuilder,
+ /// Builder for the in progress variant value, temporarily owns the buffers
+ /// from `array_builder`
+ variant_builder: VariantBuilder,
+}
+
+impl<'a> VariantBuilderExt for VariantArrayVariantBuilder<'a> {
+ fn append_value<'m, 'v>(&mut self, value: impl Into<Variant<'m, 'v>>) {
+ self.variant_builder.append_value(value);
+ }
+
+ fn new_list(&mut self) -> ListBuilder {
+ self.variant_builder.new_list()
+ }
+
+ fn new_object(&mut self) -> ObjectBuilder {
+ self.variant_builder.new_object()
+ }
+}
+
+impl<'a> VariantArrayVariantBuilder<'a> {
+ /// Constructs a new VariantArrayVariantBuilder
+ ///
+ /// Note this is not public as this is a structure that is logically
+ /// part of the [`VariantArrayBuilder`] and relies on its internal
structure
+ fn new(
+ array_builder: &'a mut VariantArrayBuilder,
+ metadata_buffer: Vec<u8>,
+ value_buffer: Vec<u8>,
+ ) -> Self {
+ let metadata_offset = metadata_buffer.len();
+ let value_offset = value_buffer.len();
+ VariantArrayVariantBuilder {
+ finished: false,
+ metadata_offset,
+ value_offset,
+ variant_builder: VariantBuilder::new_with_buffers(metadata_buffer,
value_buffer),
+ array_builder,
+ }
+ }
+
+ /// Return a reference to the underlying `VariantBuilder`
+ pub fn inner(&self) -> &VariantBuilder {
+ &self.variant_builder
+ }
+
+ /// Return a mutable reference to the underlying `VariantBuilder`
+ pub fn inner_mut(&mut self) -> &mut VariantBuilder {
+ &mut self.variant_builder
+ }
+
+ /// Called to finish the in progress variant and write it to the underlying
+ /// buffers
+ ///
+ /// Note if you do not call finish, on drop any changes made to the
+ /// underlying buffers will be rolled back.
+ pub fn finish(mut self) {
+ self.finished = true;
+ // Note: buffers are returned and replaced in the drop impl
+ }
+}
+
+impl<'a> Drop for VariantArrayVariantBuilder<'a> {
+ /// If the builder was not finished, roll back any changes made to the
+ /// underlying buffers (by truncating them)
+ fn drop(&mut self) {
Review Comment:
> They can truncate the metadata dictionary on rollback, which would
eliminate the false allocations that survive a rollback today
That is an excellent point
> We can allocate the value bytes directly in the base buffer (instead of
using a separate Vec)
That sounds like a great way to avoid the extra allocation
- https://github.com/apache/arrow-rs/issues/7899
> Once we're using splice, it opens the door to pre-allocate the space for
the value offset and field arrays, in case the caller knows how many fields or
array elements there are.
This is also a great idea 🤯
--
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
