Extract RecordBatchExt & friends out of re_redap_tests (#12039)

teh-cmc · web-flow · commit 72a788d51101 · 2025-12-02T11:43:21.000+01:00
* Extract `RecordBatchExt`, `SchemaExt` & `FieldsExt` out of `re_redap_tests` and into `re_arrow_util` instead * So that they can be used without having to depend on `re_redap_tests`, which is a guaranteed dependency cycle * Rename them to `RecordBatchTestExt`, `SchemaTestExt` & `FieldsTestExt` respectively * A) So they don't conflict with the other `RecordBatchExt` etc traits that we have * B) To communicate clearly that these traits will not shy away from unwraps() and other shortcuts --- This is a pre-requisite for the upcoming RRD footer PRs. * Sibling: rerun-io/dataplatform#2090
diff --git a/Cargo.lock b/Cargo.lock
@@ -8423,6 +8423,7 @@ dependencies = [
  "anyhow",
  "arrow",
  "comfy-table",
+ "datafusion",
  "half",
  "insta",
  "itertools 0.14.0",
@@ -9401,7 +9402,6 @@ dependencies = [
  "anyhow",
  "arrow",
  "chrono",
- "datafusion",
  "futures",
  "insta",
  "itertools 0.14.0",
diff --git a/crates/store/re_redap_tests/Cargo.toml b/crates/store/re_redap_tests/Cargo.toml
@@ -36,7 +36,6 @@ re_types.workspace = true
 
 anyhow.workspace = true
 arrow.workspace = true
-datafusion.workspace = true
 futures.workspace = true
 insta = { workspace = true, features = ["filters"] }
 itertools.workspace = true
diff --git a/crates/store/re_redap_tests/src/lib.rs b/crates/store/re_redap_tests/src/lib.rs
@@ -21,7 +21,6 @@ mod tests;
 mod utils;
 
 pub use self::utils::{
-    arrow::{FieldsExt, RecordBatchExt, SchemaExt},
     path::TempPath,
     rerun::{
         TuidPrefix, create_minimal_binary_recording_in, create_nasty_recording,
@@ -32,3 +31,5 @@ pub use self::utils::{
 };
 
 pub use self::tests::*;
+
+pub use ::re_arrow_util::{FieldsTestExt, RecordBatchTestExt, SchemaTestExt}; // convenience
diff --git a/crates/store/re_redap_tests/src/tests/common.rs b/crates/store/re_redap_tests/src/tests/common.rs
@@ -18,7 +18,7 @@ use re_protos::{
 use re_types_core::AsComponents;
 
 use crate::{
-    RecordBatchExt as _, TempPath, TuidPrefix, create_nasty_recording,
+    RecordBatchTestExt as _, TempPath, TuidPrefix, create_nasty_recording,
     create_recording_with_embeddings, create_recording_with_properties,
     create_recording_with_scalars, create_recording_with_text, create_simple_recording,
 };
diff --git a/crates/store/re_redap_tests/src/tests/create_table.rs b/crates/store/re_redap_tests/src/tests/create_table.rs
@@ -1,4 +1,4 @@
-use crate::SchemaExt as _;
+use crate::SchemaTestExt as _;
 use arrow::datatypes::{DataType, Field, Schema};
 use re_protos::cloud::v1alpha1::GetTableSchemaRequest;
 use re_protos::cloud::v1alpha1::ext::{
diff --git a/crates/store/re_redap_tests/src/tests/dataset_schema.rs b/crates/store/re_redap_tests/src/tests/dataset_schema.rs
@@ -4,7 +4,7 @@ use re_protos::{
 };
 
 use super::common::{DataSourcesDefinition, LayerDefinition, RerunCloudServiceExt as _};
-use crate::SchemaExt as _;
+use crate::SchemaTestExt as _;
 
 pub async fn simple_dataset_schema(service: impl RerunCloudService) {
     let data_sources_def = DataSourcesDefinition::new_with_tuid_prefix(
diff --git a/crates/store/re_redap_tests/src/tests/entries_table.rs b/crates/store/re_redap_tests/src/tests/entries_table.rs
@@ -8,7 +8,7 @@ use re_protos::cloud::v1alpha1::{
 };
 
 use crate::tests::common::RerunCloudServiceExt as _;
-use crate::{RecordBatchExt as _, SchemaExt as _};
+use crate::{RecordBatchTestExt as _, SchemaTestExt as _};
 
 /// We want to make sure that the "__entries" table is present and has the expected schema and data.
 pub async fn list_entries_table(service: impl RerunCloudService) {
diff --git a/crates/store/re_redap_tests/src/tests/fetch_chunks.rs b/crates/store/re_redap_tests/src/tests/fetch_chunks.rs
@@ -14,7 +14,7 @@ use re_sdk::external::re_log_encoding::ToApplication as _;
 use re_tuid::Tuid;
 use re_types_core::Loggable as _;
 
-use crate::RecordBatchExt as _;
+use crate::RecordBatchTestExt as _;
 use crate::tests::common::{
     DataSourcesDefinition, LayerDefinition, RerunCloudServiceExt as _, concat_record_batches,
 };
diff --git a/crates/store/re_redap_tests/src/tests/query_dataset.rs b/crates/store/re_redap_tests/src/tests/query_dataset.rs
@@ -11,7 +11,7 @@ use re_protos::{
 use crate::tests::common::{
     DataSourcesDefinition, LayerDefinition, RerunCloudServiceExt as _, concat_record_batches,
 };
-use crate::{FieldsExt as _, RecordBatchExt as _};
+use crate::{FieldsTestExt as _, RecordBatchTestExt as _};
 
 pub async fn query_empty_dataset(service: impl RerunCloudService) {
     let dataset_name = "dataset";
diff --git a/crates/store/re_redap_tests/src/tests/register_segment.rs b/crates/store/re_redap_tests/src/tests/register_segment.rs
@@ -18,7 +18,9 @@ use re_protos::{
 };
 
 use super::common::{DataSourcesDefinition, LayerDefinition, RerunCloudServiceExt as _, prop};
-use crate::{FieldsExt as _, RecordBatchExt as _, SchemaExt as _, create_simple_recording_in};
+use crate::{
+    FieldsTestExt as _, RecordBatchTestExt as _, SchemaTestExt as _, create_simple_recording_in,
+};
 
 pub async fn register_and_scan_simple_dataset(service: impl RerunCloudService) {
     let data_sources_def = DataSourcesDefinition::new_with_tuid_prefix(
diff --git a/crates/store/re_redap_tests/src/tests/write_table.rs b/crates/store/re_redap_tests/src/tests/write_table.rs
@@ -1,4 +1,4 @@
-use crate::RecordBatchExt as _;
+use crate::RecordBatchTestExt as _;
 use crate::tests::common::{RerunCloudServiceExt as _, concat_record_batches};
 use crate::utils::streaming::make_streaming_request;
 use crate::utils::tables::create_simple_lance_dataset;
diff --git a/crates/store/re_redap_tests/src/utils/mod.rs b/crates/store/re_redap_tests/src/utils/mod.rs
@@ -1,4 +1,3 @@
-pub mod arrow;
 pub mod path;
 pub mod rerun;
 pub mod streaming;
diff --git a/crates/utils/re_arrow_util/Cargo.toml b/crates/utils/re_arrow_util/Cargo.toml
@@ -29,6 +29,7 @@ re_tuid = { workspace = true, features = ["bytemuck"] }
 
 anyhow.workspace = true
 arrow.workspace = true
+datafusion.workspace = true
 half.workspace = true
 itertools.workspace = true
 serde_json.workspace = true
diff --git a/crates/utils/re_arrow_util/src/lib.rs b/crates/utils/re_arrow_util/src/lib.rs
@@ -5,6 +5,7 @@ mod batches;
 mod compare;
 mod format;
 mod format_data_type;
+mod test_extensions;
 
 pub use self::arrays::*;
 pub use self::batches::*;
@@ -14,6 +15,7 @@ pub use self::format::{
     format_record_batch_with_width,
 };
 pub use self::format_data_type::*;
+pub use self::test_extensions::*;
 
 // ----------------------------------------------------------------
 
diff --git a/crates/utils/re_arrow_util/src/test_extensions.rs b/crates/utils/re_arrow_util/src/test_extensions.rs
@@ -2,7 +2,7 @@ use std::collections::HashSet;
 use std::sync::Arc;
 
 use arrow::{
-    array::{ArrayRef, ListArray, StringArray},
+    array::{Array as _, ArrayRef, ListArray, StringArray},
     compute::SortOptions,
     datatypes::{DataType, Field, Fields},
 };
@@ -12,12 +12,11 @@ use datafusion::{
 };
 use itertools::Itertools as _;
 
-use re_arrow_util::{ArrowArrayDowncastRef as _, RecordBatchExt as _};
-use re_chunk::ArrowArray as _;
+use crate::{ArrowArrayDowncastRef as _, RecordBatchExt as _};
 
 // --
 
-pub trait RecordBatchExt {
+pub trait RecordBatchTestExt {
     /// Formats a record batch in a snapshot-friendly way.
     fn format_snapshot(&self, transposed: bool) -> String;
 
@@ -81,11 +80,11 @@ pub trait RecordBatchExt {
     fn remove_columns(&self, columns: &[&str]) -> Self;
 }
 
-impl RecordBatchExt for arrow::array::RecordBatch {
+impl RecordBatchTestExt for arrow::array::RecordBatch {
     fn format_snapshot(&self, transposed: bool) -> String {
-        re_arrow_util::format_record_batch_opts(
+        crate::format_record_batch_opts(
             self,
-            &re_arrow_util::RecordBatchFormatOpts {
+            &crate::RecordBatchFormatOpts {
                 transposed,
                 width: Some(800),
                 include_metadata: false,
@@ -176,6 +175,7 @@ impl RecordBatchExt for arrow::array::RecordBatch {
         }
 
         let schema = arrow::datatypes::Schema::new_with_metadata(fields, schema.metadata().clone());
+        #[cfg_attr(not(target_arch = "wasm32"), expect(clippy::disallowed_methods))] // tests
         Some(Self::try_new(Arc::new(schema), arrays).expect("creating record batch"))
     }
 
@@ -210,6 +210,7 @@ impl RecordBatchExt for arrow::array::RecordBatch {
         if schema.fields().is_empty() {
             Self::new_empty(schema)
         } else {
+            #[cfg_attr(not(target_arch = "wasm32"), expect(clippy::disallowed_methods))] // tests
             Self::try_new(schema, arrays).expect("creation should succeed")
         }
     }
@@ -309,6 +310,7 @@ impl RecordBatchExt for arrow::array::RecordBatch {
         if schema.fields().is_empty() {
             Self::new_empty(schema.clone())
         } else {
+            #[cfg_attr(not(target_arch = "wasm32"), expect(clippy::disallowed_methods))] // tests
             Self::try_new(schema.clone(), arrays).expect("creation should succeed")
         }
     }
@@ -338,12 +340,12 @@ impl RecordBatchExt for arrow::array::RecordBatch {
     }
 }
 
-pub trait SchemaExt {
+pub trait SchemaTestExt {
     /// Formats a record batch in a snapshot-friendly way.
     fn format_snapshot(&self) -> String;
 }
 
-impl SchemaExt for arrow::datatypes::Schema {
+impl SchemaTestExt for arrow::datatypes::Schema {
     fn format_snapshot(&self) -> String {
         let metadata = (!self.metadata().is_empty()).then(|| {
             format!(
@@ -363,13 +365,13 @@ impl SchemaExt for arrow::datatypes::Schema {
                 format!(
                     "{}: {}",
                     field.name(),
-                    re_arrow_util::format_data_type(field.data_type())
+                    crate::format_data_type(field.data_type())
                 )
             } else {
                 format!(
                     "{}: {} [\n    {}\n]",
                     field.name(),
-                    re_arrow_util::format_data_type(field.data_type()),
+                    crate::format_data_type(field.data_type()),
                     field
                         .metadata()
                         .iter()
@@ -384,15 +386,15 @@ impl SchemaExt for arrow::datatypes::Schema {
     }
 }
 
-pub trait FieldsExt {
+pub trait FieldsTestExt {
     /// Returns true if all the required fields are present, regardless of the order.
     fn contains_unordered(
         &self,
         required_fields: impl IntoIterator<Item = impl AsRef<Field>>,
     ) -> bool;
 }
 
-impl FieldsExt for Fields {
+impl FieldsTestExt for Fields {
     fn contains_unordered(
         &self,
         required_fields: impl IntoIterator<Item = impl AsRef<Field>>,

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-use crate::SchemaExt as _;`
	`1`	`+use crate::SchemaTestExt as _;`
`2`	`2`	`use arrow::datatypes::{DataType, Field, Schema};`
`3`	`3`	`use re_protos::cloud::v1alpha1::GetTableSchemaRequest;`
`4`	`4`	`use re_protos::cloud::v1alpha1::ext::{`