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::ArrayDeserialization;
1037 use vortex_array::ArrayPlugin;
1038 use vortex_array::IntoArray;
1039 use vortex_array::VortexSessionExecute;
1040 use vortex_array::array_session;
1041 use vortex_array::arrays::VarBinViewArray;
1042 use vortex_array::buffer::BufferHandle;
1043 use vortex_array::dtype::DType;
1044 use vortex_array::dtype::Nullability;
1045 use vortex_array::dtype::PType;
1046 use vortex_array::test_harness::check_metadata;
1047 use vortex_array::vtable::VTable as _;
1048 use vortex_buffer::Buffer;
1049 use vortex_error::VortexResult;
1050 use vortex_error::vortex_err;
1051
1052 use crate::FSST;
1053 use crate::array::FSST_SYMBOL_TABLE_LEN;
1054 use crate::array::FSSTArrayExt;
1055 use crate::array::FSSTArraySlotsExt;
1056 use crate::array::FSSTData;
1057 use crate::array::FSSTMetadata;
1058 use crate::array::FSSTSymbolTable;
1059 use crate::array::padded_symbol_table;
1060 use crate::fsst_compress;
1061 use crate::fsst_train_compressor;
1062
1063 #[test]
1064 fn slice_reuses_initialized_compressor() -> VortexResult<()> {
1065 let symbols = Buffer::<Symbol>::copy_from([
1066 Symbol::from_slice(b"abc00000"),
1067 Symbol::from_slice(b"defghijk"),
1068 ]);
1069 let symbol_lengths = Buffer::<u8>::copy_from([3, 8]);
1070
1071 let compressor = Compressor::rebuild_from(symbols.as_slice(), symbol_lengths.as_slice());
1072 let mut ctx = array_session().create_execution_ctx();
1073 let strings = VarBinViewArray::from_iter_str(["abcabcab", "defghijk", "abcxyz"]);
1074 let fsst_array = fsst_compress(&strings.into_array(), &compressor, &mut ctx)?;
1075
1076 let compressor_ptr = fsst_array.compressor() as *const Compressor;
1077 let sliced = fsst_array
1078 .slice(1..3)?
1079 .try_downcast::<FSST>()
1080 .map_err(|_| vortex_err!("slice must return an FSST array"))?;
1081 let sliced_compressor_ptr = sliced.compressor() as *const Compressor;
1082
1083 assert_eq!(compressor_ptr, sliced_compressor_ptr);
1084 Ok(())
1085 }
1086
1087 #[test]
1091 fn symbol_table_padded_on_creation() -> VortexResult<()> {
1092 let mut ctx = array_session().create_execution_ctx();
1093 let strings = VarBinViewArray::from_iter_str(["abcabcab", "defghijk", "abcxyz"]);
1094 let compressor = Compressor::rebuild_from(
1095 [
1096 Symbol::from_slice(b"abc00000"),
1097 Symbol::from_slice(b"defghijk"),
1098 ],
1099 [3u8, 8],
1100 );
1101 let fsst_array = fsst_compress(&strings.into_array(), &compressor, &mut ctx)?;
1102
1103 assert_eq!(fsst_array.padded_symbols().len(), FSST_SYMBOL_TABLE_LEN);
1104 assert_eq!(
1105 fsst_array.padded_symbol_lengths().len(),
1106 FSST_SYMBOL_TABLE_LEN
1107 );
1108 assert_eq!(fsst_array.padded_symbol_lengths().as_slice()[2..], [0; 253]);
1109
1110 assert_eq!(fsst_array.n_symbols(), 2);
1112 assert_eq!(fsst_array.symbols().len(), 2);
1113 assert_eq!(fsst_array.symbol_lengths(), &[3, 8]);
1114 assert_eq!(
1115 FSST::buffer(fsst_array.as_view(), 0).len(),
1116 2 * size_of::<Symbol>()
1117 );
1118 assert_eq!(FSST::buffer(fsst_array.as_view(), 1).len(), 2);
1119
1120 let decompressed = fsst_array
1121 .into_array()
1122 .execute::<VarBinViewArray>(&mut ctx)?;
1123 assert_eq!(decompressed.bytes_at(0).as_slice(), b"abcabcab".as_ref());
1124 assert_eq!(decompressed.bytes_at(1).as_slice(), b"defghijk".as_ref());
1125 assert_eq!(decompressed.bytes_at(2).as_slice(), b"abcxyz".as_ref());
1126 Ok(())
1127 }
1128
1129 #[test]
1132 fn symbol_table_padded_on_deserialize() -> VortexResult<()> {
1133 let mut ctx = array_session().create_execution_ctx();
1134 let input = VarBinViewArray::from_iter_str(["abcabcab", "defghijk"]).into_array();
1135 let compressor = fsst_train_compressor(&input, &mut ctx)?;
1136 let fsst_array = fsst_compress(&input, &compressor, &mut ctx)?;
1137
1138 let buffers = [
1139 BufferHandle::new_host(
1140 fsst_array
1141 .padded_symbols()
1142 .slice(0..fsst_array.n_symbols())
1143 .into_byte_buffer(),
1144 ),
1145 BufferHandle::new_host(
1146 fsst_array
1147 .padded_symbol_lengths()
1148 .slice(0..fsst_array.n_symbols())
1149 .into_byte_buffer(),
1150 ),
1151 fsst_array.codes_bytes_handle().clone(),
1152 ];
1153 assert!(buffers[1].len() < FSST_SYMBOL_TABLE_LEN);
1154
1155 let children = vec![
1156 fsst_array.uncompressed_lengths().clone(),
1157 fsst_array.codes_offsets().clone(),
1158 ];
1159
1160 let deserialized = ArrayPlugin::deserialize(
1161 &FSST,
1162 ArrayDeserialization::new(
1163 vortex_array::ArrayVTable::id(&FSST),
1164 &DType::Utf8(Nullability::NonNullable),
1165 2,
1166 &FSSTMetadata {
1167 uncompressed_lengths_ptype: fsst_array
1168 .uncompressed_lengths()
1169 .dtype()
1170 .as_ptype()
1171 .into(),
1172 codes_offsets_ptype: fsst_array.codes_offsets().dtype().as_ptype().into(),
1173 }
1174 .encode_to_vec(),
1175 &buffers,
1176 &children.as_slice(),
1177 ),
1178 &array_session(),
1179 )?;
1180
1181 let padded = deserialized
1182 .clone()
1183 .try_downcast::<FSST>()
1184 .map_err(|_| vortex_err!("deserialize must return an FSST array"))?;
1185 assert_eq!(padded.padded_symbols().len(), FSST_SYMBOL_TABLE_LEN);
1186 assert_eq!(padded.n_symbols(), fsst_array.symbols().len());
1187
1188 let decompressed = deserialized.execute::<VarBinViewArray>(&mut ctx)?;
1189 assert_eq!(decompressed.bytes_at(0).as_slice(), b"abcabcab".as_ref());
1190 assert_eq!(decompressed.bytes_at(1).as_slice(), b"defghijk".as_ref());
1191 Ok(())
1192 }
1193
1194 #[test]
1196 fn padded_constructor_does_not_repad() -> VortexResult<()> {
1197 let symbols = padded_symbol_table(&[Symbol::from_slice(b"ab000000")], Symbol::ZERO);
1198 let symbol_lengths = padded_symbol_table(&[2u8], 0);
1199 let symbols_ptr = symbols.as_slice().as_ptr();
1200 let symbol_lengths_ptr = symbol_lengths.as_slice().as_ptr();
1201
1202 let data = FSSTSymbolTable::new_padded(symbols, symbol_lengths, 1)?;
1203
1204 assert_eq!(data.padded_symbols().as_slice().as_ptr(), symbols_ptr);
1205 assert_eq!(
1206 data.padded_symbol_lengths().as_slice().as_ptr(),
1207 symbol_lengths_ptr
1208 );
1209 assert_eq!(data.symbols().len(), 1);
1210 Ok(())
1211 }
1212
1213 #[test]
1216 fn rejects_unpadded_input_to_padded_constructor() {
1217 assert!(
1218 FSSTSymbolTable::new_padded(
1219 Buffer::<Symbol>::copy_from([Symbol::from_slice(b"ab000000")]),
1220 Buffer::<u8>::copy_from([2]),
1221 1,
1222 )
1223 .is_err()
1224 );
1225 assert!(
1226 FSSTSymbolTable::new_padded(
1227 Buffer::<Symbol>::full(Symbol::ZERO, FSST_SYMBOL_TABLE_LEN),
1228 Buffer::<u8>::full(0, FSST_SYMBOL_TABLE_LEN),
1229 FSST_SYMBOL_TABLE_LEN + 1,
1230 )
1231 .is_err()
1232 );
1233 }
1234
1235 #[test]
1236 fn rejects_malformed_symbol_table() {
1237 let codes_bytes = BufferHandle::new_host(Buffer::<u8>::empty());
1238 assert!(
1239 FSSTData::try_new(
1240 Buffer::<Symbol>::copy_from([Symbol::from_slice(b"ab000000")]),
1241 Buffer::<u8>::copy_from([2, 2]),
1242 codes_bytes.clone(),
1243 0,
1244 )
1245 .is_err()
1246 );
1247 assert!(
1248 FSSTData::try_new(
1249 Buffer::<Symbol>::full(Symbol::from_slice(b"ab000000"), FSST_SYMBOL_TABLE_LEN + 1,),
1250 Buffer::<u8>::full(2, FSST_SYMBOL_TABLE_LEN + 1),
1251 codes_bytes,
1252 0,
1253 )
1254 .is_err()
1255 );
1256 }
1257
1258 #[cfg_attr(miri, ignore)]
1259 #[test]
1260 fn test_fsst_metadata() {
1261 check_metadata(
1262 "fsst.metadata",
1263 &FSSTMetadata {
1264 uncompressed_lengths_ptype: PType::U64 as i32,
1265 codes_offsets_ptype: PType::I32 as i32,
1266 }
1267 .encode_to_vec(),
1268 );
1269 }
1270
1271 #[test]
1279 fn test_back_compat() -> VortexResult<()> {
1280 let symbols = Buffer::<Symbol>::copy_from([
1281 Symbol::from_slice(b"abc00000"),
1282 Symbol::from_slice(b"defghijk"),
1283 ]);
1284 let symbol_lengths = Buffer::<u8>::copy_from([3, 8]);
1285
1286 let compressor = Compressor::rebuild_from(symbols.as_slice(), symbol_lengths.as_slice());
1287 let mut ctx = array_session().create_execution_ctx();
1288 let input = VarBinViewArray::from_iter_str(["abcabcab", "defghijk"]);
1289 let fsst_array = fsst_compress(&input.into_array(), &compressor, &mut ctx)?;
1290
1291 let compressed_codes = fsst_array.codes();
1292
1293 let buffers = [
1297 BufferHandle::new_host(symbols.into_byte_buffer()),
1298 BufferHandle::new_host(symbol_lengths.into_byte_buffer()),
1299 ];
1300
1301 let children = vec![
1305 compressed_codes.into_array(),
1306 fsst_array.uncompressed_lengths().clone(),
1307 ];
1308
1309 let fsst = ArrayPlugin::deserialize(
1310 &FSST,
1311 ArrayDeserialization::new(
1312 vortex_array::ArrayVTable::id(&FSST),
1313 &DType::Utf8(Nullability::NonNullable),
1314 2,
1315 &FSSTMetadata {
1316 uncompressed_lengths_ptype: fsst_array
1317 .uncompressed_lengths()
1318 .dtype()
1319 .as_ptype()
1320 .into(),
1321 codes_offsets_ptype: 0,
1323 }
1324 .encode_to_vec(),
1325 &buffers,
1326 &children.as_slice(),
1327 ),
1328 &array_session(),
1329 )?;
1330
1331 let decompressed =
1332 fsst.execute::<VarBinViewArray>(&mut array_session().create_execution_ctx())?;
1333 let mask = decompressed
1334 .validity()?
1335 .execute_mask(decompressed.len(), &mut ctx)?;
1336 assert!(mask.value(0));
1337 assert_eq!(decompressed.bytes_at(0).as_slice(), b"abcabcab".as_ref());
1338 assert!(mask.value(1));
1339 assert_eq!(decompressed.bytes_at(1).as_slice(), b"defghijk".as_ref());
1340 Ok(())
1341 }
1342}