[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,149 @@
+// NOTE: Largely based on the implementation of @PinkCrow007 in https://github.com/apache/arrow-rs/pull/7452
+// And the feedback there.
+use arrow_schema::ArrowError;
+use std::{array::TryFromSliceError, str};
+
+use crate::utils::{array_from_slice, invalid_utf8_err, non_empty_slice, slice_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,
+        _ => {
+            return Err(ArrowError::InvalidArgumentError(format!(
+                "unknown basic type: {}",
+                basic_type
+            )))
+        }
+    };
+    Ok(basic_type)
+}
+
+/// Extracts the primitive type from a header byte
+pub(crate) fn get_primitive_type(header: u8) -> Result<VariantPrimitiveType, ArrowError> {
+    // See https://github.com/apache/parquet-format/blob/master/VariantEncoding.md#value-encoding
+    //// Primitive type is encoded in the last 6 bits of the header byte
+    let primitive_type = header >> 2;
+    let primitive_type = match primitive_type {
+        0 => VariantPrimitiveType::Null,
+        1 => VariantPrimitiveType::BooleanTrue,
+        2 => VariantPrimitiveType::BooleanFalse,
+        3 => VariantPrimitiveType::Int8,
+        // TODO: Add types for the rest, once API is agreed upon
+        16 => VariantPrimitiveType::String,
+        _ => {
+            return Err(ArrowError::InvalidArgumentError(format!(
+                "unknown primitive type: {}",
+                primitive_type
+            )))
+        }
+    };
+    Ok(primitive_type)
+}
+
+/// 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(
+        slice_from_slice(value, 1..5)?
+            .try_into()
+            .map_err(map_try_from_slice_error)?,
+    ) as usize;
+    let string =
+        str::from_utf8(slice_from_slice(value, 5..5 + len)?).map_err(|_| invalid_utf8_err())?;
+    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 = (non_empty_slice(value)?[0] >> 2) as usize;
+
+    let string =
+        str::from_utf8(slice_from_slice(value, 1..1 + len)?).map_err(|_| invalid_utf8_err())?;
+    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,72 @@
+//! 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"]
+}
+
+#[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);
+            }
+            _ => unreachable!(),
+        }
+    }
+    Ok(())
+}

parquet-variant/src/utils.rs

@@ -0,0 +1,42 @@
+use std::{array::TryFromSliceError, ops::Range};
+
+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 non_empty_slice(slice: &[u8]) -> Result<&[u8], ArrowError> {
+    if slice.is_empty() {
+        return Err(ArrowError::InvalidArgumentError(
+            "Received empty bytes".to_string(),
+        ));
+    } else {
+        return Ok(slice);
+    }
+}
+
+/// Constructs the error message for an invalid UTF-8 string.
+pub(crate) fn invalid_utf8_err() -> ArrowError {
+    ArrowError::InvalidArgumentError("invalid UTF-8 string".to_string())
+}