[Variant] Initial API for reading Variant data and metadata

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, slice_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,55 @@
+// 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;
+
+#[inline]
+pub(crate) fn slice_from_slice(bytes: &[u8], range: Range<usize>) -> Result<&[u8], ArrowError> {
+    bytes.get(range.clone()).ok_or_else(|| {
+        ArrowError::InvalidArgumentError(format!(
+            "Tried to extract {} bytes at offset {} from {}-byte buffer",
+            range.end - range.start,
+            range.start,
+            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()))
+}