1use std::fmt::Debug;
5use std::fmt::Display;
6use std::fmt::Formatter;
7use std::hash::Hasher;
8use std::mem::MaybeUninit;
9use std::sync::Arc;
10use std::sync::OnceLock;
11
12use fsst::Compressor;
13use fsst::Decompressor;
14use fsst::Symbol;
15use num_traits::AsPrimitive;
16use prost::Message as _;
17use vortex_array::Array;
18use vortex_array::ArrayEq;
19use vortex_array::ArrayHash;
20use vortex_array::ArrayId;
21use vortex_array::ArrayParts;
22use vortex_array::ArrayRef;
23use vortex_array::ArraySlots;
24use vortex_array::ArrayView;
25use vortex_array::EqMode;
26use vortex_array::ExecutionCtx;
27use vortex_array::ExecutionResult;
28use vortex_array::TypedArrayRef;
29use vortex_array::VortexSessionExecute;
30use vortex_array::array_slots;
31use vortex_array::arrays::VarBin;
32use vortex_array::arrays::VarBinArray;
33use vortex_array::arrays::varbin::VarBinArraySlotsExt;
34use vortex_array::buffer::BufferHandle;
35use vortex_array::builders::ArrayBuilder;
36use vortex_array::builders::VarBinBuilder;
37use vortex_array::builders::VarBinViewBuilder;
38use vortex_array::dtype::DType;
39use vortex_array::dtype::Nullability;
40use vortex_array::dtype::OffsetBuilderPType;
41use vortex_array::dtype::PType;
42use vortex_array::legacy_session;
43use vortex_array::match_each_integer_ptype;
44use vortex_array::match_each_varbin_builder;
45use vortex_array::serde::ArrayChildren;
46use vortex_array::validity::Validity;
47use vortex_array::vtable::VTable;
48use vortex_array::vtable::ValidityVTable;
49use vortex_array::vtable::child_to_validity;
50use vortex_array::vtable::validity_to_child;
51use vortex_buffer::Buffer;
52use vortex_buffer::BufferMut;
53use vortex_buffer::ByteBuffer;
54use vortex_error::VortexExpect;
55use vortex_error::VortexResult;
56use vortex_error::vortex_bail;
57use vortex_error::vortex_ensure;
58use vortex_error::vortex_err;
59use vortex_error::vortex_panic;
60use vortex_session::VortexSession;
61use vortex_session::registry::CachedId;
62
63use crate::canonical::FSST_DECODE_SLACK;
64use crate::canonical::FsstDecodePlan;
65use crate::canonical::canonicalize_fsst;
66use crate::canonical::fsst_decode_bytes;
67use crate::rules::RULES;
68
69pub type FSSTArray = Array<FSST>;
71
72#[derive(Clone, prost::Message)]
73pub struct FSSTMetadata {
74 #[prost(enumeration = "PType", tag = "1")]
75 uncompressed_lengths_ptype: i32,
76
77 #[prost(enumeration = "PType", tag = "2")]
78 codes_offsets_ptype: i32,
79}
80
81impl FSSTMetadata {
82 pub fn get_uncompressed_lengths_ptype(&self) -> VortexResult<PType> {
83 PType::try_from(self.uncompressed_lengths_ptype)
84 .map_err(|_| vortex_err!("Invalid PType {}", self.uncompressed_lengths_ptype))
85 }
86}
87
88pub const FSST_SYMBOL_TABLE_LEN: usize = 255;
94
95impl ArrayHash for FSSTData {
96 fn array_hash<H: Hasher>(&self, state: &mut H, precision: EqMode) {
97 self.padded_symbols().array_hash(state, precision);
98 self.padded_symbol_lengths().array_hash(state, precision);
99 self.codes_bytes.as_host().array_hash(state, precision);
100 }
101}
102
103impl ArrayEq for FSSTData {
104 fn array_eq(&self, other: &Self, precision: EqMode) -> bool {
105 self.padded_symbols()
106 .array_eq(other.padded_symbols(), precision)
107 && self
108 .padded_symbol_lengths()
109 .array_eq(other.padded_symbol_lengths(), precision)
110 && self
111 .codes_bytes
112 .as_host()
113 .array_eq(other.codes_bytes.as_host(), precision)
114 }
115}
116
117impl VTable for FSST {
118 type TypedArrayData = FSSTData;
119 type OperationsVTable = Self;
120 type ValidityVTable = Self;
121
122 fn id(&self) -> ArrayId {
123 static ID: CachedId = CachedId::new("vortex.fsst");
124 *ID
125 }
126
127 #[allow(clippy::disallowed_methods)]
128 fn validate(
129 &self,
130 data: &Self::TypedArrayData,
131 dtype: &DType,
132 len: usize,
133 slots: &[Option<ArrayRef>],
134 ) -> VortexResult<()> {
135 let mut ctx = legacy_session().create_execution_ctx();
137 data.validate(dtype, len, slots, &mut ctx)
138 }
139
140 fn nbuffers(_array: ArrayView<'_, Self>) -> usize {
141 3
142 }
143
144 fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle {
145 match idx {
146 0 => BufferHandle::new_host(
147 array
148 .padded_symbols()
149 .slice(0..array.n_symbols())
150 .into_byte_buffer(),
151 ),
152 1 => BufferHandle::new_host(
153 array
154 .padded_symbol_lengths()
155 .slice(0..array.n_symbols())
156 .into_byte_buffer(),
157 ),
158 2 => array.codes_bytes_handle().clone(),
159 _ => vortex_panic!("FSSTArray buffer index {idx} out of bounds"),
160 }
161 }
162
163 fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option<String> {
164 match idx {
165 0 => Some("symbols".to_string()),
166 1 => Some("symbol_lengths".to_string()),
167 2 => Some("compressed_codes".to_string()),
168 _ => vortex_panic!("FSSTArray buffer_name index {idx} out of bounds"),
169 }
170 }
171
172 fn with_buffers(
173 &self,
174 array: ArrayView<'_, Self>,
175 buffers: &[BufferHandle],
176 ) -> VortexResult<ArrayParts<Self>> {
177 vortex_ensure!(
178 buffers.len() == 3,
179 "Expected 3 buffers, got {}",
180 buffers.len()
181 );
182 let symbols = Buffer::<Symbol>::from_byte_buffer(buffers[0].clone().try_to_host_sync()?);
183 let symbol_lengths = Buffer::<u8>::from_byte_buffer(buffers[1].clone().try_to_host_sync()?);
184 let data = FSSTData::try_new(symbols, symbol_lengths, buffers[2].clone(), array.len())?;
185 Ok(
186 ArrayParts::new(self.clone(), array.dtype().clone(), array.len(), data)
187 .with_slots(array.slots().iter().cloned().collect()),
188 )
189 }
190
191 fn serialize(
192 array: ArrayView<'_, Self>,
193 _session: &VortexSession,
194 ) -> VortexResult<Option<Vec<u8>>> {
195 let codes_offsets = array.codes_offsets();
196 Ok(Some(
197 FSSTMetadata {
198 uncompressed_lengths_ptype: array.uncompressed_lengths().dtype().as_ptype().into(),
199 codes_offsets_ptype: codes_offsets.dtype().as_ptype().into(),
200 }
201 .encode_to_vec(),
202 ))
203 }
204
205 fn deserialize(
230 &self,
231 dtype: &DType,
232 len: usize,
233 metadata: &[u8],
234 buffers: &[BufferHandle],
235 children: &dyn ArrayChildren,
236 session: &VortexSession,
237 ) -> VortexResult<ArrayParts<Self>> {
238 let metadata = FSSTMetadata::decode(metadata)?;
239 let symbols = Buffer::<Symbol>::from_byte_buffer(buffers[0].clone().try_to_host_sync()?);
240 let symbol_lengths = Buffer::<u8>::from_byte_buffer(buffers[1].clone().try_to_host_sync()?);
241
242 let mut ctx = session.create_execution_ctx();
243 if buffers.len() == 2 {
244 return Self::deserialize_legacy(
245 self,
246 dtype,
247 len,
248 &metadata,
249 &symbols,
250 &symbol_lengths,
251 children,
252 &mut ctx,
253 );
254 }
255
256 if buffers.len() == 3 {
257 let uncompressed_lengths = children.get(
258 0,
259 &DType::Primitive(
260 metadata.get_uncompressed_lengths_ptype()?,
261 Nullability::NonNullable,
262 ),
263 len,
264 )?;
265
266 let codes_bytes = buffers[2].clone();
267 let codes_offsets = children.get(
268 1,
269 &DType::Primitive(
270 PType::try_from(metadata.codes_offsets_ptype)?,
271 Nullability::NonNullable,
272 ),
273 len + 1,
275 )?;
276
277 let codes_validity = if children.len() == 2 {
278 Validity::from(dtype.nullability())
279 } else if children.len() == 3 {
280 let validity = children.get(2, &Validity::DTYPE, len)?;
281 Validity::Array(validity)
282 } else {
283 vortex_bail!("Expected 2 or 3 children, got {}", children.len());
284 };
285
286 FSSTData::validate_parts(
287 symbols.as_slice(),
288 symbol_lengths.as_slice(),
289 &codes_bytes,
290 &codes_offsets,
291 dtype.nullability(),
292 &uncompressed_lengths,
293 dtype,
294 len,
295 &mut ctx,
296 )?;
297 let slots = FSSTSlots {
298 uncompressed_lengths,
299 codes_offsets,
300 codes_validity: validity_to_child(&codes_validity, len),
301 }
302 .into_slots();
303 let data = FSSTData::try_new(symbols, symbol_lengths, codes_bytes, len)?;
304 return Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots));
305 }
306
307 vortex_bail!(
308 "InvalidArgument: Expected 2 or 3 buffers, got {}",
309 buffers.len()
310 );
311 }
312
313 fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String {
314 FSSTSlots::NAMES[idx].to_string()
315 }
316
317 fn execute(array: Array<Self>, ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
318 canonicalize_fsst(array.as_view(), ctx).map(ExecutionResult::done)
319 }
320
321 fn append_to_builder(
322 array: ArrayView<'_, Self>,
323 builder: &mut dyn ArrayBuilder,
324 ctx: &mut ExecutionCtx,
325 ) -> VortexResult<()> {
326 if let Some(result) =
327 match_each_varbin_builder!(builder, |builder| append_to_varbin(array, builder, ctx))
328 {
329 return result;
330 }
331
332 let Some(builder) = builder.as_any_mut().downcast_mut::<VarBinViewBuilder>() else {
337 vortex_bail!("append_to_builder for FSST requires a variable-binary builder")
338 };
339
340 let validity = array
343 .array()
344 .validity()?
345 .execute_mask(array.array().len(), ctx)?;
346 let (uncompressed_bytes, uncompressed_lens) = fsst_decode_bytes(array, ctx)?;
347 match_each_integer_ptype!(uncompressed_lens.ptype(), |P| {
348 builder.append_buffer_with_lengths(
349 uncompressed_bytes.freeze(),
350 uncompressed_lens.as_slice::<P>(),
351 &validity,
352 )
353 });
354 Ok(())
355 }
356
357 fn reduce_parent(
358 array: ArrayView<'_, Self>,
359 parent: &ArrayRef,
360 child_idx: usize,
361 ) -> VortexResult<Option<ArrayRef>> {
362 RULES.evaluate(array, parent, child_idx)
363 }
364}
365
366fn append_to_varbin<O: OffsetBuilderPType>(
371 array: ArrayView<'_, FSST>,
372 builder: &mut VarBinBuilder<O>,
373 ctx: &mut ExecutionCtx,
374) -> VortexResult<()>
375where
376 usize: AsPrimitive<O>,
377{
378 let plan = FsstDecodePlan::new(array, ctx)?;
379 let validity = array
380 .array()
381 .validity()?
382 .execute_mask(array.array().len(), ctx)?;
383 let decompressor = array.decompressor();
384 let mut decode = |out: &mut [MaybeUninit<u8>]| plan.decode_into(&decompressor, out);
387 match_each_integer_ptype!(plan.lengths.ptype(), |P| {
388 unsafe {
390 builder.append_decoded(
391 plan.total_size,
392 FSST_DECODE_SLACK,
393 plan.lengths.as_slice::<P>(),
394 &validity,
395 &mut decode,
396 )
397 }
398 })
399}
400
401#[array_slots(FSST)]
402pub struct FSSTSlots {
403 #[slot(0)]
405 pub uncompressed_lengths: ArrayRef,
406 #[slot(1)]
408 pub codes_offsets: ArrayRef,
409 #[slot(2)]
411 pub codes_validity: Option<ArrayRef>,
412}
413
414#[derive(Clone)]
423pub struct FSSTData {
424 symbol_table: Arc<FSSTSymbolTable>,
425 codes_bytes: BufferHandle,
427 len: usize,
429}
430
431impl Display for FSSTData {
432 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
433 write!(
434 f,
435 "len: {}, nsymbols: {}",
436 self.len, self.symbol_table.n_symbols
437 )
438 }
439}
440
441impl Debug for FSSTData {
442 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
443 f.debug_struct("FSSTArray")
444 .field("symbols", &self.symbols())
445 .field("symbol_lengths", &self.symbol_lengths())
446 .field("codes_bytes_len", &self.codes_bytes.len())
447 .field("len", &self.len)
448 .field("uncompressed_lengths", &"<outer slot>")
449 .field("codes_offsets", &"<outer slot>")
450 .field("codes_validity", &"<outer slot>")
451 .finish()
452 }
453}
454
455pub struct FSSTSymbolTable {
456 padded_symbols: Buffer<Symbol>,
459 padded_symbol_lengths: Buffer<u8>,
462 n_symbols: usize,
464 compressor: OnceLock<Compressor>,
466}
467
468impl FSSTSymbolTable {
469 pub fn new(symbols: Buffer<Symbol>, symbol_lengths: Buffer<u8>, n_symbols: usize) -> Self {
476 Self {
477 padded_symbols: pad_symbol_table(symbols, Symbol::ZERO),
478 padded_symbol_lengths: pad_symbol_table(symbol_lengths, 0),
479 n_symbols: n_symbols.min(FSST_SYMBOL_TABLE_LEN),
480 compressor: OnceLock::new(),
481 }
482 }
483
484 pub fn new_padded(
491 padded_symbols: Buffer<Symbol>,
492 padded_symbol_lengths: Buffer<u8>,
493 n_symbols: usize,
494 ) -> VortexResult<Self> {
495 vortex_ensure!(
496 padded_symbols.len() == FSST_SYMBOL_TABLE_LEN
497 && padded_symbol_lengths.len() == FSST_SYMBOL_TABLE_LEN,
498 InvalidArgument: "padded symbol table must have exactly {FSST_SYMBOL_TABLE_LEN} entries, found {} symbols and {} symbol lengths",
499 padded_symbols.len(),
500 padded_symbol_lengths.len()
501 );
502 vortex_ensure!(
503 n_symbols <= FSST_SYMBOL_TABLE_LEN,
504 InvalidArgument: "n_symbols must be <= {FSST_SYMBOL_TABLE_LEN}, found {n_symbols}"
505 );
506 Ok(Self {
507 padded_symbols,
508 padded_symbol_lengths,
509 n_symbols,
510 compressor: OnceLock::new(),
511 })
512 }
513
514 fn symbols(&self) -> &[Symbol] {
516 &self.padded_symbols.as_slice()[..self.n_symbols]
517 }
518
519 fn symbol_lengths(&self) -> &[u8] {
521 &self.padded_symbol_lengths.as_slice()[..self.n_symbols]
522 }
523
524 fn padded_symbols(&self) -> &Buffer<Symbol> {
527 &self.padded_symbols
528 }
529
530 fn padded_symbol_lengths(&self) -> &Buffer<u8> {
532 &self.padded_symbol_lengths
533 }
534
535 fn decompressor(&self) -> Decompressor<'_> {
540 const PADDED: &str = "FSST symbol table is padded to FSST_SYMBOL_TABLE_LEN entries";
541 let symbols = self
542 .padded_symbols
543 .as_slice()
544 .first_chunk::<FSST_SYMBOL_TABLE_LEN>()
545 .vortex_expect(PADDED);
546 let symbol_lengths = self
547 .padded_symbol_lengths
548 .as_slice()
549 .first_chunk::<FSST_SYMBOL_TABLE_LEN>()
550 .vortex_expect(PADDED);
551 Decompressor::new(symbols, symbol_lengths)
552 }
553
554 fn compressor(&self) -> &Compressor {
555 self.compressor
556 .get_or_init(|| Compressor::rebuild_from(self.symbols(), self.symbol_lengths()))
557 }
558}
559
560fn pad_symbol_table<T: Copy>(buffer: Buffer<T>, pad: T) -> Buffer<T> {
563 if buffer.len() == FSST_SYMBOL_TABLE_LEN {
564 return buffer;
565 }
566 padded_symbol_table(buffer.as_slice(), pad)
567}
568
569pub(crate) fn padded_symbol_table<T: Copy>(values: &[T], pad: T) -> Buffer<T> {
575 let populated = values.len().min(FSST_SYMBOL_TABLE_LEN);
576 let mut padded = BufferMut::with_capacity(FSST_SYMBOL_TABLE_LEN);
577 padded.extend_from_slice(&values[..populated]);
578 padded.push_n(pad, FSST_SYMBOL_TABLE_LEN - populated);
579 padded.freeze()
580}
581
582#[derive(Clone, Debug)]
583pub struct FSST;
584
585impl FSST {
586 pub fn try_new(
592 dtype: DType,
593 symbols: Buffer<Symbol>,
594 symbol_lengths: Buffer<u8>,
595 codes: VarBinArray,
596 uncompressed_lengths: ArrayRef,
597 ctx: &mut ExecutionCtx,
598 ) -> VortexResult<FSSTArray> {
599 let len = codes.len();
600 FSSTData::validate_parts_from_codes(
601 symbols.as_slice(),
602 symbol_lengths.as_slice(),
603 &codes,
604 &uncompressed_lengths,
605 &dtype,
606 len,
607 ctx,
608 )?;
609 let slots = FSSTData::make_slots(&codes, &uncompressed_lengths);
610 let codes_bytes = codes.bytes_handle().clone();
611 let data = FSSTData::try_new(symbols, symbol_lengths, codes_bytes, len)?;
612 Ok(unsafe {
613 Array::from_parts_unchecked(ArrayParts::new(FSST, dtype, len, data).with_slots(slots))
614 })
615 }
616
617 pub fn try_new_with_symbol_table(
618 dtype: DType,
619 symbol_table: Arc<FSSTSymbolTable>,
620 codes: VarBinArray,
621 uncompressed_lengths: ArrayRef,
622 ctx: &mut ExecutionCtx,
623 ) -> VortexResult<FSSTArray> {
624 let len = codes.len();
625 FSSTData::validate_parts_from_codes(
626 symbol_table.symbols(),
627 symbol_table.symbol_lengths(),
628 &codes,
629 &uncompressed_lengths,
630 &dtype,
631 len,
632 ctx,
633 )?;
634 let slots = FSSTData::make_slots(&codes, &uncompressed_lengths);
635 let codes_bytes = codes.bytes_handle().clone();
636 let data =
637 unsafe { FSSTData::new_unchecked_with_symbol_table(symbol_table, codes_bytes, len) };
638 Ok(unsafe {
639 Array::from_parts_unchecked(ArrayParts::new(FSST, dtype, len, data).with_slots(slots))
640 })
641 }
642
643 #[allow(clippy::too_many_arguments)]
647 fn deserialize_legacy(
648 &self,
649 dtype: &DType,
650 len: usize,
651 metadata: &FSSTMetadata,
652 symbols: &Buffer<Symbol>,
653 symbol_lengths: &Buffer<u8>,
654 children: &dyn ArrayChildren,
655 ctx: &mut ExecutionCtx,
656 ) -> VortexResult<ArrayParts<Self>> {
657 if children.len() != 2 {
658 vortex_bail!(InvalidArgument: "Expected 2 children, got {}", children.len());
659 }
660 let codes = children.get(0, &DType::Binary(dtype.nullability()), len)?;
661 let codes: VarBinArray = codes
662 .as_opt::<VarBin>()
663 .ok_or_else(|| {
664 vortex_err!(
665 "Expected VarBinArray for codes, got {}",
666 codes.encoding_id()
667 )
668 })?
669 .into_owned();
670 let uncompressed_lengths = children.get(
671 1,
672 &DType::Primitive(
673 metadata.get_uncompressed_lengths_ptype()?,
674 Nullability::NonNullable,
675 ),
676 len,
677 )?;
678
679 FSSTData::validate_parts_from_codes(
680 symbols.as_slice(),
681 symbol_lengths.as_slice(),
682 &codes,
683 &uncompressed_lengths,
684 dtype,
685 len,
686 ctx,
687 )?;
688 let slots = FSSTData::make_slots(&codes, &uncompressed_lengths);
689 let codes_bytes = codes.bytes_handle().clone();
690 let data = FSSTData::try_new(symbols.clone(), symbol_lengths.clone(), codes_bytes, len)?;
691 Ok(ArrayParts::new(self.clone(), dtype.clone(), len, data).with_slots(slots))
692 }
693
694 pub(crate) unsafe fn new_unchecked_with_symbol_table(
695 dtype: DType,
696 symbol_table: Arc<FSSTSymbolTable>,
697 codes: VarBinArray,
698 uncompressed_lengths: ArrayRef,
699 ) -> FSSTArray {
700 let len = codes.len();
701 let slots = FSSTData::make_slots(&codes, &uncompressed_lengths);
702 let codes_bytes = codes.bytes_handle().clone();
703 let data =
704 unsafe { FSSTData::new_unchecked_with_symbol_table(symbol_table, codes_bytes, len) };
705 unsafe {
706 Array::from_parts_unchecked(ArrayParts::new(FSST, dtype, len, data).with_slots(slots))
707 }
708 }
709}
710
711impl FSSTData {
712 fn make_slots(codes: &VarBinArray, uncompressed_lengths: &ArrayRef) -> ArraySlots {
713 FSSTSlots {
714 uncompressed_lengths: uncompressed_lengths.clone(),
715 codes_offsets: codes.offsets().clone(),
716 codes_validity: validity_to_child(
717 &codes
718 .validity()
719 .vortex_expect("FSST codes validity should be derivable"),
720 codes.len(),
721 ),
722 }
723 .into_slots()
724 }
725
726 pub fn try_new(
743 symbols: Buffer<Symbol>,
744 symbol_lengths: Buffer<u8>,
745 codes_bytes: BufferHandle,
746 len: usize,
747 ) -> VortexResult<Self> {
748 vortex_ensure!(
749 symbols.len() == symbol_lengths.len(),
750 InvalidArgument: "symbols and symbol_lengths arrays must have same length, found {} and {}",
751 symbols.len(),
752 symbol_lengths.len()
753 );
754 vortex_ensure!(
755 symbols.len() <= FSST_SYMBOL_TABLE_LEN,
756 InvalidArgument: "symbols array must have length <= {FSST_SYMBOL_TABLE_LEN}, found {}",
757 symbols.len()
758 );
759 let n_symbols = symbols.len();
760 let symbol_table = Arc::new(FSSTSymbolTable::new(symbols, symbol_lengths, n_symbols));
762 unsafe {
763 Ok(Self::new_unchecked_with_symbol_table(
764 symbol_table,
765 codes_bytes,
766 len,
767 ))
768 }
769 }
770
771 pub fn validate(
772 &self,
773 dtype: &DType,
774 len: usize,
775 slots: &[Option<ArrayRef>],
776 ctx: &mut ExecutionCtx,
777 ) -> VortexResult<()> {
778 let fsst_slots = FSSTSlotsView::from_slots(slots);
779 Self::validate_parts(
780 self.symbol_table.symbols(),
781 self.symbol_table.symbol_lengths(),
782 &self.codes_bytes,
783 fsst_slots.codes_offsets,
784 dtype.nullability(),
785 fsst_slots.uncompressed_lengths,
786 dtype,
787 len,
788 ctx,
789 )
790 }
791
792 #[expect(clippy::too_many_arguments)]
794 fn validate_parts(
795 symbols: &[Symbol],
796 symbol_lengths: &[u8],
797 codes_bytes: &BufferHandle,
798 codes_offsets: &ArrayRef,
799 codes_nullability: Nullability,
800 uncompressed_lengths: &ArrayRef,
801 dtype: &DType,
802 len: usize,
803 ctx: &mut ExecutionCtx,
804 ) -> VortexResult<()> {
805 vortex_ensure!(
806 matches!(dtype, DType::Binary(_) | DType::Utf8(_)),
807 "FSST arrays must be Binary or Utf8, found {dtype}"
808 );
809
810 if symbols.len() > FSST_SYMBOL_TABLE_LEN {
811 vortex_bail!(InvalidArgument: "symbols array must have length <= {FSST_SYMBOL_TABLE_LEN}");
812 }
813
814 if symbols.len() != symbol_lengths.len() {
815 vortex_bail!(InvalidArgument: "symbols and symbol_lengths arrays must have same length");
816 }
817
818 Self::validate_symbol_lengths(symbol_lengths)?;
819
820 let codes_len = codes_offsets.len().saturating_sub(1);
822 if codes_len != len {
823 vortex_bail!(InvalidArgument: "codes must have same len as outer array");
824 }
825
826 if uncompressed_lengths.len() != len {
827 vortex_bail!(InvalidArgument: "uncompressed_lengths must be same len as codes");
828 }
829
830 if !uncompressed_lengths.dtype().is_int() || uncompressed_lengths.dtype().is_nullable() {
831 vortex_bail!(InvalidArgument: "uncompressed_lengths must have integer type and cannot be nullable, found {}", uncompressed_lengths.dtype());
832 }
833
834 if !codes_offsets.dtype().is_int() || codes_offsets.dtype().is_nullable() {
836 vortex_bail!(InvalidArgument: "codes offsets must be non-nullable integer type, found {}", codes_offsets.dtype());
837 }
838
839 if codes_nullability != dtype.nullability() {
840 vortex_bail!(InvalidArgument: "codes nullability must match outer dtype nullability");
841 }
842
843 if codes_bytes.is_on_host() && codes_offsets.is_host() && !codes_offsets.is_empty() {
845 let last_offset: usize = (&codes_offsets
846 .execute_scalar(codes_offsets.len() - 1, ctx)
847 .vortex_expect("offsets must support scalar_at"))
848 .try_into()
849 .vortex_expect("Failed to convert offset to usize");
850 vortex_ensure!(
851 last_offset <= codes_bytes.len(),
852 InvalidArgument: "Last codes offset {} exceeds codes bytes length {}",
853 last_offset,
854 codes_bytes.len()
855 );
856 }
857
858 Ok(())
859 }
860
861 fn validate_symbol_lengths(symbol_lengths: &[u8]) -> VortexResult<()> {
862 let mut expected = 2;
863 for (idx, &len) in symbol_lengths.iter().enumerate() {
864 if len > 8 || len == 0 {
865 vortex_bail!(InvalidArgument: "symbol length at index {idx} must be between 1 and 8, found {len}");
866 }
867
868 if expected == 1 {
869 if len != 1 {
870 vortex_bail!(InvalidArgument: "symbol length at index {idx} must be 1 after one-byte symbols begin, found {len}");
871 }
872 } else {
873 if len == 1 {
874 expected = 1;
875 }
876
877 if len < expected {
878 vortex_bail!(InvalidArgument: "symbol length at index {idx} violates FSST symbol table ordering");
879 }
880 expected = len;
881 }
882 }
883
884 Ok(())
885 }
886
887 fn validate_parts_from_codes(
889 symbols: &[Symbol],
890 symbol_lengths: &[u8],
891 codes: &VarBinArray,
892 uncompressed_lengths: &ArrayRef,
893 dtype: &DType,
894 len: usize,
895 ctx: &mut ExecutionCtx,
896 ) -> VortexResult<()> {
897 Self::validate_parts(
898 symbols,
899 symbol_lengths,
900 codes.bytes_handle(),
901 codes.offsets(),
902 codes.dtype().nullability(),
903 uncompressed_lengths,
904 dtype,
905 len,
906 ctx,
907 )
908 }
909
910 pub(crate) unsafe fn new_unchecked_with_symbol_table(
911 symbol_table: Arc<FSSTSymbolTable>,
912 codes_bytes: BufferHandle,
913 len: usize,
914 ) -> Self {
915 Self {
916 symbol_table,
917 codes_bytes,
918 len,
919 }
920 }
921
922 pub fn len(&self) -> usize {
924 self.len
925 }
926
927 pub fn is_empty(&self) -> bool {
929 self.len == 0
930 }
931
932 pub fn symbols(&self) -> &[Symbol] {
938 &self.symbol_table.padded_symbols().as_slice()[0..self.symbol_table.n_symbols]
939 }
940
941 pub fn symbol_lengths(&self) -> &[u8] {
945 &self.symbol_table.padded_symbol_lengths().as_slice()[0..self.symbol_table.n_symbols]
946 }
947
948 pub fn padded_symbols(&self) -> &Buffer<Symbol> {
951 self.symbol_table.padded_symbols()
952 }
953
954 pub fn padded_symbol_lengths(&self) -> &Buffer<u8> {
957 self.symbol_table.padded_symbol_lengths()
958 }
959
960 pub fn n_symbols(&self) -> usize {
962 self.symbol_table.n_symbols
963 }
964
965 pub fn symbol_table(&self) -> Arc<FSSTSymbolTable> {
966 Arc::clone(&self.symbol_table)
967 }
968
969 pub fn codes_bytes_handle(&self) -> &BufferHandle {
971 &self.codes_bytes
972 }
973
974 pub fn codes_bytes(&self) -> &ByteBuffer {
976 self.codes_bytes.as_host()
977 }
978
979 pub fn decompressor(&self) -> Decompressor<'_> {
982 self.symbol_table.decompressor()
983 }
984
985 pub fn compressor(&self) -> &Compressor {
987 self.symbol_table.compressor()
988 }
989}
990
991pub trait FSSTArrayExt: FSSTArraySlotsExt {
992 fn uncompressed_lengths_dtype(&self) -> &DType {
993 self.uncompressed_lengths().dtype()
994 }
995
996 fn codes(&self) -> VarBinArray {
999 let offsets = self.codes_offsets().clone();
1000 let validity =
1001 child_to_validity(self.codes_validity(), self.as_ref().dtype().nullability());
1002 let codes_bytes = self.codes_bytes_handle().clone();
1003 unsafe {
1005 VarBinArray::new_unchecked_from_handle(
1006 offsets,
1007 codes_bytes,
1008 DType::Binary(self.as_ref().dtype().nullability()),
1009 validity,
1010 )
1011 }
1012 }
1013
1014 fn codes_dtype(&self) -> DType {
1016 DType::Binary(self.as_ref().dtype().nullability())
1017 }
1018}
1019
1020impl<T: TypedArrayRef<FSST>> FSSTArrayExt for T {}
1021
1022impl ValidityVTable<FSST> for FSST {
1023 fn validity(array: ArrayView<'_, FSST>) -> VortexResult<Validity> {
1024 Ok(child_to_validity(
1025 array.codes_validity(),
1026 array.dtype().nullability(),
1027 ))
1028 }
1029}
1030
1031#[cfg(test)]
1032mod test {
1033 use fsst::Compressor;
1034 use fsst::Symbol;
1035 use prost::Message;
1036 use vortex_array::ArrayPlugin;
1037 use vortex_array::IntoArray;
1038 use vortex_array::VortexSessionExecute;
1039 use vortex_array::array_session;
1040 use vortex_array::arrays::VarBinViewArray;
1041 use vortex_array::buffer::BufferHandle;
1042 use vortex_array::dtype::DType;
1043 use vortex_array::dtype::Nullability;
1044 use vortex_array::dtype::PType;
1045 use vortex_array::test_harness::check_metadata;
1046 use vortex_array::vtable::VTable as _;
1047 use vortex_buffer::Buffer;
1048 use vortex_error::VortexResult;
1049 use vortex_error::vortex_err;
1050
1051 use crate::FSST;
1052 use crate::array::FSST_SYMBOL_TABLE_LEN;
1053 use crate::array::FSSTArrayExt;
1054 use crate::array::FSSTArraySlotsExt;
1055 use crate::array::FSSTData;
1056 use crate::array::FSSTMetadata;
1057 use crate::array::FSSTSymbolTable;
1058 use crate::array::padded_symbol_table;
1059 use crate::fsst_compress;
1060 use crate::fsst_train_compressor;
1061
1062 #[test]
1063 fn slice_reuses_initialized_compressor() -> VortexResult<()> {
1064 let symbols = Buffer::<Symbol>::copy_from([
1065 Symbol::from_slice(b"abc00000"),
1066 Symbol::from_slice(b"defghijk"),
1067 ]);
1068 let symbol_lengths = Buffer::<u8>::copy_from([3, 8]);
1069
1070 let compressor = Compressor::rebuild_from(symbols.as_slice(), symbol_lengths.as_slice());
1071 let mut ctx = array_session().create_execution_ctx();
1072 let strings = VarBinViewArray::from_iter_str(["abcabcab", "defghijk", "abcxyz"]);
1073 let fsst_array = fsst_compress(&strings.into_array(), &compressor, &mut ctx)?;
1074
1075 let compressor_ptr = fsst_array.compressor() as *const Compressor;
1076 let sliced = fsst_array
1077 .slice(1..3)?
1078 .try_downcast::<FSST>()
1079 .map_err(|_| vortex_err!("slice must return an FSST array"))?;
1080 let sliced_compressor_ptr = sliced.compressor() as *const Compressor;
1081
1082 assert_eq!(compressor_ptr, sliced_compressor_ptr);
1083 Ok(())
1084 }
1085
1086 #[test]
1090 fn symbol_table_padded_on_creation() -> VortexResult<()> {
1091 let mut ctx = array_session().create_execution_ctx();
1092 let strings = VarBinViewArray::from_iter_str(["abcabcab", "defghijk", "abcxyz"]);
1093 let compressor = Compressor::rebuild_from(
1094 [
1095 Symbol::from_slice(b"abc00000"),
1096 Symbol::from_slice(b"defghijk"),
1097 ],
1098 [3u8, 8],
1099 );
1100 let fsst_array = fsst_compress(&strings.into_array(), &compressor, &mut ctx)?;
1101
1102 assert_eq!(fsst_array.padded_symbols().len(), FSST_SYMBOL_TABLE_LEN);
1103 assert_eq!(
1104 fsst_array.padded_symbol_lengths().len(),
1105 FSST_SYMBOL_TABLE_LEN
1106 );
1107 assert_eq!(fsst_array.padded_symbol_lengths().as_slice()[2..], [0; 253]);
1108
1109 assert_eq!(fsst_array.n_symbols(), 2);
1111 assert_eq!(fsst_array.symbols().len(), 2);
1112 assert_eq!(fsst_array.symbol_lengths(), &[3, 8]);
1113 assert_eq!(
1114 FSST::buffer(fsst_array.as_view(), 0).len(),
1115 2 * size_of::<Symbol>()
1116 );
1117 assert_eq!(FSST::buffer(fsst_array.as_view(), 1).len(), 2);
1118
1119 let decompressed = fsst_array
1120 .into_array()
1121 .execute::<VarBinViewArray>(&mut ctx)?;
1122 assert_eq!(decompressed.bytes_at(0).as_slice(), b"abcabcab".as_ref());
1123 assert_eq!(decompressed.bytes_at(1).as_slice(), b"defghijk".as_ref());
1124 assert_eq!(decompressed.bytes_at(2).as_slice(), b"abcxyz".as_ref());
1125 Ok(())
1126 }
1127
1128 #[test]
1131 fn symbol_table_padded_on_deserialize() -> VortexResult<()> {
1132 let mut ctx = array_session().create_execution_ctx();
1133 let input = VarBinViewArray::from_iter_str(["abcabcab", "defghijk"]).into_array();
1134 let compressor = fsst_train_compressor(&input, &mut ctx)?;
1135 let fsst_array = fsst_compress(&input, &compressor, &mut ctx)?;
1136
1137 let buffers = [
1138 BufferHandle::new_host(
1139 fsst_array
1140 .padded_symbols()
1141 .slice(0..fsst_array.n_symbols())
1142 .into_byte_buffer(),
1143 ),
1144 BufferHandle::new_host(
1145 fsst_array
1146 .padded_symbol_lengths()
1147 .slice(0..fsst_array.n_symbols())
1148 .into_byte_buffer(),
1149 ),
1150 fsst_array.codes_bytes_handle().clone(),
1151 ];
1152 assert!(buffers[1].len() < FSST_SYMBOL_TABLE_LEN);
1153
1154 let children = vec![
1155 fsst_array.uncompressed_lengths().clone(),
1156 fsst_array.codes_offsets().clone(),
1157 ];
1158
1159 let deserialized = ArrayPlugin::deserialize(
1160 &FSST,
1161 &DType::Utf8(Nullability::NonNullable),
1162 2,
1163 &FSSTMetadata {
1164 uncompressed_lengths_ptype: fsst_array
1165 .uncompressed_lengths()
1166 .dtype()
1167 .as_ptype()
1168 .into(),
1169 codes_offsets_ptype: fsst_array.codes_offsets().dtype().as_ptype().into(),
1170 }
1171 .encode_to_vec(),
1172 &buffers,
1173 &children.as_slice(),
1174 &array_session(),
1175 )?;
1176
1177 let padded = deserialized
1178 .clone()
1179 .try_downcast::<FSST>()
1180 .map_err(|_| vortex_err!("deserialize must return an FSST array"))?;
1181 assert_eq!(padded.padded_symbols().len(), FSST_SYMBOL_TABLE_LEN);
1182 assert_eq!(padded.n_symbols(), fsst_array.symbols().len());
1183
1184 let decompressed = deserialized.execute::<VarBinViewArray>(&mut ctx)?;
1185 assert_eq!(decompressed.bytes_at(0).as_slice(), b"abcabcab".as_ref());
1186 assert_eq!(decompressed.bytes_at(1).as_slice(), b"defghijk".as_ref());
1187 Ok(())
1188 }
1189
1190 #[test]
1192 fn padded_constructor_does_not_repad() -> VortexResult<()> {
1193 let symbols = padded_symbol_table(&[Symbol::from_slice(b"ab000000")], Symbol::ZERO);
1194 let symbol_lengths = padded_symbol_table(&[2u8], 0);
1195 let symbols_ptr = symbols.as_slice().as_ptr();
1196 let symbol_lengths_ptr = symbol_lengths.as_slice().as_ptr();
1197
1198 let data = FSSTSymbolTable::new_padded(symbols, symbol_lengths, 1)?;
1199
1200 assert_eq!(data.padded_symbols().as_slice().as_ptr(), symbols_ptr);
1201 assert_eq!(
1202 data.padded_symbol_lengths().as_slice().as_ptr(),
1203 symbol_lengths_ptr
1204 );
1205 assert_eq!(data.symbols().len(), 1);
1206 Ok(())
1207 }
1208
1209 #[test]
1212 fn rejects_unpadded_input_to_padded_constructor() {
1213 assert!(
1214 FSSTSymbolTable::new_padded(
1215 Buffer::<Symbol>::copy_from([Symbol::from_slice(b"ab000000")]),
1216 Buffer::<u8>::copy_from([2]),
1217 1,
1218 )
1219 .is_err()
1220 );
1221 assert!(
1222 FSSTSymbolTable::new_padded(
1223 Buffer::<Symbol>::full(Symbol::ZERO, FSST_SYMBOL_TABLE_LEN),
1224 Buffer::<u8>::full(0, FSST_SYMBOL_TABLE_LEN),
1225 FSST_SYMBOL_TABLE_LEN + 1,
1226 )
1227 .is_err()
1228 );
1229 }
1230
1231 #[test]
1232 fn rejects_malformed_symbol_table() {
1233 let codes_bytes = BufferHandle::new_host(Buffer::<u8>::empty());
1234 assert!(
1235 FSSTData::try_new(
1236 Buffer::<Symbol>::copy_from([Symbol::from_slice(b"ab000000")]),
1237 Buffer::<u8>::copy_from([2, 2]),
1238 codes_bytes.clone(),
1239 0,
1240 )
1241 .is_err()
1242 );
1243 assert!(
1244 FSSTData::try_new(
1245 Buffer::<Symbol>::full(Symbol::from_slice(b"ab000000"), FSST_SYMBOL_TABLE_LEN + 1,),
1246 Buffer::<u8>::full(2, FSST_SYMBOL_TABLE_LEN + 1),
1247 codes_bytes,
1248 0,
1249 )
1250 .is_err()
1251 );
1252 }
1253
1254 #[cfg_attr(miri, ignore)]
1255 #[test]
1256 fn test_fsst_metadata() {
1257 check_metadata(
1258 "fsst.metadata",
1259 &FSSTMetadata {
1260 uncompressed_lengths_ptype: PType::U64 as i32,
1261 codes_offsets_ptype: PType::I32 as i32,
1262 }
1263 .encode_to_vec(),
1264 );
1265 }
1266
1267 #[test]
1275 fn test_back_compat() -> VortexResult<()> {
1276 let symbols = Buffer::<Symbol>::copy_from([
1277 Symbol::from_slice(b"abc00000"),
1278 Symbol::from_slice(b"defghijk"),
1279 ]);
1280 let symbol_lengths = Buffer::<u8>::copy_from([3, 8]);
1281
1282 let compressor = Compressor::rebuild_from(symbols.as_slice(), symbol_lengths.as_slice());
1283 let mut ctx = array_session().create_execution_ctx();
1284 let input = VarBinViewArray::from_iter_str(["abcabcab", "defghijk"]);
1285 let fsst_array = fsst_compress(&input.into_array(), &compressor, &mut ctx)?;
1286
1287 let compressed_codes = fsst_array.codes();
1288
1289 let buffers = [
1293 BufferHandle::new_host(symbols.into_byte_buffer()),
1294 BufferHandle::new_host(symbol_lengths.into_byte_buffer()),
1295 ];
1296
1297 let children = vec![
1301 compressed_codes.into_array(),
1302 fsst_array.uncompressed_lengths().clone(),
1303 ];
1304
1305 let fsst = ArrayPlugin::deserialize(
1306 &FSST,
1307 &DType::Utf8(Nullability::NonNullable),
1308 2,
1309 &FSSTMetadata {
1310 uncompressed_lengths_ptype: fsst_array
1311 .uncompressed_lengths()
1312 .dtype()
1313 .as_ptype()
1314 .into(),
1315 codes_offsets_ptype: 0,
1317 }
1318 .encode_to_vec(),
1319 &buffers,
1320 &children.as_slice(),
1321 &array_session(),
1322 )?;
1323
1324 let decompressed =
1325 fsst.execute::<VarBinViewArray>(&mut array_session().create_execution_ctx())?;
1326 let mask = decompressed
1327 .validity()?
1328 .execute_mask(decompressed.len(), &mut ctx)?;
1329 assert!(mask.value(0));
1330 assert_eq!(decompressed.bytes_at(0).as_slice(), b"abcabcab".as_ref());
1331 assert!(mask.value(1));
1332 assert_eq!(decompressed.bytes_at(1).as_slice(), b"defghijk".as_ref());
1333 Ok(())
1334 }
1335}