Add re_arrow_combinators::reshape::RowMajorToColumnMajor (#12031)

Required for example to read matrices from variable sized, row-major
protobuf lists, e.g. the 3x3 or 3x4 camera matrices that are found in
`foxglove.CameraCalibration` schema.

In order to do that, we need both a cast to fixed-size-list and a
reshaping transformation.

---

This is a prerequisite for #12035.

Contributes to:
https://linear.app/rerun/issue/RR-3045
https://linear.app/rerun/issue/RR-2327

Author: MichaelGrupp

crates/store/re_arrow_combinators/src/cast.rs

@@ -1,7 +1,10 @@
 //! Transforms that cast arrays to different types.
 
+use std::sync::Arc;
+
 use arrow::array::{Array, ArrayRef, ArrowPrimitiveType, PrimitiveArray};
 use arrow::compute::cast;
+use arrow::datatypes::Field;
 
 use crate::{Error, Transform};
 
@@ -83,3 +86,54 @@ where
             .cloned()
     }
 }
+
+/// Casts a `ListArray` to a `FixedSizeListArray` with the specified value length.
+///
+/// The source `ListArray` must have lists of exactly that length (or null).
+#[derive(Clone)]
+pub struct ListToFixedSizeList {
+    value_length: i32,
+}
+
+impl ListToFixedSizeList {
+    /// Create a new `ListToFixedSizeList` transformation with an expected value length.
+    pub fn new(value_length: i32) -> Self {
+        Self { value_length }
+    }
+}
+
+impl Transform for ListToFixedSizeList {
+    type Source = arrow::array::ListArray;
+    type Target = arrow::array::FixedSizeListArray;
+
+    fn transform(&self, source: &Self::Source) -> Result<Self::Target, Error> {
+        // Check that each list has exactly the expected length (or is null).
+        let offsets = source.value_offsets();
+        let expected_length = self.value_length as usize;
+        for list_index in 0..source.len() {
+            if source.is_valid(list_index) {
+                let start = offsets[list_index] as usize;
+                let end = offsets[list_index + 1] as usize;
+                let list_length = end - start;
+                if list_length != expected_length {
+                    return Err(Error::UnexpectedListValueLength {
+                        expected: expected_length,
+                        actual: list_length,
+                    });
+                }
+            }
+        }
+
+        // Build the FixedSizeListArray.
+        let field = Arc::new(Field::new_list_field(
+            source.value_type().clone(),
+            source.is_nullable(),
+        ));
+        Ok(arrow::array::FixedSizeListArray::try_new(
+            field,
+            self.value_length,
+            source.values().clone(),
+            source.nulls().cloned(),
+        )?)
+    }
+}

crates/store/re_arrow_combinators/src/error.rs

@@ -43,6 +43,9 @@ pub enum Error {
     #[error("List contains unexpected value type: expected {expected}, got {actual:?}")]
     UnexpectedListValueType { expected: String, actual: DataType },
 
+    #[error("Expected list with {expected} elements, got {actual}")]
+    UnexpectedListValueLength { expected: usize, actual: usize },
+
     #[error("Fixed-size list contains unexpected value type: expected {expected}, got {actual:?}")]
     UnexpectedFixedSizeListValueType { expected: String, actual: DataType },
 

crates/store/re_arrow_combinators/src/reshape.rs

@@ -2,7 +2,9 @@
 
 use std::sync::Arc;
 
-use arrow::array::{Array, ArrayRef, FixedSizeListArray, ListArray, StructArray, UInt32Array};
+use arrow::array::{
+    Array, ArrayRef, FixedSizeListArray, ListArray, StructArray, UInt32Array, UInt64Array,
+};
 use arrow::buffer::{NullBuffer, OffsetBuffer};
 use arrow::datatypes::Field;
 
@@ -383,3 +385,77 @@ impl Transform for Explode {
         ))
     }
 }
+
+/// Reorders a `FixedSizeListArray`, where each `FixedSizeList` stores matrix elements
+/// in flat row-major order, to `FixedSizeList`s in column-major order.
+///
+/// The source array is expected to have a value length of `output_rows * output_columns`.
+#[derive(Clone, Debug)]
+pub struct RowMajorToColumnMajor {
+    output_rows: usize,
+    output_columns: usize,
+    permutation_per_list: Vec<usize>,
+}
+
+impl RowMajorToColumnMajor {
+    /// Create a new row-major to column-major transformation for the desired output shape.
+    pub fn new(output_rows: usize, output_columns: usize) -> Self {
+        let mut permutation = Vec::with_capacity(output_rows * output_columns);
+        for column in 0..output_columns {
+            for row in 0..output_rows {
+                let row_major_pos = row * output_columns + column;
+                permutation.push(row_major_pos);
+            }
+        }
+        Self {
+            output_rows,
+            output_columns,
+            permutation_per_list: permutation,
+        }
+    }
+}
+
+impl Transform for RowMajorToColumnMajor {
+    type Source = FixedSizeListArray;
+    type Target = FixedSizeListArray;
+
+    fn transform(&self, source: &Self::Source) -> Result<Self::Target, Error> {
+        // First, check that the input array has the expected value length.
+        let expected_list_size = self.output_rows * self.output_columns;
+        let value_length = source.value_length() as usize;
+        if value_length != expected_list_size {
+            return Err(Error::UnexpectedListValueLength {
+                expected: expected_list_size,
+                actual: value_length,
+            });
+        }
+
+        // Create indices for extracting column-major values as row-major, for all input lists.
+        let total_values = source.values().len();
+        let indices_to_take: UInt64Array = (0..total_values)
+            .map(|value_index| {
+                let list_index = value_index / expected_list_size;
+                let value_index_within_list = value_index % expected_list_size;
+                let next_index_to_take = list_index * expected_list_size
+                    + self.permutation_per_list[value_index_within_list];
+                next_index_to_take as u64
+            })
+            .collect();
+
+        // Reorder values into a new FixedSizeListArray.
+        // We explicitly allow `take` here because we care about nulls.
+        #[expect(clippy::disallowed_methods)]
+        let reordered_values = arrow::compute::take(source.values(), &indices_to_take, None)?;
+
+        let field = Arc::new(Field::new_list_field(
+            source.value_type().clone(),
+            source.is_nullable(),
+        ));
+        Ok(FixedSizeListArray::new(
+            field,
+            source.value_length(),
+            reordered_values,
+            source.nulls().cloned(),
+        ))
+    }
+}

crates/store/re_arrow_combinators/tests/snapshots/transform__row_major_to_col_major.snap

@@ -0,0 +1,15 @@
+---
+source: crates/utils/re_arrow_combinators/tests/transform.rs
+expression: "format!(\"{}\", DisplayRB(result.clone()))"
+---
+┌──────────────────────────────────────────────────┐
+│ col                                              │
+│ ---                                              │
+│ type: nullable FixedSizeList[nullable i32; 12]   │
+╞══════════════════════════════════════════════════╡
+│ [1, 4, 7, 10, null, 5, 8, 11, 3, 6, null, 12]    │
+├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+│ null                                             │
+├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
+│ [13, 16, 19, 22, 14, 17, 20, 23, 15, 18, 21, 24] │
+└──────────────────────────────────────────────────┘

crates/store/re_arrow_combinators/tests/transform.rs

@@ -6,13 +6,16 @@ use std::sync::Arc;
 
 use re_arrow_combinators::{
     Transform as _,
-    cast::PrimitiveCast,
+    cast::{ListToFixedSizeList, PrimitiveCast},
     map::{MapFixedSizeList, MapList, MapPrimitive, ReplaceNull},
-    reshape::{Flatten, GetField, StructToFixedList},
+    reshape::{Flatten, GetField, RowMajorToColumnMajor, StructToFixedList},
 };
 
 use arrow::{
-    array::{Float32Array, Float64Array, Float64Builder, ListArray, ListBuilder, StructBuilder},
+    array::{
+        Float32Array, Float64Array, Float64Builder, Int32Builder, ListArray, ListBuilder,
+        StructBuilder,
+    },
     datatypes::{DataType, Field, Fields},
 };
 
@@ -295,3 +298,69 @@ fn test_flatten_multiple_elements() {
         format!("{}", DisplayRB(result.clone()))
     );
 }
+
+#[test]
+fn test_row_major_to_col_major() {
+    let inner_builder = Int32Builder::new();
+    let mut outer_builder = ListBuilder::new(inner_builder);
+
+    // First list represents a 4x3 matrix in row-major order with some null elements.
+    // Row 0
+    outer_builder.values().append_value(1);
+    outer_builder.values().append_null();
+    outer_builder.values().append_value(3);
+    // Row 1
+    outer_builder.values().append_value(4);
+    outer_builder.values().append_value(5);
+    outer_builder.values().append_value(6);
+    // Row 2
+    outer_builder.values().append_value(7);
+    outer_builder.values().append_value(8);
+    outer_builder.values().append_null();
+    // Row 3
+    outer_builder.values().append_value(10);
+    outer_builder.values().append_value(11);
+    outer_builder.values().append_value(12);
+    outer_builder.append(true);
+
+    // Second list is invalid / null.
+    for _ in 0..12 {
+        // Add dummy values for Arrow's fixed-size requirements.
+        // See: https://docs.rs/arrow/latest/arrow/array/struct.FixedSizeListArray.html#representation
+        outer_builder.values().append_value(0);
+    }
+    outer_builder.append(false);
+
+    // Third list represents a 4x3 matrix in row-major order without null elements.
+    // Row 0
+    outer_builder.values().append_value(13);
+    outer_builder.values().append_value(14);
+    outer_builder.values().append_value(15);
+    // Row 1
+    outer_builder.values().append_value(16);
+    outer_builder.values().append_value(17);
+    outer_builder.values().append_value(18);
+    // Row 2
+    outer_builder.values().append_value(19);
+    outer_builder.values().append_value(20);
+    outer_builder.values().append_value(21);
+    // Row 3
+    outer_builder.values().append_value(22);
+    outer_builder.values().append_value(23);
+    outer_builder.values().append_value(24);
+    outer_builder.append(true);
+
+    let input_array = outer_builder.finish();
+
+    // Cast to `FixedSizeListArray` and convert to column-major order.
+    let fixed_size_list_array = ListToFixedSizeList::new(12)
+        .transform(&input_array)
+        .unwrap();
+    let result = RowMajorToColumnMajor::new(4, 3)
+        .transform(&fixed_size_list_array)
+        .unwrap();
+    insta::assert_snapshot!(
+        "row_major_to_col_major",
+        format!("{}", DisplayRB(result.clone()))
+    );
+}