Skip to content

Commit ceb5796

Browse files
mprammerclaude
authored and
robert3005
committed
refactor[fsst]: take ArrayRef in fsst_compress, decide offset width upfront
fsst_compress and fsst_train_compressor take an ArrayRef and dispatch internally on VarBinView/VarBin encoding, so callers don't need to know which string encoding they hold. The codes-offsets width (i32 vs i64) is chosen upfront from the FSST worst-case bound (2 * input bytes + 7 per non-null row), which always covers the actual output. Follow-up to #7853 addressing review feedback on #7833. Signed-off-by: mprammer <martin@spiraldb.com> Signed-off-by: Robert Kruszewski <github@robertk.io> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent 0dd6db7 commit ceb5796

23 files changed

Lines changed: 636 additions & 540 deletions

File tree

‎Cargo.lock‎

Lines changed: 2 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎encodings/fsst/Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ workspace = true
1818

1919
[dependencies]
2020
fsst-rs = { workspace = true }
21+
num-traits = { workspace = true }
2122
prost = { workspace = true }
2223
rand = { workspace = true, optional = true }
2324
vortex-array = { workspace = true }

‎encodings/fsst/benches/fsst_compress.rs‎

Lines changed: 29 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -53,28 +53,22 @@ const BENCH_ARGS: &[(usize, usize, u8)] = &[
5353

5454
#[divan::bench(args = BENCH_ARGS)]
5555
fn compress_fsst(bencher: Bencher, (string_count, avg_len, unique_chars): (usize, usize, u8)) {
56-
let array = generate_test_data(string_count, avg_len, unique_chars);
57-
let compressor = fsst_train_compressor(&array);
56+
let array = generate_test_data(string_count, avg_len, unique_chars).into_array();
57+
let compressor =
58+
fsst_train_compressor(array.clone(), &mut SESSION.create_execution_ctx()).unwrap();
5859
bencher
59-
.with_inputs(|| (&array, &compressor, SESSION.create_execution_ctx()))
60+
.with_inputs(|| (array.clone(), &compressor, SESSION.create_execution_ctx()))
6061
.bench_refs(|(array, compressor, ctx)| {
61-
fsst_compress(*array, array.len(), array.dtype(), compressor, ctx)
62+
fsst_compress(array.clone(), compressor, ctx).unwrap()
6263
})
6364
}
6465

6566
#[divan::bench(args = BENCH_ARGS)]
6667
fn decompress_fsst(bencher: Bencher, (string_count, avg_len, unique_chars): (usize, usize, u8)) {
67-
let array = generate_test_data(string_count, avg_len, unique_chars);
68-
let compressor = fsst_train_compressor(&array);
69-
let len = array.len();
70-
let dtype = array.dtype().clone();
71-
let encoded = fsst_compress(
72-
array,
73-
len,
74-
&dtype,
75-
&compressor,
76-
&mut SESSION.create_execution_ctx(),
77-
);
68+
let array = generate_test_data(string_count, avg_len, unique_chars).into_array();
69+
let mut ctx = SESSION.create_execution_ctx();
70+
let compressor = fsst_train_compressor(array.clone(), &mut ctx).unwrap();
71+
let encoded = fsst_compress(array, &compressor, &mut ctx).unwrap();
7872

7973
bencher
8074
.with_inputs(|| (&encoded, SESSION.create_execution_ctx()))
@@ -83,24 +77,20 @@ fn decompress_fsst(bencher: Bencher, (string_count, avg_len, unique_chars): (usi
8377

8478
#[divan::bench(args = BENCH_ARGS)]
8579
fn train_compressor(bencher: Bencher, (string_count, avg_len, unique_chars): (usize, usize, u8)) {
86-
let array = generate_test_data(string_count, avg_len, unique_chars);
80+
let array = generate_test_data(string_count, avg_len, unique_chars).into_array();
8781
bencher
88-
.with_inputs(|| &array)
89-
.bench_refs(|array| fsst_train_compressor(array))
82+
.with_inputs(|| (array.clone(), SESSION.create_execution_ctx()))
83+
.bench_refs(|(array, ctx)| fsst_train_compressor(array.clone(), ctx).unwrap())
9084
}
9185

9286
#[divan::bench(args = BENCH_ARGS)]
9387
fn pushdown_compare(bencher: Bencher, (string_count, avg_len, unique_chars): (usize, usize, u8)) {
94-
let array = generate_test_data(string_count, avg_len, unique_chars);
95-
let compressor = fsst_train_compressor(&array);
96-
let fsst_array = fsst_compress(
97-
&array,
98-
array.len(),
99-
array.dtype(),
100-
&compressor,
101-
&mut SESSION.create_execution_ctx(),
102-
);
103-
let constant = ConstantArray::new(Scalar::from(&b"const"[..]), array.len());
88+
let array = generate_test_data(string_count, avg_len, unique_chars).into_array();
89+
let len = array.len();
90+
let mut ctx = SESSION.create_execution_ctx();
91+
let compressor = fsst_train_compressor(array.clone(), &mut ctx).unwrap();
92+
let fsst_array = fsst_compress(array, &compressor, &mut ctx).unwrap();
93+
let constant = ConstantArray::new(Scalar::from(&b"const"[..]), len);
10494

10595
bencher
10696
.with_inputs(|| (&fsst_array, &constant, SESSION.create_execution_ctx()))
@@ -120,16 +110,12 @@ fn canonicalize_compare(
120110
bencher: Bencher,
121111
(string_count, avg_len, unique_chars): (usize, usize, u8),
122112
) {
123-
let array = generate_test_data(string_count, avg_len, unique_chars);
124-
let compressor = fsst_train_compressor(&array);
125-
let fsst_array = fsst_compress(
126-
&array,
127-
array.len(),
128-
array.dtype(),
129-
&compressor,
130-
&mut SESSION.create_execution_ctx(),
131-
);
132-
let constant = ConstantArray::new(Scalar::from(&b"const"[..]), array.len());
113+
let array = generate_test_data(string_count, avg_len, unique_chars).into_array();
114+
let len = array.len();
115+
let mut ctx = SESSION.create_execution_ctx();
116+
let compressor = fsst_train_compressor(array.clone(), &mut ctx).unwrap();
117+
let fsst_array = fsst_compress(array, &compressor, &mut ctx).unwrap();
118+
let constant = ConstantArray::new(Scalar::from(&b"const"[..]), len);
133119

134120
bencher
135121
.with_inputs(|| (&fsst_array, &constant, SESSION.create_execution_ctx()))
@@ -225,11 +211,11 @@ fn generate_chunked_test_data(
225211
let mut ctx = SESSION.create_execution_ctx();
226212
(0..chunk_size)
227213
.map(|_| {
228-
let array = generate_test_data(string_count, avg_len, unique_chars);
229-
let compressor = fsst_train_compressor(&array);
230-
let len = array.len();
231-
let dtype = array.dtype().clone();
232-
fsst_compress(array, len, &dtype, &compressor, &mut ctx).into_array()
214+
let array = generate_test_data(string_count, avg_len, unique_chars).into_array();
215+
let compressor = fsst_train_compressor(array.clone(), &mut ctx).unwrap();
216+
fsst_compress(array, &compressor, &mut ctx)
217+
.unwrap()
218+
.into_array()
233219
})
234220
.collect::<ChunkedArray>()
235221
}

‎encodings/fsst/benches/fsst_url_compare.rs‎

Lines changed: 81 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ use vortex_array::RecursiveCanonical;
1212
use vortex_array::VortexSessionExecute;
1313
use vortex_array::arrays::ConstantArray;
1414
use vortex_array::arrays::VarBinArray;
15+
use vortex_array::arrays::VarBinViewArray;
1516
use vortex_array::builtins::ArrayBuiltins;
1617
use vortex_array::expr::like;
1718
use vortex_array::expr::lit;
@@ -38,6 +39,16 @@ const NUM_URLS: usize = NUM_STRINGS;
3839

3940
static URL_DATA: LazyLock<VarBinArray> = LazyLock::new(generate_url_data);
4041

42+
/// `URL_DATA` as a `VarBinViewArray`, used by the view-input bench variants.
43+
static URL_VIEW_DATA: LazyLock<VarBinViewArray> = LazyLock::new(|| {
44+
let mut ctx = SESSION.create_execution_ctx();
45+
URL_DATA
46+
.clone()
47+
.into_array()
48+
.execute::<VarBinViewArray>(&mut ctx)
49+
.unwrap()
50+
});
51+
4152
// ---------------------------------------------------------------------------
4253
// Eq compare benchmarks (FSST pushdown vs canonicalize)
4354
// ---------------------------------------------------------------------------
@@ -56,14 +67,10 @@ fn pick_url_with_domain(data: &VarBinArray, domain: &str) -> String {
5667
#[divan::bench]
5768
fn eq_pushdown_high_match(bencher: Bencher) {
5869
let data = &*URL_DATA;
59-
let compressor = fsst_train_compressor(data);
60-
let fsst_array = fsst_compress(
61-
data,
62-
data.len(),
63-
data.dtype(),
64-
&compressor,
65-
&mut SESSION.create_execution_ctx(),
66-
);
70+
let mut ctx = SESSION.create_execution_ctx();
71+
let array = data.clone().into_array();
72+
let compressor = fsst_train_compressor(array.clone(), &mut ctx).unwrap();
73+
let fsst_array = fsst_compress(array, &compressor, &mut ctx).unwrap();
6774
let match_url = pick_url_with_domain(data, HIGH_MATCH_DOMAIN);
6875
let constant = ConstantArray::new(Scalar::from(match_url.as_str()), NUM_URLS);
6976

@@ -83,14 +90,10 @@ fn eq_pushdown_high_match(bencher: Bencher) {
8390
#[divan::bench]
8491
fn eq_pushdown_low_match(bencher: Bencher) {
8592
let data = &*URL_DATA;
86-
let compressor = fsst_train_compressor(data);
87-
let fsst_array = fsst_compress(
88-
data,
89-
data.len(),
90-
data.dtype(),
91-
&compressor,
92-
&mut SESSION.create_execution_ctx(),
93-
);
93+
let mut ctx = SESSION.create_execution_ctx();
94+
let array = data.clone().into_array();
95+
let compressor = fsst_train_compressor(array.clone(), &mut ctx).unwrap();
96+
let fsst_array = fsst_compress(array, &compressor, &mut ctx).unwrap();
9497
let match_url = pick_url_with_domain(data, LOW_MATCH_DOMAIN);
9598
let constant = ConstantArray::new(Scalar::from(match_url.as_str()), NUM_URLS);
9699

@@ -107,17 +110,59 @@ fn eq_pushdown_low_match(bencher: Bencher) {
107110
});
108111
}
109112

113+
#[divan::bench]
114+
fn eq_pushdown_high_match_view(bencher: Bencher) {
115+
let data = &*URL_VIEW_DATA;
116+
let mut ctx = SESSION.create_execution_ctx();
117+
let array = data.clone().into_array();
118+
let compressor = fsst_train_compressor(array.clone(), &mut ctx).unwrap();
119+
let fsst_array = fsst_compress(array, &compressor, &mut ctx).unwrap();
120+
let match_url = pick_url_with_domain(&URL_DATA, HIGH_MATCH_DOMAIN);
121+
let constant = ConstantArray::new(Scalar::from(match_url.as_str()), NUM_URLS);
122+
123+
bencher
124+
.with_inputs(|| (&fsst_array, &constant, SESSION.create_execution_ctx()))
125+
.bench_refs(|(fsst_array, constant, ctx)| {
126+
fsst_array
127+
.clone()
128+
.into_array()
129+
.binary(constant.clone().into_array(), Operator::Eq)
130+
.unwrap()
131+
.execute::<RecursiveCanonical>(ctx)
132+
.unwrap()
133+
});
134+
}
135+
136+
#[divan::bench]
137+
fn eq_pushdown_low_match_view(bencher: Bencher) {
138+
let data = &*URL_VIEW_DATA;
139+
let mut ctx = SESSION.create_execution_ctx();
140+
let array = data.clone().into_array();
141+
let compressor = fsst_train_compressor(array.clone(), &mut ctx).unwrap();
142+
let fsst_array = fsst_compress(array, &compressor, &mut ctx).unwrap();
143+
let match_url = pick_url_with_domain(&URL_DATA, LOW_MATCH_DOMAIN);
144+
let constant = ConstantArray::new(Scalar::from(match_url.as_str()), NUM_URLS);
145+
146+
bencher
147+
.with_inputs(|| (&fsst_array, &constant, SESSION.create_execution_ctx()))
148+
.bench_refs(|(fsst_array, constant, ctx)| {
149+
fsst_array
150+
.clone()
151+
.into_array()
152+
.binary(constant.clone().into_array(), Operator::Eq)
153+
.unwrap()
154+
.execute::<RecursiveCanonical>(ctx)
155+
.unwrap()
156+
});
157+
}
158+
110159
#[divan::bench]
111160
fn eq_canonicalize_high_match(bencher: Bencher) {
112161
let data = &*URL_DATA;
113-
let compressor = fsst_train_compressor(data);
114-
let fsst_array = fsst_compress(
115-
data,
116-
data.len(),
117-
data.dtype(),
118-
&compressor,
119-
&mut SESSION.create_execution_ctx(),
120-
);
162+
let mut ctx = SESSION.create_execution_ctx();
163+
let array = data.clone().into_array();
164+
let compressor = fsst_train_compressor(array.clone(), &mut ctx).unwrap();
165+
let fsst_array = fsst_compress(array, &compressor, &mut ctx).unwrap();
121166
let match_url = pick_url_with_domain(data, HIGH_MATCH_DOMAIN);
122167
let constant = ConstantArray::new(Scalar::from(match_url.as_str()), NUM_URLS);
123168

@@ -140,14 +185,10 @@ fn eq_canonicalize_high_match(bencher: Bencher) {
140185
#[divan::bench]
141186
fn eq_canonicalize_low_match(bencher: Bencher) {
142187
let data = &*URL_DATA;
143-
let compressor = fsst_train_compressor(data);
144-
let fsst_array = fsst_compress(
145-
data,
146-
data.len(),
147-
data.dtype(),
148-
&compressor,
149-
&mut SESSION.create_execution_ctx(),
150-
);
188+
let mut ctx = SESSION.create_execution_ctx();
189+
let array = data.clone().into_array();
190+
let compressor = fsst_train_compressor(array.clone(), &mut ctx).unwrap();
191+
let fsst_array = fsst_compress(array, &compressor, &mut ctx).unwrap();
151192
let match_url = pick_url_with_domain(data, LOW_MATCH_DOMAIN);
152193
let constant = ConstantArray::new(Scalar::from(match_url.as_str()), NUM_URLS);
153194

@@ -174,14 +215,10 @@ fn eq_canonicalize_low_match(bencher: Bencher) {
174215
#[divan::bench]
175216
fn like_substr_high_match(bencher: Bencher) {
176217
let data = &*URL_DATA;
177-
let compressor = fsst_train_compressor(data);
178-
let fsst_array = fsst_compress(
179-
data,
180-
data.len(),
181-
data.dtype(),
182-
&compressor,
183-
&mut SESSION.create_execution_ctx(),
184-
);
218+
let mut ctx = SESSION.create_execution_ctx();
219+
let array = data.clone().into_array();
220+
let compressor = fsst_train_compressor(array.clone(), &mut ctx).unwrap();
221+
let fsst_array = fsst_compress(array, &compressor, &mut ctx).unwrap();
185222
let pattern = format!("%{HIGH_MATCH_DOMAIN}%");
186223
let expr = like(root(), lit(pattern.as_str()));
187224

@@ -201,14 +238,10 @@ fn like_substr_high_match(bencher: Bencher) {
201238
#[divan::bench]
202239
fn like_substr_low_match(bencher: Bencher) {
203240
let data = &*URL_DATA;
204-
let compressor = fsst_train_compressor(data);
205-
let fsst_array = fsst_compress(
206-
data,
207-
data.len(),
208-
data.dtype(),
209-
&compressor,
210-
&mut SESSION.create_execution_ctx(),
211-
);
241+
let mut ctx = SESSION.create_execution_ctx();
242+
let array = data.clone().into_array();
243+
let compressor = fsst_train_compressor(array.clone(), &mut ctx).unwrap();
244+
let fsst_array = fsst_compress(array, &compressor, &mut ctx).unwrap();
212245
let pattern = format!("%{LOW_MATCH_DOMAIN}%");
213246
let expr = like(root(), lit(pattern.as_str()));
214247

‎encodings/fsst/src/array.rs‎

Lines changed: 5 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -866,7 +866,7 @@ mod test {
866866
use crate::FSST;
867867
use crate::array::FSSTArrayExt;
868868
use crate::array::FSSTMetadata;
869-
use crate::fsst_compress_iter;
869+
use crate::fsst_compress;
870870

871871
#[test]
872872
fn slice_reuses_initialized_compressor() -> VortexResult<()> {
@@ -878,18 +878,8 @@ mod test {
878878

879879
let compressor = Compressor::rebuild_from(symbols.as_slice(), symbol_lengths.as_slice());
880880
let mut ctx = LEGACY_SESSION.create_execution_ctx();
881-
let fsst_array = fsst_compress_iter(
882-
[
883-
Some(b"abcabcab".as_ref()),
884-
Some(b"defghijk".as_ref()),
885-
Some(b"abcxyz".as_ref()),
886-
]
887-
.into_iter(),
888-
3,
889-
DType::Utf8(Nullability::NonNullable),
890-
&compressor,
891-
&mut ctx,
892-
);
881+
let strings = VarBinViewArray::from_iter_str(["abcabcab", "defghijk", "abcxyz"]);
882+
let fsst_array = fsst_compress(strings.into_array(), &compressor, &mut ctx)?;
893883

894884
let compressor_ptr = fsst_array.compressor() as *const Compressor;
895885
let sliced = fsst_array
@@ -932,13 +922,8 @@ mod test {
932922

933923
let compressor = Compressor::rebuild_from(symbols.as_slice(), symbol_lengths.as_slice());
934924
let mut ctx = LEGACY_SESSION.create_execution_ctx();
935-
let fsst_array = fsst_compress_iter(
936-
[Some(b"abcabcab".as_ref()), Some(b"defghijk".as_ref())].into_iter(),
937-
2,
938-
DType::Utf8(Nullability::NonNullable),
939-
&compressor,
940-
&mut ctx,
941-
);
925+
let input = VarBinViewArray::from_iter_str(["abcabcab", "defghijk"]);
926+
let fsst_array = fsst_compress(input.into_array(), &compressor, &mut ctx).unwrap();
942927

943928
let compressed_codes = fsst_array.codes();
944929

0 commit comments

Comments
 (0)