[VARIANT] Add support for the json_to_variant API

parquet-variant/Cargo.toml

@@ -35,7 +35,10 @@ rust-version = "1.83"
 [dependencies]
 arrow-schema = { workspace = true }
 chrono = { workspace = true }
-serde_json = "1.0"
+# "arbitrary_precision" allows us to manually parse numbers. "preserve_order" does not automatically
+# sort object keys
+serde_json = { version = "1.0", features = ["arbitrary_precision", "preserve_order"] }
+rust_decimal = "1.37.2"
 base64 = "0.21"
 
 [lib]

parquet-variant/examples/variant_from_json_examples.rs

@@ -0,0 +1,58 @@
+// 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.
+
+//! Example showing how to convert Variant values to JSON
+
+use parquet_variant::{
+    json_to_variant, variant_to_json, variant_to_json_string, variant_to_json_value,
+    SampleVecBasedVariantBufferManager,
+};
+
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    // The caller must provide an object implementing the `VariantBufferManager` trait to the library.
+    // This allows the library to write the constructed variant to buffers provided by the caller.
+    // This way, the caller has direct control over the output buffers.
+    let mut variant_buffer_manager = SampleVecBasedVariantBufferManager {
+        value_buffer: vec![0u8; 1],
+        metadata_buffer: vec![0u8; 1],
+    };
+
+    let person_string = "{\"name\":\"Alice\", \"age\":30, ".to_string()
+        + "\"email\":\"[email protected]\", \"is_active\": true, \"score\": 95.7,"
+        + "\"additional_info\": null}";
+    let (metadata_size, value_size) = json_to_variant(&person_string, &mut variant_buffer_manager)?;
+
+    let variant = parquet_variant::Variant::try_new(
+        &variant_buffer_manager.metadata_buffer[..metadata_size],
+        &variant_buffer_manager.value_buffer[..value_size],
+    )?;
+
+    let json_string = variant_to_json_string(&variant)?;
+    let json_value = variant_to_json_value(&variant)?;
+    let pretty_json = serde_json::to_string_pretty(&json_value)?;
+    println!("{}", pretty_json);
+
+    let mut buffer = Vec::new();
+    variant_to_json(&mut buffer, &variant)?;
+    let buffer_result = String::from_utf8(buffer)?;
+
+    // Verify all methods produce the same result
+    assert_eq!(json_string, buffer_result);
+    assert_eq!(json_string, serde_json::to_string(&json_value)?);
+
+    Ok(())
+}

parquet-variant/src/builder.rs

@@ -721,6 +721,42 @@ impl<'a, 'b> ObjectBuilder<'a, 'b> {
     }
 }
 
+pub(crate) trait AppendVariantHelper {
+    fn append_value_helper<'m, 'd, T: Into<Variant<'m, 'd>>>(&mut self, value: T);
+
+    fn new_list_helper(&mut self) -> ListBuilder;
+
+    fn new_object_helper(&mut self) -> ObjectBuilder;
+}
+
+impl AppendVariantHelper for ListBuilder<'_> {
+    fn append_value_helper<'m, 'd, T: Into<Variant<'m, 'd>>>(&mut self, value: T) {
+        self.append_value(value);
+    }
+
+    fn new_list_helper(&mut self) -> ListBuilder {
+        self.new_list()
+    }
+
+    fn new_object_helper(&mut self) -> ObjectBuilder {
+        self.new_object()
+    }
+}
+
+impl AppendVariantHelper for VariantBuilder {
+    fn append_value_helper<'m, 'd, T: Into<Variant<'m, 'd>>>(&mut self, value: T) {
+        self.append_value(value);
+    }
+
+    fn new_list_helper(&mut self) -> ListBuilder {
+        self.new_list()
+    }
+
+    fn new_object_helper(&mut self) -> ObjectBuilder {
+        self.new_object()
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;

parquet-variant/src/from_json.rs

@@ -0,0 +1,143 @@
+pub use crate::variant::{VariantDecimal4, VariantDecimal8};
+use crate::variant_buffer_manager::VariantBufferManager;
+use crate::{AppendVariantHelper, ListBuilder, ObjectBuilder, Variant, VariantBuilder};
+use arrow_schema::ArrowError;
+use rust_decimal::prelude::*;
+use serde_json::{Map, Number, Value};
+
+/// Eventually, internal writes should also be performed using VariantBufferManager instead of
+/// ValueBuffer and MetadataBuffer so the caller has control of the memory.
+/// Returns a pair <value_size, metadata_size>
+pub fn json_to_variant(
+    json: &str,
+    variant_buffer_manager: &mut impl VariantBufferManager,
+) -> Result<(usize, usize), ArrowError> {
+    let mut builder = VariantBuilder::new();
+    let json: Value = serde_json::from_str(json)
+        .map_err(|e| ArrowError::InvalidArgumentError(format!("JSON format error: {}", e)))?;
+
+    build_json(&json, &mut builder)?;
+    let (metadata, value) = builder.finish();
+    let value_size = value.len();
+    let metadata_size = metadata.len();
+
+    // Write to caller's buffers - Remove this when the library internally writes to the caller's
+    // buffers anyway
+    let caller_metadata_buffer =
+        variant_buffer_manager.ensure_size_and_borrow_metadata_buffer(metadata_size)?;
+    caller_metadata_buffer[..metadata_size].copy_from_slice(metadata.as_slice());
+    let caller_value_buffer =
+        variant_buffer_manager.ensure_size_and_borrow_value_buffer(value_size)?;
+    caller_value_buffer[..value_size].copy_from_slice(value.as_slice());
+    Ok((metadata_size, value_size))
+}
+
+fn build_json(json: &Value, builder: &mut VariantBuilder) -> Result<(), ArrowError> {
+    append_json(json, builder)?;
+    Ok(())
+}
+
+fn variant_from_number<'a, 'b>(n: &Number) -> Result<Variant<'a, 'b>, ArrowError> {
+    if let Some(i) = n.as_i64() {
+        // Find minimum Integer width to fit
+        if i as i8 as i64 == i {
+            Ok((i as i8).into())
+        } else if i as i16 as i64 == i {
+            Ok((i as i16).into())
+        } else if i as i32 as i64 == i {
+            Ok((i as i32).into())
+        } else {
+            Ok(i.into())
+        }
+    } else {
+        // Try decimal
+        // TODO: Replace with custom decimal parsing as the rust_decimal library only supports
+        // a max unscaled value of 2^96.
+        match Decimal::from_str_exact(n.as_str()) {
+            Ok(dec) => {
+                let unscaled: i128 = dec.mantissa();
+                let scale = dec.scale() as u8;
+                if unscaled.abs() <= VariantDecimal4::MAX_UNSCALED_VALUE as i128
+                    && scale <= VariantDecimal4::MAX_PRECISION as u8
+                {
+                    (unscaled as i32, scale).try_into()
+                } else if unscaled.abs() <= VariantDecimal8::MAX_UNSCALED_VALUE as i128
+                    && scale <= VariantDecimal8::MAX_PRECISION as u8
+                {
+                    (unscaled as i64, scale).try_into()
+                } else {
+                    (unscaled, scale).try_into()
+                }
+            }
+            Err(_) => {
+                // Try double
+                match n.as_f64() {
+                    Some(f) => return Ok(f.into()),
+                    None => Err(ArrowError::InvalidArgumentError(format!(
+                        "Failed to parse {} as number",
+                        n.as_str()
+                    ))),
+                }?
+            }
+        }
+    }
+}
+
+fn append_json(json: &Value, builder: &mut impl AppendVariantHelper) -> Result<(), ArrowError> {
+    match json {
+        Value::Null => builder.append_value_helper(Variant::Null),
+        Value::Bool(b) => builder.append_value_helper(*b),
+        Value::Number(n) => {
+            builder.append_value_helper(variant_from_number(n)?);
+        }
+        Value::String(s) => builder.append_value_helper(s.as_str()),
+        Value::Array(arr) => {
+            let mut list_builder = builder.new_list_helper();
+            build_list(arr, &mut list_builder)?;
+            list_builder.finish();
+        }
+        Value::Object(obj) => {
+            let mut obj_builder = builder.new_object_helper();
+            build_object(obj, &mut obj_builder)?;
+            obj_builder.finish();
+        }
+    };
+    Ok(())
+}
+
+fn build_list(arr: &[Value], builder: &mut ListBuilder) -> Result<(), ArrowError> {
+    for val in arr {
+        append_json(val, builder)?;
+    }
+    Ok(())
+}
+
+fn build_object<'a, 'b>(
+    obj: &'b Map<String, Value>,
+    builder: &mut ObjectBuilder<'a, 'b>,
+) -> Result<(), ArrowError> {
+    for (key, value) in obj.iter() {
+        let mut field_builder = ObjectFieldBuilder { key, builder };
+        append_json(value, &mut field_builder)?;
+    }
+    Ok(())
+}
+
+struct ObjectFieldBuilder<'a, 'b, 'c> {
+    key: &'a str,
+    builder: &'b mut ObjectBuilder<'c, 'a>,
+}
+
+impl AppendVariantHelper for ObjectFieldBuilder<'_, '_, '_> {
+    fn append_value_helper<'m, 'd, T: Into<Variant<'m, 'd>>>(&mut self, value: T) {
+        self.builder.insert(self.key, value);
+    }
+
+    fn new_list_helper(&mut self) -> ListBuilder {
+        self.builder.new_list(self.key)
+    }
+
+    fn new_object_helper(&mut self) -> ObjectBuilder {
+        self.builder.new_object(self.key)
+    }
+}

parquet-variant/src/lib.rs

@@ -33,10 +33,14 @@ mod decoder;
 mod variant;
 // TODO: dead code removal
 mod builder;
+mod from_json;
 mod to_json;
 #[allow(dead_code)]
 mod utils;
+mod variant_buffer_manager;
 
 pub use builder::*;
+pub use from_json::json_to_variant;
 pub use to_json::{variant_to_json, variant_to_json_string, variant_to_json_value};
 pub use variant::*;
+pub use variant_buffer_manager::{SampleVecBasedVariantBufferManager, VariantBufferManager};

parquet-variant/src/variant/decimal.rs

@@ -31,8 +31,8 @@ pub struct VariantDecimal4 {
 }
 
 impl VariantDecimal4 {
-    const MAX_PRECISION: u32 = 9;
-    const MAX_UNSCALED_VALUE: u32 = 10_u32.pow(Self::MAX_PRECISION) - 1;
+    pub(crate) const MAX_PRECISION: u32 = 9;
+    pub(crate) const MAX_UNSCALED_VALUE: u32 = 10_u32.pow(Self::MAX_PRECISION) - 1;
 
     pub fn try_new(integer: i32, scale: u8) -> Result<Self, ArrowError> {
         // Validate that scale doesn't exceed precision
@@ -73,8 +73,8 @@ pub struct VariantDecimal8 {
 }
 
 impl VariantDecimal8 {
-    const MAX_PRECISION: u32 = 18;
-    const MAX_UNSCALED_VALUE: u64 = 10_u64.pow(Self::MAX_PRECISION) - 1;
+    pub(crate) const MAX_PRECISION: u32 = 18;
+    pub(crate) const MAX_UNSCALED_VALUE: u64 = 10_u64.pow(Self::MAX_PRECISION) - 1;
 
     pub fn try_new(integer: i64, scale: u8) -> Result<Self, ArrowError> {
         // Validate that scale doesn't exceed precision

parquet-variant/src/variant_buffer_manager.rs

@@ -0,0 +1,57 @@
+use arrow_schema::ArrowError;
+
+pub trait VariantBufferManager {
+    /// Returns the slice where the variant metadata needs to be written to. This method may be
+    /// called several times during the construction of a new `metadata` field in a variant. The
+    /// implementation must make sure that on every call, all the data written to the metadata
+    /// buffer so far are preserved.
+    /// The implementation must also make sure that the length of the slice being returned is at
+    /// least `size` bytes. The implementation may throw an error if it is unable to fulfill its
+    /// requirements.
+    fn ensure_size_and_borrow_metadata_buffer(
+        &mut self,
+        size: usize,
+    ) -> Result<&mut [u8], ArrowError>;
+
+    /// Returns the slice where value needs to be written to. This method may be called several
+    /// times during the construction of a new `value` field in a variant. The implementation must
+    /// make sure that on every call, all the data written to the value buffer so far are preserved.
+    /// The implementation must also make sure that the length of the slice being returned is at
+    /// least `size` bytes. The implementation may throw an error if it is unable to fulfill its
+    /// requirements.
+    fn ensure_size_and_borrow_value_buffer(&mut self, size: usize)
+        -> Result<&mut [u8], ArrowError>;
+}
+
+pub struct SampleVecBasedVariantBufferManager {
+    pub value_buffer: Vec<u8>,
+    pub metadata_buffer: Vec<u8>,
+}
+
+impl VariantBufferManager for SampleVecBasedVariantBufferManager {
+    fn ensure_size_and_borrow_metadata_buffer(
+        &mut self,
+        size: usize,
+    ) -> Result<&mut [u8], ArrowError> {
+        let cur_len = self.metadata_buffer.len();
+        if size > cur_len {
+            // Reallocate larger buffer
+            let new_len = size.next_power_of_two();
+            self.metadata_buffer.resize(new_len, 0);
+        }
+        Ok(&mut self.metadata_buffer)
+    }
+
+    fn ensure_size_and_borrow_value_buffer(
+        &mut self,
+        size: usize,
+    ) -> Result<&mut [u8], ArrowError> {
+        let cur_len = self.value_buffer.len();
+        if size > cur_len {
+            // Reallocate larger buffer
+            let new_len = size.next_power_of_two();
+            self.value_buffer.resize(new_len, 0);
+        }
+        Ok(&mut self.value_buffer)
+    }
+}