Initial API for reading Variant data and metadata (#7535)

# Which issue does this PR close?

<!--
We generally require a GitHub issue to be filed for all bug fixes and
enhancements and this helps us generate change logs for our releases.
You can link an issue to this PR using the GitHub syntax. For example
`Closes #123` indicates that this PR will close issue #123.
-->

Closes #7423 

# Rationale for this change
 
<!--
Why are you proposing this change? If this is already explained clearly
in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand
your changes and offer better suggestions for fixes.
-->
We need to agree on an API for reading Variant metadata. Based on the
work and discussions in #7452, in this PR we propose an API plus an
implementation (WIP while draft) for reading variant metadata in the
parquet-variant crate.
A lot of the work is based on the work in #7452 by @PinkCrow007 and
feedback from @alamb, @scovich, and @Weijun-H.

# What changes are included in this PR?
- Adds Variant enum (and associated structs)
- Adds an API for parsing and reading metadata
- Adds an API for parsing and reading Variant values of various types

We attempt to be result- and validation driven while ensuring
zero-allocations, and we do so avoiding `serde_json`.

We tried to keep the Variant API similar to the `Json::Value` api.

<!--
There is no need to duplicate the description in the issue here but it
is sometimes worth providing a summary of the individual changes in this
PR.
-->

# Are there any user-facing changes?
The new API's added in parquet-variant will be user facing.

---------

Co-authored-by: Andrew Lamb <[email protected]>
Co-authored-by: Ryan Johnson <[email protected]>

Author: 3 people

parquet-variant/Cargo.toml

@@ -31,6 +31,7 @@ edition = { workspace = true }
 rust-version = { workspace = true }
 
 [dependencies]
+arrow-schema = "55.1.0"
 
 [lib]
 

parquet-variant/src/decoder.rs

@@ -0,0 +1,158 @@
+// 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.
+use arrow_schema::ArrowError;
+use std::array::TryFromSliceError;
+
+use crate::utils::{array_from_slice, first_byte_from_slice, string_from_slice};
+
+#[derive(Debug, Clone, Copy)]
+pub enum VariantBasicType {
+    Primitive = 0,
+    ShortString = 1,
+    Object = 2,
+    Array = 3,
+}
+
+#[derive(Debug, Clone, Copy)]
+pub enum VariantPrimitiveType {
+    Null = 0,
+    BooleanTrue = 1,
+    BooleanFalse = 2,
+    Int8 = 3,
+    // TODO: Add types for the rest of primitives, once API is agreed upon
+    String = 16,
+}
+
+/// Extracts the basic type from a header byte
+pub(crate) fn get_basic_type(header: u8) -> Result<VariantBasicType, ArrowError> {
+    // See https://github.com/apache/parquet-format/blob/master/VariantEncoding.md#value-encoding
+    let basic_type = header & 0x03; // Basic type is encoded in the first 2 bits
+    let basic_type = match basic_type {
+        0 => VariantBasicType::Primitive,
+        1 => VariantBasicType::ShortString,
+        2 => VariantBasicType::Object,
+        3 => VariantBasicType::Array,
+        _ => {
+            //NOTE:  A 2-bit value has a max of 4 different values (0-3), hence this is unreachable as we
+            // masked `basic_type` with 0x03 above.
+            unreachable!();
+        }
+    };
+    Ok(basic_type)
+}
+
+impl TryFrom<u8> for VariantPrimitiveType {
+    type Error = ArrowError;
+
+    fn try_from(value: u8) -> Result<Self, Self::Error> {
+        match value {
+            0 => Ok(VariantPrimitiveType::Null),
+            1 => Ok(VariantPrimitiveType::BooleanTrue),
+            2 => Ok(VariantPrimitiveType::BooleanFalse),
+            3 => Ok(VariantPrimitiveType::Int8),
+            // TODO: Add types for the rest, once API is agreed upon
+            16 => Ok(VariantPrimitiveType::String),
+            _ => Err(ArrowError::InvalidArgumentError(format!(
+                "unknown primitive type: {}",
+                value
+            ))),
+        }
+    }
+}
+/// Extract the primitive type from a Variant value-header byte
+pub(crate) fn get_primitive_type(header: u8) -> Result<VariantPrimitiveType, ArrowError> {
+    // last 6 bits contain the primitive-type, see spec
+    VariantPrimitiveType::try_from(header >> 2)
+}
+
+/// To be used in `map_err` when unpacking an integer from a slice of bytes.
+fn map_try_from_slice_error(e: TryFromSliceError) -> ArrowError {
+    ArrowError::InvalidArgumentError(e.to_string())
+}
+
+/// Decodes an Int8 from the value section of a variant.
+pub(crate) fn decode_int8(value: &[u8]) -> Result<i8, ArrowError> {
+    let value = i8::from_le_bytes(array_from_slice(value, 1)?);
+    Ok(value)
+}
+
+/// Decodes a long string from the value section of a variant.
+pub(crate) fn decode_long_string(value: &[u8]) -> Result<&str, ArrowError> {
+    let len = u32::from_le_bytes(array_from_slice(value, 1)?) as usize;
+    let string = string_from_slice(value, 5..5 + len)?;
+    Ok(string)
+}
+
+/// Decodes a short string from the value section of a variant.
+pub(crate) fn decode_short_string(value: &[u8]) -> Result<&str, ArrowError> {
+    let len = (first_byte_from_slice(value)? >> 2) as usize;
+
+    let string = string_from_slice(value, 1..1 + len)?;
+    Ok(string)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_i8() -> Result<(), ArrowError> {
+        let value = [
+            0 | 3 << 2, // Primitive type for i8
+            42,
+        ];
+        let result = decode_int8(&value)?;
+        assert_eq!(result, 42);
+        Ok(())
+    }
+
+    #[test]
+    fn test_short_string() -> Result<(), ArrowError> {
+        let value = [
+            1 | 5 << 2, // Basic type for short string | length of short string
+            'H' as u8,
+            'e' as u8,
+            'l' as u8,
+            'l' as u8,
+            'o' as u8,
+            'o' as u8,
+        ];
+        let result = decode_short_string(&value)?;
+        assert_eq!(result, "Hello");
+        Ok(())
+    }
+
+    #[test]
+    fn test_string() -> Result<(), ArrowError> {
+        let value = [
+            0 | 16 << 2, // Basic type for short string | length of short string
+            5,
+            0,
+            0,
+            0, // Length of string
+            'H' as u8,
+            'e' as u8,
+            'l' as u8,
+            'l' as u8,
+            'o' as u8,
+            'o' as u8,
+        ];
+        let result = decode_long_string(&value)?;
+        assert_eq!(result, "Hello");
+        Ok(())
+    }
+}

parquet-variant/src/lib.rs

@@ -26,3 +26,16 @@
 //! If you are interested in helping, you can find more information on the GitHub [Variant issue]
 //!
 //! [Variant issue]: https://github.com/apache/arrow-rs/issues/6736
+
+// TODO: dead code removal
+#[allow(dead_code)]
+mod decoder;
+// TODO: dead code removal
+#[allow(dead_code)]
+mod variant;
+// TODO: dead code removal
+#[allow(dead_code)]
+mod utils;
+
+#[cfg(test)]
+mod test_variant;

parquet-variant/src/test_variant.rs

@@ -0,0 +1,100 @@
+// 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.
+
+//! End-to-end check: (almost) every sample from apache/parquet-testing/variant
+//! can be parsed into our `Variant`.
+
+// NOTE: We keep this file separate rather than a test mod inside variant.rs because it should be
+// moved to the test folder later
+use std::fs;
+use std::path::{Path, PathBuf};
+
+use crate::variant::{Variant, VariantMetadata};
+use arrow_schema::ArrowError;
+
+fn cases_dir() -> PathBuf {
+    Path::new(env!("CARGO_MANIFEST_DIR"))
+        .join("..")
+        .join("parquet-testing")
+        .join("variant")
+}
+
+fn load_case(name: &str) -> Result<(Vec<u8>, Vec<u8>), ArrowError> {
+    let root = cases_dir();
+    let meta = fs::read(root.join(format!("{name}.metadata")))?;
+    let val = fs::read(root.join(format!("{name}.value")))?;
+    Ok((meta, val))
+}
+
+fn get_primitive_cases() -> Vec<(&'static str, Variant<'static, 'static>)> {
+    vec![
+    ("primitive_boolean_false", Variant::BooleanFalse),
+    ("primitive_boolean_true", Variant::BooleanTrue),
+    ("primitive_int8", Variant::Int8(42)),
+    // Using the From<String> trait
+    ("primitive_string", Variant::from("This string is longer than 64 bytes and therefore does not fit in a short_string and it also includes several non ascii characters such as 🐢, 💖, ♥\u{fe0f}, 🎣 and 🤦!!")),
+    // Using the From<String> trait
+    ("short_string", Variant::from("Less than 64 bytes (❤\u{fe0f} with utf8)")), 
+    // TODO Reenable when https://github.com/apache/parquet-testing/issues/81 is fixed
+    // ("primitive_null", Variant::Null),
+    ]
+}
+
+fn get_non_primitive_cases() -> Vec<&'static str> {
+    vec!["object_primitive", "array_primitive"]
+}
+
+#[test]
+fn variant_primitive() -> Result<(), ArrowError> {
+    let cases = get_primitive_cases();
+    for (case, want) in cases {
+        let (metadata_bytes, value) = load_case(case)?;
+        let metadata = VariantMetadata::try_new(&metadata_bytes)?;
+        let got = Variant::try_new(&metadata, &value)?;
+        assert_eq!(got, want);
+    }
+    Ok(())
+}
+
+#[test]
+fn variant_non_primitive() -> Result<(), ArrowError> {
+    let cases = get_non_primitive_cases();
+    for case in cases {
+        let (metadata, value) = load_case(case)?;
+        let metadata = VariantMetadata::try_new(&metadata)?;
+        let variant = Variant::try_new(&metadata, &value)?;
+        match case {
+            "object_primitive" => {
+                assert!(matches!(variant, Variant::Object(_)));
+                assert_eq!(metadata.dictionary_size(), 7);
+                let dict_val = metadata.get_field_by(0)?;
+                assert_eq!(dict_val, "int_field");
+            }
+            "array_primitive" => match variant {
+                Variant::Array(arr) => {
+                    let v = arr.get(0)?;
+                    assert!(matches!(v, Variant::Int8(2)));
+                    let v = arr.get(1)?;
+                    assert!(matches!(v, Variant::Int8(1)));
+                }
+                _ => panic!("expected an array"),
+            },
+            _ => unreachable!(),
+        }
+    }
+    Ok(())
+}

parquet-variant/src/utils.rs

@@ -0,0 +1,60 @@
+// 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.
+use std::{array::TryFromSliceError, ops::Range, str};
+
+use arrow_schema::ArrowError;
+
+use std::fmt::Debug;
+use std::slice::SliceIndex;
+
+#[inline]
+
+pub(crate) fn slice_from_slice<I: SliceIndex<[u8]> + Clone + Debug>(
+    bytes: &[u8],
+    index: I,
+) -> Result<&I::Output, ArrowError> {
+    bytes.get(index.clone()).ok_or_else(|| {
+        ArrowError::InvalidArgumentError(format!(
+            "Tried to extract byte(s) {index:?} from {}-byte buffer",
+            bytes.len(),
+        ))
+    })
+}
+pub(crate) fn array_from_slice<const N: usize>(
+    bytes: &[u8],
+    offset: usize,
+) -> Result<[u8; N], ArrowError> {
+    let bytes = slice_from_slice(bytes, offset..offset + N)?;
+    bytes.try_into().map_err(map_try_from_slice_error)
+}
+
+/// To be used in `map_err` when unpacking an integer from a slice of bytes.
+pub(crate) fn map_try_from_slice_error(e: TryFromSliceError) -> ArrowError {
+    ArrowError::InvalidArgumentError(e.to_string())
+}
+
+pub(crate) fn first_byte_from_slice(slice: &[u8]) -> Result<&u8, ArrowError> {
+    slice
+        .get(0)
+        .ok_or_else(|| ArrowError::InvalidArgumentError("Received empty bytes".to_string()))
+}
+
+/// Helper to get a &str from a slice based on range, if it's valid or an error otherwise
+pub(crate) fn string_from_slice(slice: &[u8], range: Range<usize>) -> Result<&str, ArrowError> {
+    str::from_utf8(slice_from_slice(slice, range)?)
+        .map_err(|_| ArrowError::InvalidArgumentError("invalid UTF-8 string".to_string()))
+}