Repository: MaterializeInc/materialize
Pull request #32094: timely-util: fix cloning `Column::Align`
Author: login: teskje
Created: 2025-04-04T16:56:13Z | Merged: 2025-04-08T07:52:20Z

The implementation of `Column::clone` for the `Align` variant would first allocate a zero-filled region with `alloc_aligned_zeroed`, and then attempt to _append_ the bytes from the original region to its end using `Region::extend_from_slice`. The correct thing to do is to `copy_from_slice` instead, overwriting the zeroes.

Furthermore, the `Column` implementation calls `copy_from_slice` on a region returned by `alloc_aligned_zeroed` in a couple places. The returned region can be larger than the requested length, in which case `copy_from_slice` would panic because it expects both source and destination to have the same length. This is fixed here by slicing the destination to the source length.

### Motivation

  * This PR fixes a recognized bug.

Fixes https://github.com/MaterializeInc/database-issues/issues/9088

### Checklist

- [x] This PR has adequate test coverage / QA involvement has been duly considered. ([trigger-ci for additional test/nightly runs](https://trigger-ci.dev.materialize.com/))
- [x] This PR has an associated up-to-date [design doc](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/design/README.md), is a design doc ([template](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/design/00000000_template.md)), or is sufficiently small to not require a design.
  <!-- Reference the design in the description. -->
- [x] If this PR evolves [an existing `$T ⇔ Proto$T` mapping](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/command-and-response-binary-encoding.md) (possibly in a backwards-incompatible way), then it is tagged with a `T-proto` label.
- [x] If this PR will require changes to cloud orchestration or tests, there is a companion cloud PR to account for those changes that is tagged with the release-blocker label ([example](https://github.com/MaterializeInc/cloud/pull/5021)).
  <!-- Ask in #team-cloud on Slack if you need help preparing the cloud PR. -->
- [x] If this PR includes major [user-facing behavior changes](https://github.com/MaterializeInc/materialize/blob/main/doc/developer/guide-changes.md#what-changes-require-a-release-note), I have pinged the relevant PM to schedule a changelog post.



Diff:
diff --git a/src/timely-util/Cargo.toml b/src/timely-util/Cargo.toml
index a8f0e568a54e1..4461cb172e666 100644
--- a/src/timely-util/Cargo.toml
+++ b/src/timely-util/Cargo.toml
@@ -19,7 +19,7 @@ differential-dataflow = "0.14.0"
 either = "1"
 futures-util = "0.3.31"
 lgalloc = "0.5"
-mz-ore = { path = "../ore", features = ["async", "process", "tracing", "test", "num-traits"] }
+mz-ore = { path = "../ore", features = ["async", "process", "tracing", "test", "num-traits", "region", "differential-dataflow"] }
 num-traits = "0.2"
 proptest = { version = "1.6.0", default-features = false, features = ["std"] }
 serde = { version = "1.0.219", features = ["derive"] }
diff --git a/src/timely-util/src/containers.rs b/src/timely-util/src/containers.rs
index a5ef6883998eb..c6bf7d05fba2d 100644
--- a/src/timely-util/src/containers.rs
+++ b/src/timely-util/src/containers.rs
@@ -126,12 +126,13 @@ mod container {
                 Column::Bytes(b) => {
                     assert_eq!(b.len() % 8, 0);
                     let mut alloc: Region<u64> = super::alloc_aligned_zeroed(b.len() / 8);
-                    bytemuck::cast_slice_mut(&mut alloc[..]).copy_from_slice(&b[..]);
+                    let alloc_bytes = bytemuck::cast_slice_mut(&mut alloc);
+                    alloc_bytes[..b.len()].copy_from_slice(b);
                     Self::Align(alloc)
                 }
                 Column::Align(a) => {
                     let mut alloc = super::alloc_aligned_zeroed(a.len());
-                    alloc.extend_from_slice(&a[..]);
+                    alloc[..a.len()].copy_from_slice(a);
                     Column::Align(alloc)
                 }
             }
@@ -198,7 +199,8 @@ mod container {
             } else {
                 // We failed to cast the slice, so we'll reallocate.
                 let mut alloc: Region<u64> = super::alloc_aligned_zeroed(bytes.len() / 8);
-                bytemuck::cast_slice_mut(&mut alloc[..]).copy_from_slice(&bytes[..]);
+                let alloc_bytes = bytemuck::cast_slice_mut(&mut alloc);
+                alloc_bytes[..bytes.len()].copy_from_slice(&bytes);
                 Self::Align(alloc)
             }
         }
@@ -474,3 +476,75 @@ mod provided_builder {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use mz_ore::region::Region;
+    use timely::bytes::arc::BytesMut;
+    use timely::dataflow::channels::ContainerBytes;
+    use timely::Container;
+
+    use super::*;
+
+    /// Produce some bytes that are in columnar format.
+    fn raw_columnar_bytes() -> Vec<u8> {
+        let mut raw = Vec::new();
+        raw.extend(12_u64.to_le_bytes()); // length
+        raw.extend(1_i32.to_le_bytes());
+        raw.extend(2_i32.to_le_bytes());
+        raw.extend(3_i32.to_le_bytes());
+        raw.extend([0, 0, 0, 0]); // padding
+        raw
+    }
+
+    #[mz_ore::test]
+    fn test_column_clone() {
+        let columns = Columnar::as_columns([1, 2, 3].iter());
+        let column_typed: Column<i32> = Column::Typed(columns);
+        let column_typed2 = column_typed.clone();
+
+        assert_eq!(column_typed2.iter().collect::<Vec<_>>(), vec![&1, &2, &3]);
+
+        let bytes = BytesMut::from(raw_columnar_bytes()).freeze();
+        let column_bytes: Column<i32> = Column::Bytes(bytes);
+        let column_bytes2 = column_bytes.clone();
+
+        assert_eq!(column_bytes2.iter().collect::<Vec<_>>(), vec![&1, &2, &3]);
+
+        let raw = raw_columnar_bytes();
+        let mut region: Region<u64> = alloc_aligned_zeroed(raw.len() / 8);
+        let region_bytes = bytemuck::cast_slice_mut(&mut region);
+        region_bytes[..raw.len()].copy_from_slice(&raw);
+        let column_align: Column<i32> = Column::Align(region);
+        let column_align2 = column_align.clone();
+
+        assert_eq!(column_align2.iter().collect::<Vec<_>>(), vec![&1, &2, &3]);
+    }
+
+    #[mz_ore::test]
+    fn test_column_from_bytes() {
+        let raw = raw_columnar_bytes();
+
+        let buf = vec![0; raw.len() + 8];
+        let align = buf.as_ptr().align_offset(std::mem::size_of::<u64>());
+        let mut bytes_mut = BytesMut::from(buf);
+        let _ = bytes_mut.extract_to(align);
+        bytes_mut[..raw.len()].copy_from_slice(&raw);
+        let aligned_bytes = bytes_mut.extract_to(raw.len());
+
+        let column: Column<i32> = Column::from_bytes(aligned_bytes);
+        assert!(matches!(column, Column::Bytes(_)));
+        assert_eq!(column.iter().collect::<Vec<_>>(), vec![&1, &2, &3]);
+
+        let buf = vec![0; raw.len() + 8];
+        let align = buf.as_ptr().align_offset(std::mem::size_of::<u64>());
+        let mut bytes_mut = BytesMut::from(buf);
+        let _ = bytes_mut.extract_to(align + 1);
+        bytes_mut[..raw.len()].copy_from_slice(&raw);
+        let unaligned_bytes = bytes_mut.extract_to(raw.len());
+
+        let column: Column<i32> = Column::from_bytes(unaligned_bytes);
+        assert!(matches!(column, Column::Align(_)));
+        assert_eq!(column.iter().collect::<Vec<_>>(), vec![&1, &2, &3]);
+    }
+}


Reviews:
user: login: teskje
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T16:57:34Z
user: login: teskje
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-04T16:58:13Z
user: login: antiguru
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-04T17:26:33Z
user: login: antiguru
body: Left some comments inline. Thanks very much for diagnosing and fixing!
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-04-07T11:30:32Z
user: login: teskje
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-07T12:52:17Z
user: login: teskje
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-07T12:54:19Z
user: login: antiguru
state: APPROVED
author_association: MEMBER
submitted_at: 2025-04-07T17:46:00Z
user: login: cursor[bot]
body: <!-- BUGBOT_REVIEW -->
✅ Bugbot reviewed your changes and found no new issues!

_Comment `@cursor review` or `bugbot run` to trigger another review on this PR_
state: COMMENTED
author_association: NONE
submitted_at: 2026-03-19T19:10:48Z

Review comments:
diff_hunk: @@ -474,3 +476,75 @@ mod provided_builder {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use mz_ore::region::Region;
+    use timely::bytes::arc::BytesMut;
+    use timely::dataflow::channels::ContainerBytes;
+    use timely::Container;
+
+    use super::*;
+
+    /// Produce some bytes that are in columnar format.
+    fn raw_columnar_bytes() -> Vec<u8> {
+        let mut raw = Vec::new();
+        raw.extend(12_u64.to_le_bytes()); // length
+        raw.extend(1_i32.to_le_bytes());
+        raw.extend(2_i32.to_le_bytes());
+        raw.extend(3_i32.to_le_bytes());
+        raw.extend([0, 0, 0, 0]); // padding
+        raw
path: src/timely-util/src/containers.rs
user: login: teskje
body: I tried using `Sequence::encode` here, but it gives me a `Vec<u64>` and I don't know how to convert that into a `Vec<u8>` without unsafe, so I just went the manual way.
created_at: 2025-04-04T16:57:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -474,3 +476,75 @@ mod provided_builder {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use mz_ore::region::Region;
+    use timely::bytes::arc::BytesMut;
+    use timely::dataflow::channels::ContainerBytes;
+    use timely::Container;
+
+    use super::*;
+
+    /// Produce some bytes that are in columnar format.
+    fn raw_columnar_bytes() -> Vec<u8> {
+        let mut raw = Vec::new();
+        raw.extend(12_u64.to_le_bytes()); // length
+        raw.extend(1_i32.to_le_bytes());
+        raw.extend(2_i32.to_le_bytes());
+        raw.extend(3_i32.to_le_bytes());
+        raw.extend([0, 0, 0, 0]); // padding
+        raw
+    }
+
+    #[test]
+    fn test_column_clone() {
+        let columns = Columnar::as_columns([1, 2, 3].iter());
+        let column_typed: Column<i32> = Column::Typed(columns);
+        let column_typed2 = column_typed.clone();
+
+        assert_eq!(column_typed2.iter().collect::<Vec<_>>(), vec![&1, &2, &3]);
+
+        let bytes = BytesMut::from(raw_columnar_bytes()).freeze();
+        let column_bytes: Column<i32> = Column::Bytes(bytes);
+        let column_bytes2 = column_bytes.clone();
+
+        assert_eq!(column_bytes2.iter().collect::<Vec<_>>(), vec![&1, &2, &3]);
+
+        let raw = raw_columnar_bytes();
+        let mut region: Region<u64> = alloc_aligned_zeroed(raw.len() / 8);
+        let region_bytes = bytemuck::cast_slice_mut(&mut region);
+        region_bytes[..raw.len()].copy_from_slice(&raw);
+        let column_align: Column<i32> = Column::Align(region);
+        let column_align2 = column_align.clone();
+
+        assert_eq!(column_align2.iter().collect::<Vec<_>>(), vec![&1, &2, &3]);
+    }
+
+    #[test]
+    fn test_column_from_bytes() {
+        let raw = raw_columnar_bytes();
+
+        let buf = vec![0; raw.len() + 8];
+        let align = buf.as_ptr().align_offset(std::mem::size_of::<u64>());
+        let mut bytes_mut = BytesMut::from(buf);
+        let _ = bytes_mut.extract_to(align);
+        bytes_mut[..raw.len()].copy_from_slice(&raw);
+        let aligned_bytes = bytes_mut.extract_to(raw.len());
+
+        let column: Column<i32> = Column::from_bytes(aligned_bytes);
+        assert!(matches!(column, Column::Bytes(_)));
+        assert_eq!(column.iter().collect::<Vec<_>>(), vec![&1, &2, &3]);
+
+        let buf = vec![0; raw.len() + 8];
+        let align = buf.as_ptr().align_offset(std::mem::size_of::<u64>());
+        let mut bytes_mut = BytesMut::from(buf);
+        let _ = bytes_mut.extract_to(align + 1);
+        bytes_mut[..raw.len()].copy_from_slice(&raw);
+        let unaligned_bytes = bytes_mut.extract_to(raw.len());
path: src/timely-util/src/containers.rs
user: login: teskje
body: Is there an easier way to get aligned/unaligned `Bytes`?
created_at: 2025-04-04T16:58:12Z
author_association: CONTRIBUTOR
diff_hunk: @@ -474,3 +476,75 @@ mod provided_builder {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use mz_ore::region::Region;
+    use timely::bytes::arc::BytesMut;
+    use timely::dataflow::channels::ContainerBytes;
+    use timely::Container;
+
+    use super::*;
+
+    /// Produce some bytes that are in columnar format.
+    fn raw_columnar_bytes() -> Vec<u8> {
+        let mut raw = Vec::new();
+        raw.extend(12_u64.to_le_bytes()); // length
+        raw.extend(1_i32.to_le_bytes());
+        raw.extend(2_i32.to_le_bytes());
+        raw.extend(3_i32.to_le_bytes());
+        raw.extend([0, 0, 0, 0]); // padding
+        raw
path: src/timely-util/src/containers.rs
user: login: antiguru
body: You could use bytemuck: https://docs.rs/bytemuck/latest/bytemuck/fn.cast_slice.html
created_at: 2025-04-04T17:26:32Z
author_association: MEMBER
diff_hunk: @@ -126,12 +126,13 @@ mod container {
                 Column::Bytes(b) => {
                     assert_eq!(b.len() % 8, 0);
                     let mut alloc: Region<u64> = super::alloc_aligned_zeroed(b.len() / 8);
-                    bytemuck::cast_slice_mut(&mut alloc[..]).copy_from_slice(&b[..]);
+                    let alloc_bytes = bytemuck::cast_slice_mut(&mut alloc);
+                    alloc_bytes[..b.len()].copy_from_slice(b);
path: src/timely-util/src/containers.rs
user: login: antiguru
body: This shouldn't be necessary, right? `alloc_bytes` should have capacity for all of `b`, so it should copy the whole of `b`, even if it has more capacity.
created_at: 2025-04-07T11:24:36Z
author_association: MEMBER
diff_hunk: @@ -198,7 +199,8 @@ mod container {
             } else {
                 // We failed to cast the slice, so we'll reallocate.
                 let mut alloc: Region<u64> = super::alloc_aligned_zeroed(bytes.len() / 8);
-                bytemuck::cast_slice_mut(&mut alloc[..]).copy_from_slice(&bytes[..]);
+                let alloc_bytes = bytemuck::cast_slice_mut(&mut alloc);
+                alloc_bytes[..bytes.len()].copy_from_slice(&bytes);
path: src/timely-util/src/containers.rs
user: login: antiguru
body: Here again, this should be needed, right?
created_at: 2025-04-07T11:25:46Z
author_association: MEMBER
diff_hunk: @@ -126,12 +126,13 @@ mod container {
                 Column::Bytes(b) => {
                     assert_eq!(b.len() % 8, 0);
                     let mut alloc: Region<u64> = super::alloc_aligned_zeroed(b.len() / 8);
-                    bytemuck::cast_slice_mut(&mut alloc[..]).copy_from_slice(&b[..]);
+                    let alloc_bytes = bytemuck::cast_slice_mut(&mut alloc);
+                    alloc_bytes[..b.len()].copy_from_slice(b);
                     Self::Align(alloc)
                 }
                 Column::Align(a) => {
                     let mut alloc = super::alloc_aligned_zeroed(a.len());
-                    alloc.extend_from_slice(&a[..]);
+                    alloc[..a.len()].copy_from_slice(a);
path: src/timely-util/src/containers.rs
user: login: antiguru
body: It's a bug, thanks for catching!

We might not want to call `alloc_aligned_zeroed` here because it sets the memory twice. Instead, we'd want to allocate an uninitialized region and use `extend_from_slice` as before. This needs a new API because `alloc_aligned_zeroed` checks whether to use lgalloc or not, and we don't have an equivalent non-initializing variant.
created_at: 2025-04-07T11:29:04Z
author_association: MEMBER
diff_hunk: @@ -126,12 +126,13 @@ mod container {
                 Column::Bytes(b) => {
                     assert_eq!(b.len() % 8, 0);
                     let mut alloc: Region<u64> = super::alloc_aligned_zeroed(b.len() / 8);
-                    bytemuck::cast_slice_mut(&mut alloc[..]).copy_from_slice(&b[..]);
+                    let alloc_bytes = bytemuck::cast_slice_mut(&mut alloc);
+                    alloc_bytes[..b.len()].copy_from_slice(b);
path: src/timely-util/src/containers.rs
user: login: teskje
body: `copy_from_slice` panics if the slices have different lengths ([docs](https://doc.rust-lang.org/std/primitive.slice.html#method.copy_from_slice)). Since `alloc_bytes` can be longer than `b`, we need to truncate it to avoid the panic.
created_at: 2025-04-07T12:52:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -126,12 +126,13 @@ mod container {
                 Column::Bytes(b) => {
                     assert_eq!(b.len() % 8, 0);
                     let mut alloc: Region<u64> = super::alloc_aligned_zeroed(b.len() / 8);
-                    bytemuck::cast_slice_mut(&mut alloc[..]).copy_from_slice(&b[..]);
+                    let alloc_bytes = bytemuck::cast_slice_mut(&mut alloc);
+                    alloc_bytes[..b.len()].copy_from_slice(b);
                     Self::Align(alloc)
                 }
                 Column::Align(a) => {
                     let mut alloc = super::alloc_aligned_zeroed(a.len());
-                    alloc.extend_from_slice(&a[..]);
+                    alloc[..a.len()].copy_from_slice(a);
path: src/timely-util/src/containers.rs
user: login: teskje
body: Makes sense! I'd like to treat that out of scope for this PR though.
created_at: 2025-04-07T12:54:19Z
author_association: CONTRIBUTOR

PR comments:
user: login: antiguru
created_at: 2025-04-04T17:30:30Z
body: Thank you! I'll check in more detail on Monday. We shouldn't use this path as we're not using the zero-copy allocator and don't enable lgalloc for columnar if I'm not mistaken.
author_association: MEMBER
user: login: teskje
created_at: 2025-04-05T11:56:03Z
body: > We shouldn't use this path as we're not using the zero-copy allocator and don't enable lgalloc for columnar if I'm not mistaken.

The `enable_columnar_lgalloc` dyncfg defaults to `true` and we're not setting a different value in LD, so I think we do enable it. The [panic](https://materializeinc.sentry.io/issues/6401706298/?project=6780145&referrer=github_integration) only occurred two times in @jubrad's staging environment, so it seems like a path that we hit very rarely. Not sure what starts need to align for that.
author_association: CONTRIBUTOR
user: login: teskje
created_at: 2025-04-08T07:52:17Z
body: TFTR!
author_association: CONTRIBUTOR
user: login: bosconi
created_at: 2026-03-19T19:05:54Z
body: bugbot run
author_association: MEMBER
user: login: cursor[bot]
created_at: 2026-03-19T19:06:03Z
body: ## PR Summary

<!-- CURSOR_SUMMARY -->
**Low Risk**
Low risk: fixes a clear copy/clone bug and adds targeted unit tests; behavior changes are limited to internal columnar buffer copying and alignment handling.

**Overview**
Fixes `Column` cloning and `from_bytes` relocation to **copy into aligned `Region<u64>` buffers correctly**: `Column::Align` now overwrites the pre-zeroed allocation instead of appending, and byte copies into newly allocated regions are sliced to the source length to avoid panics when the region is larger than requested.

Adds unit tests covering cloning for `Typed`/`Bytes`/`Align` variants and `from_bytes` behavior for aligned vs unaligned input, and enables the needed `mz-ore` features (`region`, `differential-dataflow`) for these tests.

<sup>Written by [Cursor Bugbot](https://cursor.com/dashboard?tab=bugbot) for commit 96a85b3a06c0f88f8786b1b3f7e3e6b4292c492f. This will update automatically on new commits. Configure [here](https://cursor.com/dashboard?tab=bugbot).</sup>
<!-- /CURSOR_SUMMARY -->
author_association: NONE