|
1 | 1 | // SPDX-License-Identifier: Apache-2.0 |
2 | 2 | // SPDX-FileCopyrightText: Copyright the Vortex contributors |
3 | 3 |
|
4 | | -use rand::SeedableRng; |
5 | | -use rand::rngs::StdRng; |
6 | | -use rand::seq::IndexedRandom; |
7 | 4 | use vortex_array::ArrayRef; |
8 | 5 | use vortex_array::IntoArray; |
9 | 6 | use vortex_array::LEGACY_SESSION; |
10 | 7 | use vortex_array::VortexSessionExecute; |
11 | 8 | use vortex_array::arrays::VarBinViewArray; |
| 9 | +use vortex_array::arrays::varbin::VarBinArrayExt; |
12 | 10 | use vortex_array::arrays::varbin::builder::VarBinBuilder; |
13 | 11 | use vortex_array::assert_arrays_eq; |
14 | 12 | use vortex_array::assert_nth_scalar; |
15 | 13 | use vortex_array::dtype::DType; |
16 | 14 | use vortex_array::dtype::Nullability; |
| 15 | +use vortex_array::dtype::PType; |
17 | 16 | use vortex_buffer::buffer; |
18 | 17 | use vortex_mask::Mask; |
19 | 18 |
|
20 | 19 | use crate::FSST; |
| 20 | +use crate::FSSTArrayExt; |
21 | 21 | use crate::fsst_compress; |
22 | 22 | use crate::fsst_train_compressor; |
23 | 23 |
|
@@ -111,60 +111,28 @@ fn test_fsst_array_ops() { |
111 | 111 | assert_arrays_eq!(fsst_array, canonical_array); |
112 | 112 | } |
113 | 113 |
|
114 | | -// TODO(someone): ideally CI would run this in release mode as well since debug builds make the |
115 | | -// allocation and compression loop substantially slower. |
116 | | -/// Regression for #7833: [`fsst_compress`] must accept inputs whose cumulative compressed |
117 | | -/// bytes exceed [`i32::MAX`]. Before the fix, [`fsst_compress_iter`] hardcoded |
118 | | -/// [`VarBinBuilder<i32>`] for the FSST output and panicked in |
119 | | -/// [`VarBinBuilder::append_value`] once cumulative compressed bytes crossed the boundary. |
| 114 | +/// Regression coverage for #7833. |
120 | 115 | /// |
121 | | -/// The input is built with [`VarBinBuilder<i64>`] so the input itself does not panic, which |
122 | | -/// confirms the overflow is on the FSST output side. After the fix the test must succeed |
123 | | -/// with the row count preserved. |
124 | | -/// |
125 | | -/// Allocates ~2.5 GiB for the input and ~2.5 GiB for the FSST output (~5 GiB total), so it |
126 | | -/// is gated to CI runs and skipped when `VORTEX_SKIP_SLOW_TESTS` is set. To run it locally: |
127 | | -/// |
128 | | -/// ```text |
129 | | -/// CI=1 cargo test --release -p vortex-fsst fsst_compress_offsets |
130 | | -/// ``` |
| 116 | +/// The bug was that [`fsst_compress_iter`] built compressed codes with [`VarBinBuilder<i32>`], |
| 117 | +/// which panicked once cumulative compressed bytes crossed [`i32::MAX`]. The end-to-end |
| 118 | +/// behavior is "large FSST outputs do not overflow", but exercising that literally requires |
| 119 | +/// multi-GiB input. This test checks the implementation detail that currently guarantees that |
| 120 | +/// behavior: compressed code offsets are widened to [`i64`]. |
131 | 121 | /// |
132 | 122 | /// [`fsst_compress_iter`]: crate::compress::fsst_compress_iter |
133 | | -#[test_with::env(CI)] |
134 | | -#[test_with::no_env(VORTEX_SKIP_SLOW_TESTS)] |
135 | | -fn fsst_compress_offsets_overflow_i32() { |
136 | | - // High-entropy ASCII strings sliced from a random pool. FSST is a symbol-table |
137 | | - // compressor; pseudo-random data with no recurring byte sequences resists compression, |
138 | | - // so the compressed output stays close to input size and crosses the i32 boundary. |
139 | | - const STRING_LEN: usize = 64 * 1024; |
140 | | - const TOTAL_BYTES: usize = (1usize << 31) + (512 << 20); // ~2.5 GiB |
141 | | - const N: usize = TOTAL_BYTES / STRING_LEN; |
142 | | - const POOL_LEN: usize = 64 * 1024 * 1024; |
143 | | - |
144 | | - // Printable ASCII alphabet so the result is valid UTF-8. |
145 | | - const ALPHABET: &[u8; 95] = |
146 | | - b" !\"#$%&'()*+,-./0123456789:;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]^_`abcdefghijklmnopqrstuvwxyz{|}~"; |
147 | | - |
148 | | - let mut rng = StdRng::seed_from_u64(0xC0DE_C011_B711); |
149 | | - let pool: Vec<u8> = (0..POOL_LEN) |
150 | | - .map(|_| *ALPHABET.choose(&mut rng).unwrap()) |
151 | | - .collect(); |
152 | | - |
153 | | - println!("building large VarBinArray"); |
154 | | - let mut builder = VarBinBuilder::<i64>::with_capacity(N); |
155 | | - for i in 0..N { |
156 | | - let off = i.wrapping_mul(31337) % (POOL_LEN - STRING_LEN); |
157 | | - builder.append_value(&pool[off..off + STRING_LEN]); |
158 | | - } |
| 123 | +#[test] |
| 124 | +fn fsst_compress_uses_i64_code_offsets() { |
| 125 | + let mut builder = VarBinBuilder::<i32>::with_capacity(2); |
| 126 | + builder.append_value(b"hello world"); |
| 127 | + builder.append_value(b"hello vortex"); |
159 | 128 | let array = builder.finish(DType::Utf8(Nullability::NonNullable)); |
160 | 129 |
|
161 | | - println!("training FSST compressor"); |
162 | 130 | let compressor = fsst_train_compressor(&array); |
163 | 131 | let len = array.len(); |
164 | 132 | let dtype = array.dtype().clone(); |
165 | 133 | let mut ctx = LEGACY_SESSION.create_execution_ctx(); |
166 | 134 |
|
167 | | - println!("compressing to FSST"); |
168 | 135 | let compressed = fsst_compress(array, len, &dtype, &compressor, &mut ctx); |
169 | 136 | assert_eq!(compressed.len(), len); |
| 137 | + assert_eq!(compressed.codes().offsets().dtype().as_ptype(), PType::I64); |
170 | 138 | } |
0 commit comments