Assume rerun.kind: data when rerun.kind is missing (#9467)

### What
- default `rerun.kind` to `"data"`

This should solve the problem of the partition table not being Sorbet
compliant, by making more things sorbet compliant.

### TODO
* [x] Test with a real partition table

<img width="1693" alt="image"
src="https://github.com/user-attachments/assets/401ac9c4-aa0f-47cd-a8ba-8743468160cc"
/>

---------

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

Author: abey79

crates/store/re_sorbet/src/column_descriptor.rs

@@ -9,15 +9,18 @@ use arrow::datatypes::{
 
 use re_log_types::EntityPath;
 
-use crate::{ComponentColumnDescriptor, IndexColumnDescriptor, MetadataExt as _};
+use crate::{ColumnKind, ComponentColumnDescriptor, IndexColumnDescriptor};
 
 #[derive(thiserror::Error, Debug)]
 pub enum ColumnError {
     #[error(transparent)]
     MissingFieldMetadata(#[from] crate::MissingFieldMetadata),
 
+    #[error(transparent)]
+    UnknownColumnKind(#[from] crate::UnknownColumnKind),
+
     #[error("Unsupported column rerun.kind: {kind:?}. Expected one of: index, data")]
-    UnsupportedColumnKind { kind: String },
+    UnsupportedColumnKind { kind: ColumnKind },
 
     #[error(transparent)]
     UnsupportedTimeType(#[from] crate::UnsupportedTimeType),
@@ -114,17 +117,16 @@ impl ColumnDescriptor {
         chunk_entity_path: Option<&EntityPath>,
         field: &ArrowField,
     ) -> Result<Self, ColumnError> {
-        let kind = field.get_or_err("rerun.kind")?;
-        match kind {
-            "index" | "time" => Ok(Self::Time(IndexColumnDescriptor::try_from(field)?)),
+        match ColumnKind::try_from(field)? {
+            ColumnKind::RowId => Err(ColumnError::UnsupportedColumnKind {
+                kind: ColumnKind::RowId,
+            }),
 
-            "data" => Ok(Self::Component(
+            ColumnKind::Index => Ok(Self::Time(IndexColumnDescriptor::try_from(field)?)),
+
+            ColumnKind::Component => Ok(Self::Component(
                 ComponentColumnDescriptor::from_arrow_field(chunk_entity_path, field),
             )),
-
-            _ => Err(ColumnError::UnsupportedColumnKind {
-                kind: kind.to_owned(),
-            }),
         }
     }
 }

crates/store/re_sorbet/src/column_kind.rs

@@ -1,26 +1,54 @@
 use arrow::datatypes::Field as ArrowField;
 
-use crate::{MetadataExt as _, SorbetError};
+use crate::MetadataExt as _;
+
+#[derive(thiserror::Error, Debug)]
+#[error("Unknown `rerun.kind` {kind:?} in column {column_name:?}. Expect one of `row_id`, `index`, or `component`.")]
+pub struct UnknownColumnKind {
+    pub kind: String,
+    pub column_name: String,
+}
 
 /// The type of column in a sorbet batch.
-#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)]
 pub enum ColumnKind {
+    /// Row ID
     RowId,
+
+    /// Timeline
     Index,
+
+    /// Data (also the default when unknown)
+    #[default]
     Component,
 }
 
+impl std::fmt::Display for ColumnKind {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::RowId => write!(f, "control"),
+            Self::Index => write!(f, "index"),
+            Self::Component => write!(f, "data"),
+        }
+    }
+}
+
 impl TryFrom<&ArrowField> for ColumnKind {
-    type Error = SorbetError;
+    type Error = UnknownColumnKind;
 
-    fn try_from(fields: &ArrowField) -> Result<Self, Self::Error> {
-        let kind = fields.get_or_err("rerun.kind")?;
+    fn try_from(field: &ArrowField) -> Result<Self, Self::Error> {
+        let Some(kind) = field.get_opt("rerun.kind") else {
+            return Ok(Self::default());
+        };
         match kind {
             "control" | "row_id" => Ok(Self::RowId),
             "index" | "time" => Ok(Self::Index),
             "component" | "data" => Ok(Self::Component),
 
-            _ => Err(SorbetError::custom(format!("Unknown column kind: {kind}"))),
+            _ => Err(UnknownColumnKind {
+                kind: kind.to_owned(),
+                column_name: field.name().to_owned(),
+            }),
         }
     }
 }

crates/store/re_sorbet/src/component_column_descriptor.rs

@@ -3,7 +3,7 @@ use arrow::datatypes::{DataType as ArrowDatatype, Field as ArrowField};
 use re_log_types::{ComponentPath, EntityPath};
 use re_types_core::{ArchetypeFieldName, ArchetypeName, ComponentDescriptor, ComponentName};
 
-use crate::{ArrowFieldMetadata, BatchType, MetadataExt as _};
+use crate::{ArrowFieldMetadata, BatchType, ColumnKind, MetadataExt as _};
 
 /// Describes a data/component column, such as `Position3D`, in a dataframe.
 ///
@@ -196,7 +196,7 @@ impl ComponentColumnDescriptor {
         // TODO(#6889): This needs some proper sorbetization -- I just threw these names randomly.
         // We use the long names for the archetype and component names so that they roundtrip properly!
         let mut metadata = std::collections::HashMap::from([
-            ("rerun.kind".to_owned(), "data".to_owned()),
+            ("rerun.kind".to_owned(), ColumnKind::Component.to_string()),
             (
                 "rerun.component".to_owned(),
                 component_name.full_name().to_owned(),

crates/store/re_sorbet/src/error.rs

@@ -2,6 +2,9 @@ use arrow::error::ArrowError;
 
 #[derive(thiserror::Error, Debug)]
 pub enum SorbetError {
+    #[error(transparent)]
+    UnknownColumnKind(#[from] crate::UnknownColumnKind),
+
     #[error(transparent)]
     MissingMetadataKey(#[from] crate::MissingMetadataKey),
 

crates/store/re_sorbet/src/index_column_descriptor.rs

@@ -105,7 +105,10 @@ impl IndexColumnDescriptor {
         let nullable = true; // Time column must be nullable since static data doesn't have a time.
 
         let mut metadata = std::collections::HashMap::from([
-            ("rerun.kind".to_owned(), "index".to_owned()),
+            (
+                "rerun.kind".to_owned(),
+                crate::ColumnKind::Index.to_string(),
+            ),
             ("rerun.index_name".to_owned(), timeline.name().to_string()),
         ]);
         if *is_sorted {

crates/store/re_sorbet/src/lib.rs

@@ -36,7 +36,7 @@ pub use self::{
     chunk_schema::ChunkSchema,
     column_descriptor::{ColumnDescriptor, ColumnError},
     column_descriptor_ref::ColumnDescriptorRef,
-    column_kind::ColumnKind,
+    column_kind::{ColumnKind, UnknownColumnKind},
     component_column_descriptor::ComponentColumnDescriptor,
     error::SorbetError,
     index_column_descriptor::{IndexColumnDescriptor, UnsupportedTimeType},

crates/store/re_sorbet/src/row_id_column_descriptor.rs

@@ -53,7 +53,10 @@ impl RowIdColumnDescriptor {
         let Self { is_sorted } = self;
 
         let mut metadata = std::collections::HashMap::from([
-            ("rerun.kind".to_owned(), "control".to_owned()),
+            (
+                "rerun.kind".to_owned(),
+                crate::ColumnKind::RowId.to_string(),
+            ),
             (
                 "ARROW:extension:name".to_owned(),
                 re_tuid::Tuid::ARROW_EXTENSION_NAME.to_owned(),

crates/store/re_sorbet/src/sorbet_batch.rs

@@ -10,8 +10,8 @@ use arrow::{
 use re_arrow_util::{into_arrow_ref, ArrowArrayDowncastRef as _};
 
 use crate::{
-    ArrowBatchMetadata, ColumnDescriptorRef, ComponentColumnDescriptor, IndexColumnDescriptor,
-    RowIdColumnDescriptor, SorbetError, SorbetSchema,
+    ArrowBatchMetadata, ColumnDescriptorRef, ColumnKind, ComponentColumnDescriptor,
+    IndexColumnDescriptor, RowIdColumnDescriptor, SorbetError, SorbetSchema,
 };
 
 /// Any rerun-compatible [`ArrowRecordBatch`].
@@ -205,10 +205,8 @@ fn make_all_data_columns_list_arrays(batch: &ArrowRecordBatch) -> ArrowRecordBat
 
     for (field, array) in itertools::izip!(batch.schema().fields(), batch.columns()) {
         let is_list_array = array.downcast_array_ref::<ArrowListArray>().is_some();
-        let is_data_column = field
-            .metadata()
-            .get("rerun.kind")
-            .is_some_and(|kind| kind == "data");
+        let is_data_column =
+            ColumnKind::try_from(field.as_ref()).is_ok_and(|kind| kind == ColumnKind::Component);
         if is_data_column && !is_list_array {
             let (field, array) = re_arrow_util::wrap_in_list_array(field, array.clone());
             fields.push(field.into());