impl FromAvro/ToAvro for all encodings

spiraldb · Jan 8, 2025 · 884c4f1 · 884c4f1
1 parent 26a40aa
commit 884c4f1
Show file tree

Hide file tree

Showing 34 changed files with 301 additions and 163 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/encodings/bytebool/Cargo.toml b/encodings/bytebool/Cargo.toml
@@ -21,6 +21,7 @@ arrow-buffer = { workspace = true }
 num-traits = { workspace = true }
 serde = { workspace = true, features = ["derive"] }
 vortex-array = { workspace = true }
+vortex-avro = { workspace = true }
 vortex-buffer = { workspace = true }
 vortex-dtype = { workspace = true }
 vortex-error = { workspace = true }

diff --git a/encodings/bytebool/src/array.rs b/encodings/bytebool/src/array.rs
@@ -10,13 +10,13 @@ use vortex_array::validity::{LogicalValidity, Validity, ValidityMetadata, Validi
 use vortex_array::variants::{BoolArrayTrait, VariantsVTable};
 use vortex_array::visitor::{ArrayVisitor, VisitorVTable};
 use vortex_array::{impl_encoding, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoCanonical};
+use vortex_avro::{FromAvro, ToAvro};
 use vortex_buffer::ByteBuffer;
 use vortex_dtype::DType;
 use vortex_error::{VortexExpect as _, VortexResult};
-
 impl_encoding!("vortex.bytebool", ids::BYTE_BOOL, ByteBool);
 
-#[derive(Clone, Debug, Serialize, Deserialize)]
+#[derive(Clone, Debug, Serialize, Deserialize, FromAvro, ToAvro)]
 pub struct ByteBoolMetadata {
     validity: ValidityMetadata,
 }

diff --git a/encodings/datetime-parts/Cargo.toml b/encodings/datetime-parts/Cargo.toml
@@ -19,6 +19,7 @@ workspace = true
 [dependencies]
 serde = { workspace = true, features = ["derive"] }
 vortex-array = { workspace = true }
+vortex-avro = { workspace = true }
 vortex-buffer = { workspace = true }
 vortex-datetime-dtype = { workspace = true }
 vortex-dtype = { workspace = true }

diff --git a/encodings/datetime-parts/src/array.rs b/encodings/datetime-parts/src/array.rs
@@ -12,14 +12,14 @@ use vortex_array::{
     impl_encoding, ArrayDType, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoArrayData,
     IntoCanonical,
 };
+use vortex_avro::{FromAvro, ToAvro};
 use vortex_dtype::{DType, PType};
 use vortex_error::{vortex_bail, VortexExpect as _, VortexResult, VortexUnwrap};
 
 use crate::compute::decode_to_temporal;
-
 impl_encoding!("vortex.datetimeparts", ids::DATE_TIME_PARTS, DateTimeParts);
 
-#[derive(Clone, Debug, Serialize, Deserialize)]
+#[derive(Clone, Debug, Serialize, Deserialize, FromAvro, ToAvro)]
 pub struct DateTimePartsMetadata {
     // Validity lives in the days array
     // TODO(ngates): we should actually model this with a Tuple array when we have one.

diff --git a/encodings/dict/Cargo.toml b/encodings/dict/Cargo.toml
@@ -19,6 +19,7 @@ hashbrown = { workspace = true }
 num-traits = { workspace = true }
 serde = { workspace = true }
 vortex-array = { workspace = true }
+vortex-avro = { workspace = true }
 vortex-buffer = { workspace = true }
 vortex-dtype = { workspace = true }
 vortex-error = { workspace = true }

diff --git a/encodings/dict/src/array.rs b/encodings/dict/src/array.rs
@@ -13,15 +13,15 @@ use vortex_array::{
     impl_encoding, ArrayDType, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoArrayData,
     IntoArrayVariant, IntoCanonical,
 };
+use vortex_avro::{FromAvro, ToAvro};
 use vortex_dtype::{match_each_integer_ptype, DType, PType};
 use vortex_error::{vortex_bail, vortex_panic, VortexExpect as _, VortexResult};
-
 impl_encoding!("vortex.dict", ids::DICT, Dict);
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, FromAvro, ToAvro)]
 pub struct DictMetadata {
     codes_ptype: PType,
-    values_len: usize,
+    values_len: u64,
 }
 
 impl Display for DictMetadata {
@@ -41,7 +41,7 @@ impl DictArray {
             DictMetadata {
                 codes_ptype: PType::try_from(codes.dtype())
                     .vortex_expect("codes dtype must be uint"),
-                values_len: values.len(),
+                values_len: values.len() as u64,
             },
             [codes, values].into(),
             StatsSet::default(),
@@ -58,7 +58,7 @@ impl DictArray {
     #[inline]
     pub fn values(&self) -> ArrayData {
         self.as_ref()
-            .child(1, self.dtype(), self.metadata().values_len)
+            .child(1, self.dtype(), self.metadata().values_len as usize)
             .vortex_expect("DictArray is missing its values child array")
     }
 }
@@ -135,7 +135,7 @@ mod test {
             "dict.metadata",
             DictMetadata {
                 codes_ptype: PType::U64,
-                values_len: usize::MAX,
+                values_len: u64::MAX,
             },
         );
     }

diff --git a/encodings/fastlanes/Cargo.toml b/encodings/fastlanes/Cargo.toml
@@ -24,6 +24,7 @@ itertools = { workspace = true }
 num-traits = { workspace = true }
 serde = { workspace = true }
 vortex-array = { workspace = true }
+vortex-avro = { workspace = true }
 vortex-buffer = { workspace = true }
 vortex-dtype = { workspace = true }
 vortex-error = { workspace = true }

diff --git a/encodings/fastlanes/src/bitpacking/mod.rs b/encodings/fastlanes/src/bitpacking/mod.rs
@@ -14,16 +14,16 @@ use vortex_array::visitor::{ArrayVisitor, VisitorVTable};
 use vortex_array::{
     impl_encoding, ArrayDType, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoCanonical,
 };
+use vortex_avro::{FromAvro, ToAvro};
 use vortex_buffer::ByteBuffer;
 use vortex_dtype::{DType, NativePType, PType};
 use vortex_error::{vortex_bail, vortex_err, VortexExpect as _, VortexResult};
-
 mod compress;
 mod compute;
 
 impl_encoding!("fastlanes.bitpacked", ids::FL_BITPACKED, BitPacked);
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, FromAvro, ToAvro)]
 pub struct BitPackedMetadata {
     validity: ValidityMetadata,
     bit_width: u8,

diff --git a/encodings/fastlanes/src/delta/mod.rs b/encodings/fastlanes/src/delta/mod.rs
@@ -12,6 +12,7 @@ use vortex_array::{
     impl_encoding, ArrayDType, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoArrayData,
     IntoCanonical,
 };
+use vortex_avro::{FromAvro, ToAvro};
 use vortex_buffer::Buffer;
 use vortex_dtype::{match_each_unsigned_integer_ptype, NativePType};
 use vortex_error::{vortex_bail, vortex_panic, VortexExpect as _, VortexResult};
@@ -21,7 +22,7 @@ mod compute;
 
 impl_encoding!("fastlanes.delta", ids::FL_DELTA, Delta);
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, FromAvro, ToAvro)]
 pub struct DeltaMetadata {
     validity: ValidityMetadata,
     deltas_len: u64,

diff --git a/encodings/fastlanes/src/for/mod.rs b/encodings/fastlanes/src/for/mod.rs
@@ -10,16 +10,16 @@ use vortex_array::visitor::{ArrayVisitor, VisitorVTable};
 use vortex_array::{
     impl_encoding, ArrayDType, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoCanonical,
 };
+use vortex_avro::{FromAvro, ToAvro};
 use vortex_dtype::DType;
 use vortex_error::{vortex_bail, VortexExpect as _, VortexResult};
 use vortex_scalar::{PValue, Scalar};
-
 mod compress;
 mod compute;
 
 impl_encoding!("fastlanes.for", ids::FL_FOR, FoR);
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, FromAvro, ToAvro)]
 pub struct FoRMetadata {
     reference: PValue,
     shift: u8,

diff --git a/encodings/fsst/Cargo.toml b/encodings/fsst/Cargo.toml
@@ -22,6 +22,7 @@ fsst-rs = { workspace = true }
 serde = { workspace = true }
 
 vortex-array = { workspace = true }
+vortex-avro = { workspace = true }
 vortex-buffer = { workspace = true }
 vortex-dtype = { workspace = true }
 vortex-error = { workspace = true }

diff --git a/encodings/fsst/src/array.rs b/encodings/fsst/src/array.rs
@@ -10,6 +10,7 @@ use vortex_array::validity::{ArrayValidity, LogicalValidity, Validity, ValidityV
 use vortex_array::variants::{BinaryArrayTrait, Utf8ArrayTrait, VariantsVTable};
 use vortex_array::visitor::{ArrayVisitor, VisitorVTable};
 use vortex_array::{impl_encoding, ArrayDType, ArrayData, ArrayLen, ArrayTrait, IntoCanonical};
+use vortex_avro::{FromAvro, ToAvro};
 use vortex_dtype::{DType, Nullability, PType};
 use vortex_error::{vortex_bail, VortexExpect, VortexResult};
 
@@ -18,9 +19,9 @@ impl_encoding!("vortex.fsst", ids::FSST, FSST);
 static SYMBOLS_DTYPE: DType = DType::Primitive(PType::U64, Nullability::NonNullable);
 static SYMBOL_LENS_DTYPE: DType = DType::Primitive(PType::U8, Nullability::NonNullable);
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, FromAvro, ToAvro)]
 pub struct FSSTMetadata {
-    symbols_len: usize,
+    symbols_len: u64,
     codes_nullability: Nullability,
     uncompressed_lengths_ptype: PType,
 }
@@ -95,7 +96,7 @@ impl FSSTArray {
             dtype,
             len,
             FSSTMetadata {
-                symbols_len,
+                symbols_len: symbols_len as u64,
                 codes_nullability,
                 uncompressed_lengths_ptype,
             },
@@ -107,14 +108,14 @@ impl FSSTArray {
     /// Access the symbol table array
     pub fn symbols(&self) -> ArrayData {
         self.as_ref()
-            .child(0, &SYMBOLS_DTYPE, self.metadata().symbols_len)
+            .child(0, &SYMBOLS_DTYPE, self.metadata().symbols_len as usize)
             .vortex_expect("FSSTArray symbols child")
     }
 
     /// Access the symbol table array
     pub fn symbol_lengths(&self) -> ArrayData {
         self.as_ref()
-            .child(1, &SYMBOL_LENS_DTYPE, self.metadata().symbols_len)
+            .child(1, &SYMBOL_LENS_DTYPE, self.metadata().symbols_len as usize)
             .vortex_expect("FSSTArray symbol_lengths child")
     }
 
@@ -239,7 +240,7 @@ mod test {
         check_metadata(
             "fsst.metadata",
             FSSTMetadata {
-                symbols_len: usize::MAX,
+                symbols_len: u64::MAX,
                 codes_nullability: Nullability::Nullable,
                 uncompressed_lengths_ptype: PType::U64,
             },

diff --git a/encodings/roaring/Cargo.toml b/encodings/roaring/Cargo.toml
@@ -19,6 +19,7 @@ croaring = { workspace = true }
 num-traits = { workspace = true }
 serde = { workspace = true }
 vortex-array = { workspace = true }
+vortex-avro = { workspace = true }
 vortex-buffer = { workspace = true }
 vortex-dtype = { workspace = true }
 vortex-error = { workspace = true }

diff --git a/encodings/roaring/src/boolean/mod.rs b/encodings/roaring/src/boolean/mod.rs
@@ -15,6 +15,7 @@ use vortex_array::visitor::{ArrayVisitor, VisitorVTable};
 use vortex_array::{
     impl_encoding, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoArrayData, IntoCanonical,
 };
+use vortex_avro::{FromAvro, ToAvro};
 use vortex_buffer::ByteBuffer;
 use vortex_dtype::{DType, Nullability};
 use vortex_error::{vortex_bail, vortex_err, VortexExpect as _, VortexResult};
@@ -25,7 +26,7 @@ mod stats;
 
 impl_encoding!("vortex.roaring_bool", ids::ROARING_BOOL, RoaringBool);
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, FromAvro, ToAvro)]
 pub struct RoaringBoolMetadata;
 
 impl Display for RoaringBoolMetadata {

diff --git a/encodings/roaring/src/integer/mod.rs b/encodings/roaring/src/integer/mod.rs
@@ -15,17 +15,17 @@ use vortex_array::{
     impl_encoding, ArrayDType as _, ArrayData, ArrayLen, ArrayTrait, Canonical, IntoArrayData,
     IntoCanonical,
 };
+use vortex_avro::{FromAvro, ToAvro};
 use vortex_buffer::{Buffer, ByteBuffer};
 use vortex_dtype::Nullability::NonNullable;
 use vortex_dtype::{DType, PType};
 use vortex_error::{vortex_bail, VortexExpect as _, VortexResult};
-
 mod compress;
 mod compute;
 
 impl_encoding!("vortex.roaring_int", ids::ROARING_INT, RoaringInt);
 
-#[derive(Debug, Clone, Serialize, Deserialize)]
+#[derive(Debug, Clone, Serialize, Deserialize, FromAvro, ToAvro)]
 pub struct RoaringIntMetadata {
     ptype: PType,
 }

diff --git a/encodings/runend-bool/Cargo.toml b/encodings/runend-bool/Cargo.toml
@@ -19,6 +19,7 @@ itertools = { workspace = true }
 num-traits = { workspace = true }
 serde = { workspace = true }
 vortex-array = { workspace = true }
+vortex-avro = { workspace = true }
 vortex-buffer = { workspace = true }
 vortex-dtype = { workspace = true }
 vortex-error = { workspace = true }