align frames to buffer

Signed-off-by: Alexander Droste <[email protected]>

Author: 0ax1

encodings/zstd/src/array.rs

@@ -331,10 +331,6 @@ impl ZstdArray {
                 .copied()
                 .unwrap_or(value_bytes.len());
 
-            // Use slice_unaligned instead of slice to handle cases where frame boundaries
-            // don't align with the buffer's alignment requirements (e.g., 64-byte aligned
-            // buffers with 1-byte frame boundaries). ZSTD compression works on raw bytes
-            // and doesn't require memory alignment.
             let uncompressed = &value_bytes.slice(frame_byte_starts[i]..frame_byte_end);
             let compressed = compressor
                 .compress(uncompressed)
@@ -371,8 +367,12 @@ impl ZstdArray {
         };
 
         let value_bytes = values.byte_buffer();
+        // Align frames to buffer alignment. This is necessary for overaligned buffers.
+        let alignment = *value_bytes.alignment();
+        let step_width = (values_per_frame * byte_width).div_ceil(alignment) * alignment;
+
         let frame_byte_starts = (0..n_values * byte_width)
-            .step_by(values_per_frame * byte_width)
+            .step_by(step_width)
             .collect::<Vec<_>>();
         let Frames {
             dictionary,

encodings/zstd/src/test.rs

@@ -204,13 +204,13 @@ fn test_sliced_array_children() {
     sliced.children();
 }
 
+/// Tests that each beginning of a frame in ZSTD matches
+/// the buffer alignment when compressing primitive arrays.
 #[test]
-fn test_zstd_alignment_issue() {
-    // Scenario where a u8 array has a 64-byte alignment.
-    // See: https://github.com/vortex-data/vortex/issues/5715
+fn test_zstd_frame_start_buffer_alignment() {
     let data = vec![0u8; 2];
-
     let aligned_buffer = Buffer::copy_from_aligned(&data, Alignment::new(8));
+    // u8 array now has a 8-byte alignment.
     let array = PrimitiveArray::new(aligned_buffer, Validity::NonNullable);
     let compressed = ZstdArray::from_primitive(&array, 0, 1);