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