re_chunk/

shuffle.rs

use arrow::array::{Array as ArrowArray, ListArray as ArrowListArray};
use arrow::buffer::{OffsetBuffer as ArrowOffsets, ScalarBuffer as ArrowScalarBuffer};
use itertools::Itertools as _;
use re_log_types::TimelineName;

use crate::{Chunk, ChunkId, TimeColumn};

// ---

impl Chunk {
    /// Is the chunk currently ascendingly sorted by [`crate::RowId`]?
    ///
    /// This is O(1) (cached).
    ///
    /// See also [`Self::is_sorted_uncached`].
    #[inline]
    pub fn is_sorted(&self) -> bool {
        self.is_sorted
    }

    /// For debugging purposes.
    #[doc(hidden)]
    #[inline]
    pub fn is_sorted_uncached(&self) -> bool {
        re_tracing::profile_function!();

        self.row_ids()
            .tuple_windows::<(_, _)>()
            .all(|row_ids| row_ids.0 <= row_ids.1)
    }

    /// Is the chunk ascendingly sorted by time, for all of its timelines?
    ///
    /// This is O(1) (cached).
    #[inline]
    pub fn is_time_sorted(&self) -> bool {
        self.timelines
            .values()
            .all(|time_column| time_column.is_sorted())
    }

    /// Is the chunk ascendingly sorted by time, for a specific timeline?
    ///
    /// This is O(1) (cached).
    ///
    /// See also [`Self::is_timeline_sorted_uncached`].
    #[inline]
    pub fn is_timeline_sorted(&self, timeline: &TimelineName) -> bool {
        self.is_static()
            || self
                .timelines
                .get(timeline)
                .is_some_and(|time_column| time_column.is_sorted())
    }

    /// For debugging purposes.
    #[doc(hidden)]
    #[inline]
    pub fn is_timeline_sorted_uncached(&self, timeline: &TimelineName) -> bool {
        self.is_static()
            || self
                .timelines
                .get(timeline)
                .is_some_and(|time_column| time_column.is_sorted_uncached())
    }

    /// Sort the chunk, if needed.
    ///
    /// The underlying arrow data will be copied and shuffled in memory in order to make it contiguous.
    ///
    /// If the chunk changes, it is given a new unique [`ChunkId`].
    #[inline]
    pub fn sort_if_unsorted(&mut self) {
        if self.is_sorted() {
            return;
        }

        re_tracing::profile_function!();

        self.id = ChunkId::new();

        #[cfg(not(target_arch = "wasm32"))]
        let now = std::time::Instant::now();

        let swaps = {
            re_tracing::profile_scope!("swaps");
            let row_ids = self.row_ids_slice();
            let mut swaps = (0..row_ids.len()).collect::<Vec<_>>();
            swaps.sort_by_key(|&i| row_ids[i]);
            swaps
        };

        self.shuffle_with(&swaps);

        #[cfg(not(target_arch = "wasm32"))]
        re_log::trace!(
            entity_path = %self.entity_path,
            num_rows = self.row_ids.len(),
            elapsed = ?now.elapsed(),
            "chunk sorted",
        );

        #[cfg(debug_assertions)]
        #[expect(clippy::unwrap_used)] // dev only
        self.sanity_check().unwrap();
    }

    /// Returns a new [`Chunk`] that is sorted by `(<timeline>, RowId)`.
    ///
    /// The underlying arrow data will be copied and shuffled in memory in order to make it contiguous.
    ///
    /// This is a no-op if the underlying timeline is already sorted appropriately (happy path).
    ///
    /// If the chunk is already sorted, the original [`crate::ChunkId`] is preserved.
    /// Otherwise, the returned chunk gets a new unique [`crate::ChunkId`].
    #[must_use]
    pub fn sorted_by_timeline_if_unsorted(&self, timeline: &TimelineName) -> Self {
        let Some(time_column) = self.timelines.get(timeline) else {
            return self.clone_with_same_id();
        };

        if time_column.is_sorted() {
            return self.clone_with_same_id();
        }

        let mut chunk = self.clone_with_new_id();

        re_tracing::profile_function!();

        #[cfg(not(target_arch = "wasm32"))]
        let now = std::time::Instant::now();

        let swaps = {
            re_tracing::profile_scope!("swaps");
            let row_ids = chunk.row_ids_slice();
            let times = time_column.times_raw().to_vec();
            let mut swaps = (0..times.len()).collect::<Vec<_>>();
            swaps.sort_by_key(|&i| (times[i], row_ids[i]));
            swaps
        };

        chunk.shuffle_with(&swaps);

        #[cfg(not(target_arch = "wasm32"))]
        re_log::trace!(
            entity_path = %chunk.entity_path,
            num_rows = chunk.row_ids.len(),
            elapsed = ?now.elapsed(),
            "chunk sorted",
        );

        #[cfg(debug_assertions)]
        #[expect(clippy::unwrap_used)] // dev only
        chunk.sanity_check().unwrap();

        chunk
    }

    /// Randomly shuffles the chunk using the given `seed`.
    ///
    /// The underlying arrow data will be copied and shuffled in memory in order to make it contiguous.
    #[inline]
    #[cfg(debug_assertions)] // only for tests
    pub fn shuffle_random(&mut self, seed: u64) {
        re_tracing::profile_function!();

        #[cfg(not(target_arch = "wasm32"))]
        let now = std::time::Instant::now();

        use rand::SeedableRng as _;
        use rand::seq::SliceRandom as _;
        let mut rng = rand::rngs::StdRng::seed_from_u64(seed);

        let swaps = {
            re_tracing::profile_scope!("swaps");
            let mut swaps = (0..self.row_ids.len()).collect::<Vec<_>>();
            swaps.shuffle(&mut rng);
            swaps
        };

        self.shuffle_with(&swaps);

        #[cfg(not(target_arch = "wasm32"))]
        re_log::trace!(
            entity_path = %self.entity_path,
            num_rows = self.row_ids.len(),
            elapsed = ?now.elapsed(),
            "chunk shuffled",
        );
    }

    /// Shuffle the chunk according to the specified `swaps`.
    ///
    /// `swaps` is a slice that maps an implicit destination index to its explicit source index.
    /// E.g. `swap[0] = 3` means that the entry at index `3` in the original chunk should be move to index `0`.
    ///
    /// The underlying arrow data will be copied and shuffled in memory in order to make it contiguous.
    //
    // TODO(RR-3865): Use arrow's `ListView` to only shuffle offsets instead of the data itself.
    pub(crate) fn shuffle_with(&mut self, swaps: &[usize]) {
        re_tracing::profile_function!();

        self.id = ChunkId::new();

        // Row IDs
        {
            re_tracing::profile_scope!("row ids");

            let row_ids = self.row_ids_slice();

            let mut sorted_row_ids = row_ids.to_vec();
            for (to, from) in swaps.iter().copied().enumerate() {
                sorted_row_ids[to] = row_ids[from];
            }
            self.row_ids = re_types_core::RowId::arrow_from_slice(&sorted_row_ids);
        }

        let Self {
            id: _,
            entity_path: _,
            heap_size_bytes: _,
            is_sorted: _,
            row_ids: _,
            timelines,
            components,
        } = self;

        // Timelines
        {
            re_tracing::profile_scope!("timelines");

            for info in timelines.values_mut() {
                let TimeColumn {
                    timeline: _,
                    times,
                    is_sorted,
                    time_range: _,
                } = info;

                let mut sorted = times.to_vec();
                for (to, from) in swaps.iter().copied().enumerate() {
                    sorted[to] = times[from];
                }

                *is_sorted = sorted.windows(2).all(|times| times[0] <= times[1]);
                *times = ArrowScalarBuffer::from(sorted);
            }
        }

        // Components
        //
        // Reminder: these are all `ListArray`s.
        re_tracing::profile_scope!("components (offsets & data)");
        {
            for original in components.list_arrays_mut() {
                let sorted_arrays = swaps
                    .iter()
                    .copied()
                    .map(|from| original.value(from))
                    .collect_vec();
                let sorted_arrays = sorted_arrays
                    .iter()
                    .map(|array| &**array as &dyn ArrowArray)
                    .collect_vec();

                let datatype = original.data_type().clone();
                let offsets =
                    ArrowOffsets::from_lengths(sorted_arrays.iter().map(|array| array.len()));
                #[expect(clippy::unwrap_used)] // these are slices of the same outer array
                let values = re_arrow_util::concat_arrays(&sorted_arrays).unwrap();
                let validity = original
                    .nulls()
                    .map(|validity| swaps.iter().map(|&from| validity.is_valid(from)).collect());

                let field = match datatype {
                    arrow::datatypes::DataType::List(field) => field.clone(),
                    _ => unreachable!("This is always s list array"),
                };
                *original = ArrowListArray::new(field, offsets, values, validity);
            }
        }

        self.is_sorted = self.is_sorted_uncached();
    }
}

impl TimeColumn {
    /// Is the timeline sorted?
    ///
    /// This is O(1) (cached).
    ///
    /// See also [`Self::is_sorted_uncached`].
    #[inline]
    pub fn is_sorted(&self) -> bool {
        self.is_sorted
    }

    /// Like [`Self::is_sorted`], but actually checks the entire dataset rather than relying on the
    /// cached value.
    ///
    /// O(n). Useful for tests/debugging, or when you just don't know.
    ///
    /// See also [`Self::is_sorted`].
    #[inline]
    pub fn is_sorted_uncached(&self) -> bool {
        re_tracing::profile_function!();
        self.times_raw()
            .windows(2)
            .all(|times| times[0] <= times[1])
    }
}

#[cfg(test)]
mod tests {
    use re_log_types::example_components::{MyColor, MyPoint, MyPoints};
    use re_log_types::{EntityPath, Timeline};
    use re_types_core::ComponentBatch as _;

    use super::*;
    use crate::{ChunkId, RowId};

    #[test]
    fn sort() -> anyhow::Result<()> {
        let entity_path: EntityPath = "a/b/c".into();

        let timeline1 = Timeline::new_duration("log_time");
        let timeline2 = Timeline::new_sequence("frame_nr");

        let points1 = vec![
            MyPoint::new(1.0, 2.0),
            MyPoint::new(3.0, 4.0),
            MyPoint::new(5.0, 6.0),
        ];
        let points3 = vec![MyPoint::new(10.0, 20.0)];
        let points4 = vec![MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)];

        let colors1 = vec![
            MyColor::from_rgb(1, 2, 3),
            MyColor::from_rgb(4, 5, 6),
            MyColor::from_rgb(7, 8, 9),
        ];
        let colors2 = vec![MyColor::from_rgb(10, 20, 30)];
        let colors4 = vec![
            MyColor::from_rgb(101, 102, 103),
            MyColor::from_rgb(104, 105, 106),
        ];

        {
            let chunk_sorted = Chunk::builder(entity_path.clone())
                .with_sparse_component_batches(
                    RowId::new(),
                    [(timeline1, 1000), (timeline2, 42)],
                    [
                        (MyPoints::descriptor_points(), Some(&points1 as _)),
                        (MyPoints::descriptor_colors(), Some(&colors1 as _)),
                    ],
                )
                .with_sparse_component_batches(
                    RowId::new(),
                    [(timeline1, 1001), (timeline2, 43)],
                    [
                        (MyPoints::descriptor_points(), None),
                        (MyPoints::descriptor_colors(), Some(&colors2 as _)),
                    ],
                )
                .with_sparse_component_batches(
                    RowId::new(),
                    [(timeline1, 1002), (timeline2, 44)],
                    [
                        (MyPoints::descriptor_points(), Some(&points3 as _)),
                        (MyPoints::descriptor_colors(), None),
                    ],
                )
                .with_sparse_component_batches(
                    RowId::new(),
                    [(timeline1, 1003), (timeline2, 45)],
                    [
                        (MyPoints::descriptor_points(), Some(&points4 as _)),
                        (MyPoints::descriptor_colors(), Some(&colors4 as _)),
                    ],
                )
                .build()?;

            eprintln!("{chunk_sorted}");

            assert!(chunk_sorted.is_sorted());
            assert!(chunk_sorted.is_sorted_uncached());

            let chunk_shuffled = {
                let mut chunk_shuffled = chunk_sorted.clone();
                chunk_shuffled.shuffle_random(666);
                chunk_shuffled
            };

            eprintln!("{chunk_shuffled}");

            assert!(!chunk_shuffled.is_sorted());
            assert!(!chunk_shuffled.is_sorted_uncached());
            assert_ne!(chunk_sorted, chunk_shuffled);

            let chunk_resorted = {
                let mut chunk_resorted = chunk_shuffled.clone();
                chunk_resorted.sort_if_unsorted();
                chunk_resorted
            };

            eprintln!("{chunk_resorted}");

            assert!(chunk_resorted.is_sorted());
            assert!(chunk_resorted.is_sorted_uncached());
            assert_eq!(chunk_sorted, chunk_resorted);
        }

        Ok(())
    }

    #[test]
    fn sort_time() -> anyhow::Result<()> {
        let entity_path: EntityPath = "a/b/c".into();

        let timeline1 = Timeline::new_duration("log_time");
        let timeline2 = Timeline::new_sequence("frame_nr");

        let chunk_id = ChunkId::new();
        let row_id1 = RowId::new();
        let row_id2 = RowId::new();
        let row_id3 = RowId::new();
        let row_id4 = RowId::new();

        let points1 = vec![
            MyPoint::new(1.0, 2.0),
            MyPoint::new(3.0, 4.0),
            MyPoint::new(5.0, 6.0),
        ];
        let points3 = vec![MyPoint::new(10.0, 20.0)];
        let points4 = vec![MyPoint::new(100.0, 200.0), MyPoint::new(300.0, 400.0)];

        let colors1 = vec![
            MyColor::from_rgb(1, 2, 3),
            MyColor::from_rgb(4, 5, 6),
            MyColor::from_rgb(7, 8, 9),
        ];
        let colors2 = vec![MyColor::from_rgb(10, 20, 30)];
        let colors4 = vec![
            MyColor::from_rgb(101, 102, 103),
            MyColor::from_rgb(104, 105, 106),
        ];

        {
            let chunk_unsorted_timeline2 = Chunk::builder_with_id(chunk_id, entity_path.clone())
                .with_sparse_component_batches(
                    row_id1,
                    [(timeline1, 1000), (timeline2, 45)],
                    [
                        (MyPoints::descriptor_points(), Some(&points1 as _)),
                        (MyPoints::descriptor_colors(), Some(&colors1 as _)),
                    ],
                )
                .with_sparse_component_batches(
                    row_id2,
                    [(timeline1, 1001), (timeline2, 44)],
                    [
                        (MyPoints::descriptor_points(), None),
                        (MyPoints::descriptor_colors(), Some(&colors2 as _)),
                    ],
                )
                .with_sparse_component_batches(
                    row_id3,
                    [(timeline1, 1002), (timeline2, 43)],
                    [
                        (MyPoints::descriptor_points(), Some(&points3 as _)),
                        (MyPoints::descriptor_colors(), None),
                    ],
                )
                .with_sparse_component_batches(
                    row_id4,
                    [(timeline1, 1003), (timeline2, 42)],
                    [
                        (MyPoints::descriptor_points(), Some(&points4 as _)),
                        (MyPoints::descriptor_colors(), Some(&colors4 as _)),
                    ],
                )
                .build()?;

            eprintln!("unsorted:\n{chunk_unsorted_timeline2}");

            assert!(chunk_unsorted_timeline2.is_sorted());
            assert!(chunk_unsorted_timeline2.is_sorted_uncached());

            assert!(
                chunk_unsorted_timeline2
                    .timelines()
                    .get(timeline1.name())
                    .unwrap()
                    .is_sorted()
            );
            assert!(
                chunk_unsorted_timeline2
                    .timelines()
                    .get(timeline1.name())
                    .unwrap()
                    .is_sorted_uncached()
            );

            assert!(
                !chunk_unsorted_timeline2
                    .timelines()
                    .get(timeline2.name())
                    .unwrap()
                    .is_sorted()
            );
            assert!(
                !chunk_unsorted_timeline2
                    .timelines()
                    .get(timeline2.name())
                    .unwrap()
                    .is_sorted_uncached()
            );

            let chunk_sorted_timeline2 =
                chunk_unsorted_timeline2.sorted_by_timeline_if_unsorted(timeline2.name());

            eprintln!("sorted:\n{chunk_sorted_timeline2}");

            assert!(!chunk_sorted_timeline2.is_sorted());
            assert!(!chunk_sorted_timeline2.is_sorted_uncached());

            assert!(
                !chunk_sorted_timeline2
                    .timelines()
                    .get(timeline1.name())
                    .unwrap()
                    .is_sorted()
            );
            assert!(
                !chunk_sorted_timeline2
                    .timelines()
                    .get(timeline1.name())
                    .unwrap()
                    .is_sorted_uncached()
            );

            assert!(
                chunk_sorted_timeline2
                    .timelines()
                    .get(timeline2.name())
                    .unwrap()
                    .is_sorted()
            );
            assert!(
                chunk_sorted_timeline2
                    .timelines()
                    .get(timeline2.name())
                    .unwrap()
                    .is_sorted_uncached()
            );

            let chunk_sorted_timeline2_expected =
                Chunk::builder_with_id(chunk_id, entity_path.clone())
                    .with_sparse_component_batches(
                        row_id4,
                        [(timeline1, 1003), (timeline2, 42)],
                        [
                            (MyPoints::descriptor_points(), Some(&points4 as _)),
                            (MyPoints::descriptor_colors(), Some(&colors4 as _)),
                        ],
                    )
                    .with_sparse_component_batches(
                        row_id3,
                        [(timeline1, 1002), (timeline2, 43)],
                        [
                            (MyPoints::descriptor_points(), Some(&points3 as _)),
                            (MyPoints::descriptor_colors(), None),
                        ],
                    )
                    .with_sparse_component_batches(
                        row_id2,
                        [(timeline1, 1001), (timeline2, 44)],
                        [
                            (MyPoints::descriptor_points(), None),
                            (MyPoints::descriptor_colors(), Some(&colors2 as _)),
                        ],
                    )
                    .with_sparse_component_batches(
                        row_id1,
                        [(timeline1, 1000), (timeline2, 45)],
                        [
                            (MyPoints::descriptor_points(), Some(&points1 as _)),
                            (MyPoints::descriptor_colors(), Some(&colors1 as _)),
                        ],
                    )
                    .build()?;

            eprintln!("expected:\n{chunk_sorted_timeline2}");

            assert_eq!(
                chunk_sorted_timeline2_expected,
                chunk_sorted_timeline2,
                "{}",
                similar_asserts::SimpleDiff::from_str(
                    &format!("{chunk_sorted_timeline2_expected}"),
                    &format!("{chunk_sorted_timeline2}"),
                    "got",
                    "expected",
                ),
            );
        }

        Ok(())
    }

    /// `Chunk::from_auto_row_ids` should reorder its inputs so that the time columns become
    /// sorted as well as possible, to avoid "out-of-order" chunks where some time columns are not sorted with respect to `RowId`.
    #[test]
    fn from_auto_row_ids_sorts_lexicographically() -> anyhow::Result<()> {
        let entity_path: EntityPath = "a/b/c".into();

        // Two timelines named such that "alpha" sorts before "beta".
        // Construct deliberately unsorted inputs:
        //   row | alpha | beta | color
        //   ----|-------|------|------
        //    0  |   2   |   5  |  100
        //    1  |   1   |   9  |  200
        //    2  |   1   |   7  |  300
        //    3  |   2   |   3  |  400
        //    4  |   1   |   9  |  500   (duplicate of row 1's key — tests stability)
        //
        // After lex-sort by (alpha, beta) the expected row order is:
        //   2 (1,7,300), 1 (1,9,200), 4 (1,9,500), 3 (2,3,400), 0 (2,5,100)
        let alpha = TimeColumn::new_sequence("alpha", [2_i64, 1, 1, 2, 1]);
        let beta = TimeColumn::new_sequence("beta", [5_i64, 9, 7, 3, 9]);

        let colors = vec![
            MyColor(100),
            MyColor(200),
            MyColor(300),
            MyColor(400),
            MyColor(500),
        ];
        let colors_array = colors.to_arrow_list_array()?;

        // `Chunk::from_columns` uses `Chunk::from_auto_row_ids`.
        let chunk = Chunk::from_columns(
            entity_path,
            [alpha, beta],
            [(MyPoints::descriptor_colors(), colors_array)],
        )?;

        eprintln!("{chunk}");

        assert!(chunk.is_sorted());
        assert!(chunk.is_sorted_uncached());

        let alpha = chunk.timelines().get(&"alpha".into()).unwrap();
        let beta = chunk.timelines().get(&"beta".into()).unwrap();

        // The primary timeline (alphabetically first) must be globally sorted; the secondary
        // one is only sorted within each primary-key group.
        assert!(alpha.is_sorted());
        assert!(!beta.is_sorted());

        assert_eq!(alpha.times_raw().to_vec(), vec![1, 1, 1, 2, 2]);
        assert_eq!(beta.times_raw().to_vec(), vec![7, 9, 9, 3, 5]);

        // Verify the components were permuted in lockstep with the time columns,
        // and that ties (rows 1 and 4) preserved their original order (stable sort).
        let got_colors: Vec<u32> = chunk
            .iter_slices::<u32>(MyPoints::descriptor_colors().component)
            .flat_map(<[u32]>::to_vec)
            .collect();
        assert_eq!(got_colors, vec![300, 200, 500, 400, 100]);

        // RowIds must be sequential ascending.
        let row_ids: Vec<_> = chunk.row_ids().collect();
        for w in row_ids.windows(2) {
            assert!(w[0] < w[1], "row_ids must be strictly ascending");
        }

        Ok(())
    }
}