Fix wrong datatype displayed for non-list column in the table view (#11140)

Co-authored-by: Emil Ernerfeldt <[email protected]>

Author: abey79

crates/store/re_sorbet/src/migrations/mod.rs

@@ -10,7 +10,7 @@ use std::cmp::Ordering;
 
 use arrow::{array::RecordBatch, datatypes::SchemaRef};
 
-use crate::SorbetSchema;
+use crate::{BatchType, SorbetSchema};
 
 mod make_list_arrays;
 
@@ -105,11 +105,13 @@ fn maybe_apply<M: Migration>(
 
 /// Migrate a sorbet record batch of unknown version to the latest version.
 #[tracing::instrument(level = "debug", skip_all)]
-pub fn migrate_record_batch(mut batch: RecordBatch) -> RecordBatch {
+pub fn migrate_record_batch(mut batch: RecordBatch, batch_type: BatchType) -> RecordBatch {
     batch = migrate_record_batch_impl(batch);
 
-    // TODO(emilk): only call make_all_data_columns_list_arrays for chunk batches?
-    make_list_arrays::make_all_data_columns_list_arrays(&batch)
+    match batch_type {
+        BatchType::Chunk => make_list_arrays::make_all_data_columns_list_arrays(&batch),
+        BatchType::Dataframe => batch,
+    }
 }
 
 fn migrate_record_batch_impl(mut batch: RecordBatch) -> RecordBatch {

crates/store/re_sorbet/src/sorbet_batch.rs

@@ -178,7 +178,7 @@ impl SorbetBatch {
         re_tracing::profile_function!();
 
         // First migrate the incoming batch to the latest format:
-        let batch = crate::migrations::migrate_record_batch(batch.clone());
+        let batch = crate::migrations::migrate_record_batch(batch.clone(), batch_type);
 
         let sorbet_schema =
             SorbetSchema::try_from_migrated_arrow_schema(batch.schema_ref().as_ref())?;

crates/store/re_sorbet/src/sorbet_schema.rs

@@ -1,9 +1,11 @@
-use arrow::datatypes::Schema as ArrowSchema;
+use arrow::datatypes::{Schema as ArrowSchema, SchemaRef as ArrowSchemaRef};
 
 use re_log_types::EntityPath;
 use re_types_core::ChunkId;
 
-use crate::{ArrowBatchMetadata, SorbetColumnDescriptors, SorbetError, TimestampMetadata};
+use crate::{
+    ArrowBatchMetadata, SorbetColumnDescriptors, SorbetError, TimestampMetadata, migrate_schema_ref,
+};
 
 // ----------------------------------------------------------------------------
 
@@ -116,6 +118,11 @@ impl SorbetSchema {
 }
 
 impl SorbetSchema {
+    /// Parse an arbitrary arrow schema by first migrating it to the Rerun schema.
+    pub fn try_from_raw_arrow_schema(arrow_schema: ArrowSchemaRef) -> Result<Self, SorbetError> {
+        Self::try_from_migrated_arrow_schema(&migrate_schema_ref(arrow_schema))
+    }
+
     /// Parse an already migrated Arrow schema.
     #[tracing::instrument(level = "trace", skip_all)]
     pub(crate) fn try_from_migrated_arrow_schema(

crates/viewer/re_dataframe_ui/Cargo.toml

@@ -41,6 +41,7 @@ datafusion.workspace = true
 egui.workspace = true
 egui_dnd.workspace = true
 egui_table.workspace = true
+itertools.workspace = true
 nohash-hasher.workspace = true
 parking_lot.workspace = true
 serde.workspace = true

crates/viewer/re_dataframe_ui/src/datafusion_adapter.rs

@@ -1,6 +1,6 @@
 use std::sync::Arc;
 
-use arrow::datatypes::{DataType, Fields};
+use arrow::datatypes::{DataType, SchemaRef};
 use datafusion::common::{DataFusionError, TableReference};
 use datafusion::functions::expr_fn::concat;
 use datafusion::logical_expr::{binary_expr, col as datafusion_col, lit};
@@ -61,6 +61,19 @@ impl From<&TableBlueprint> for DataFusionQueryData {
     }
 }
 
+/// Result of the async datafusion query process.
+#[derive(Debug, Clone)]
+pub struct DataFusionQueryResult {
+    /// The record batches to display.
+    pub sorbet_batches: Vec<SorbetBatch>,
+
+    /// The schema of the record batches.
+    pub original_schema: SchemaRef,
+
+    /// The migrated schema of the record batches (useful when the list of batches is empty).
+    pub sorbet_schema: re_sorbet::SorbetSchema,
+}
+
 /// A table blueprint along with the context required to execute the corresponding datafusion query.
 #[derive(Clone)]
 struct DataFusionQuery {
@@ -83,12 +96,11 @@ impl DataFusionQuery {
         }
     }
 
-    /// Execute the query and produce a vector of [`SorbetBatch`]s along with physical columns
-    /// names.
+    /// Execute the query to produce the data to display.
     ///
     /// Note: the future returned by this function must be `'static`, so it takes `self`. Use
     /// `clone()` as required.
-    async fn execute(self) -> Result<(Vec<SorbetBatch>, Fields), DataFusionError> {
+    async fn execute(self) -> Result<DataFusionQueryResult, DataFusionError> {
         let mut dataframe = self.session_ctx.table(self.table_ref).await?;
 
         let DataFusionQueryData {
@@ -182,20 +194,16 @@ impl DataFusionQuery {
         }
 
         //
-        // Collect
+        // Collect record batches
         //
 
+        let original_schema = Arc::clone(dataframe.schema().inner());
         let record_batches = dataframe.collect().await?;
 
-        // TODO(#10421) IMPORTANT: fields must be copied here *before* converting to `SorbetBatch`,
-        // because that conversion modifies the field names. As a result, the schema contained in a
-        // `SorbetBatch` cannot be used to derive the physical column names as seen by DataFusion.
-        let fields = record_batches
-            .first()
-            .map(|record_batch| record_batch.schema().fields.clone())
-            .unwrap_or_default();
+        //
+        // Convert to `SorbetBatch`
+        //
 
-        // convert to SorbetBatch
         let sorbet_batches = record_batches
             .iter()
             .map(|record_batch| {
@@ -204,7 +212,23 @@ impl DataFusionQuery {
             .collect::<Result<Vec<_>, _>>()
             .map_err(|err| DataFusionError::External(err.into()))?;
 
-        Ok((sorbet_batches, fields))
+        //
+        // Get (or create) `SorbetSchema`
+        //
+
+        let sorbet_schema = sorbet_batches
+            .first()
+            .map(|batch| Ok(batch.sorbet_schema().clone()))
+            .unwrap_or_else(|| {
+                re_sorbet::SorbetSchema::try_from_raw_arrow_schema(Arc::clone(&original_schema))
+                    .map_err(|err| DataFusionError::External(err.into()))
+            })?;
+
+        Ok(DataFusionQueryResult {
+            sorbet_batches,
+            original_schema,
+            sorbet_schema,
+        })
     }
 }
 
@@ -222,8 +246,8 @@ impl PartialEq for DataFusionQuery {
     }
 }
 
-type RequestedSorbetBatches =
-    RequestedObject<Result<(Vec<SorbetBatch>, arrow::datatypes::Fields), DataFusionError>>;
+type RequestedDataFusionQueryResult =
+    RequestedObject<Result<DataFusionQueryResult, DataFusionError>>;
 
 /// Helper struct to manage the datafusion async query and the resulting `SorbetBatch`.
 #[derive(Clone)]
@@ -237,11 +261,11 @@ pub struct DataFusionAdapter {
     query: DataFusionQuery,
 
     // Used to have something to display while the new dataframe is being queried.
-    pub last_sorbet_batches: Option<(Vec<SorbetBatch>, Fields)>,
+    pub last_query_results: Option<DataFusionQueryResult>,
 
     // TODO(ab, lucasmerlin): this `Mutex` is only needed because of the `Clone` bound in egui
     // so we should clean that up if the bound is lifted.
-    pub requested_sorbet_batches: Arc<Mutex<RequestedSorbetBatches>>,
+    pub requested_query_result: Arc<Mutex<RequestedDataFusionQueryResult>>,
 
     pub queried_at: Timestamp,
 }
@@ -271,13 +295,13 @@ impl DataFusionAdapter {
             let table_state = Self {
                 id,
                 blueprint: initial_blueprint,
-                requested_sorbet_batches: Arc::new(Mutex::new(RequestedObject::new_with_repaint(
+                requested_query_result: Arc::new(Mutex::new(RequestedObject::new_with_repaint(
                     runtime,
                     ui.ctx().clone(),
                     query.clone().execute(),
                 ))),
                 query,
-                last_sorbet_batches: None,
+                last_query_results: None,
                 queried_at: Timestamp::now(),
             };
 
@@ -288,7 +312,7 @@ impl DataFusionAdapter {
             table_state
         });
 
-        adapter.requested_sorbet_batches.lock().on_frame_start();
+        adapter.requested_query_result.lock().on_frame_start();
 
         adapter
     }
@@ -314,10 +338,10 @@ impl DataFusionAdapter {
         if self.query.query_data != new_query_data {
             self.query.query_data = new_query_data;
 
-            let mut dataframe = self.requested_sorbet_batches.lock();
+            let mut dataframe = self.requested_query_result.lock();
 
             if let Some(Ok(sorbet_batches)) = dataframe.try_as_ref() {
-                self.last_sorbet_batches = Some(sorbet_batches.clone());
+                self.last_query_results = Some(sorbet_batches.clone());
             }
 
             *dataframe = RequestedObject::new_with_repaint(