Port codegened arrow serialization

rerun-io · Nov 22, 2024 · 2b31881 · 2b31881
1 parent 53571fd
commit 2b31881
Show file tree

Hide file tree

Showing 193 changed files with 3,759 additions and 2,895 deletions.
diff --git a/crates/build/re_types_builder/src/codegen/rust/api.rs b/crates/build/re_types_builder/src/codegen/rust/api.rs
@@ -958,13 +958,13 @@ fn quote_trait_impls_for_datatype_or_component(
 
     let quoted_serializer = if let Some(forwarded_type) = forwarded_type.as_ref() {
         quote! {
-            fn to_arrow2_opt<'a>(
+            fn to_arrow_opt<'a>(
                 data: impl IntoIterator<Item = Option<impl Into<::std::borrow::Cow<'a, Self>>>>,
-            ) -> SerializationResult<Box<dyn arrow2::array::Array>>
+            ) -> SerializationResult<arrow::array::ArrayRef>
             where
                 Self: Clone + 'a,
             {
-                #forwarded_type::to_arrow2_opt(data.into_iter().map(|datum| {
+                #forwarded_type::to_arrow_opt(data.into_iter().map(|datum| {
                     datum.map(|datum| match datum.into() {
                         ::std::borrow::Cow::Borrowed(datum) => ::std::borrow::Cow::Borrowed(&datum.0),
                         ::std::borrow::Cow::Owned(datum) => ::std::borrow::Cow::Owned(datum.0),
@@ -978,9 +978,9 @@ fn quote_trait_impls_for_datatype_or_component(
 
         quote! {
             // NOTE: Don't inline this, this gets _huge_.
-            fn to_arrow2_opt<'a>(
+            fn to_arrow_opt<'a>(
                 data: impl IntoIterator<Item = Option<impl Into<::std::borrow::Cow<'a, Self>>>>,
-            ) -> SerializationResult<Box<dyn arrow2::array::Array>>
+            ) -> SerializationResult<arrow::array::ArrayRef>
             where
                 Self: Clone + 'a
             {
@@ -989,10 +989,14 @@ fn quote_trait_impls_for_datatype_or_component(
 
                 #![allow(clippy::wildcard_imports)]
                 #![allow(clippy::manual_is_variant_and)]
-                use arrow::datatypes::*;
-                use arrow2::array::*;
+                use arrow::{array::*, buffer::*, datatypes::*};
                 use ::re_types_core::{Loggable as _, ResultExt as _};
 
+                #[allow(unused)]
+                fn as_array_ref<T: Array + 'static>(t: T) -> ArrayRef {
+                    std::sync::Arc::new(t) as ArrayRef
+                }
+
                 Ok(#quoted_serializer)
             }
         }

diff --git a/crates/build/re_types_builder/src/codegen/rust/arrow.rs b/crates/build/re_types_builder/src/codegen/rust/arrow.rs
@@ -164,3 +164,22 @@ pub fn is_backed_by_arrow_buffer(typ: &DataType) -> bool {
             | DataType::Float64
     )
 }
+
+pub fn quoted_arrow_primitive_type(datatype: &DataType) -> TokenStream {
+    match datatype {
+        DataType::Null => quote!(arrow::datatypes::NullType),
+        DataType::Boolean => quote!(arrow::datatypes::BooleanType),
+        DataType::Int8 => quote!(arrow::datatypes::Int8Type),
+        DataType::Int16 => quote!(arrow::datatypes::Int16Type),
+        DataType::Int32 => quote!(arrow::datatypes::Int32Type),
+        DataType::Int64 => quote!(arrow::datatypes::Int64Type),
+        DataType::UInt8 => quote!(arrow::datatypes::UInt8Type),
+        DataType::UInt16 => quote!(arrow::datatypes::UInt16Type),
+        DataType::UInt32 => quote!(arrow::datatypes::UInt32Type),
+        DataType::UInt64 => quote!(arrow::datatypes::UInt64Type),
+        DataType::Float16 => quote!(arrow::datatypes::Float16Type),
+        DataType::Float32 => quote!(arrow::datatypes::Float32Type),
+        DataType::Float64 => quote!(arrow::datatypes::Float64Type),
+        _ => unimplemented!("Not a primitive type: {datatype:#?}"),
+    }
+}