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