fix: zstd frame step width for alignment (#5730)

The beginning of each frame in zstd needs to align with the buffer
alignment.

---------

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

Author: 0ax1

encodings/zstd/src/array.rs

@@ -330,6 +330,7 @@ impl ZstdArray {
                 .get(i + 1)
                 .copied()
                 .unwrap_or(value_bytes.len());
+
             let uncompressed = &value_bytes.slice(frame_byte_starts[i]..frame_byte_end);
             let compressed = compressor
                 .compress(uncompressed)
@@ -366,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

@@ -9,6 +9,7 @@ use vortex_array::arrays::VarBinViewArray;
 use vortex_array::assert_arrays_eq;
 use vortex_array::validity::Validity;
 use vortex_array::vtable::ValidityHelper;
+use vortex_buffer::Alignment;
 use vortex_buffer::Buffer;
 use vortex_dtype::DType;
 use vortex_dtype::Nullability;
@@ -202,3 +203,16 @@ fn test_sliced_array_children() {
     let sliced = compressed.slice(0..4);
     sliced.children();
 }
+
+/// Tests that each beginning of a frame in ZSTD matches
+/// the buffer alignment when compressing primitive arrays.
+#[test]
+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);
+
+    assert!(compressed.is_ok());
+}