DuckDB FSST Export

Signed-off-by: Nicholas Gates <nick@nickgates.com>

Author: gatesn

vortex-duckdb/cpp/include/duckdb_vx/vector.h

@@ -55,15 +55,13 @@ void duckdb_vx_sequence_vector(duckdb_vector c_vector, int64_t start, int64_t st
 
 // Finalize a vector as an FSST vector using externally managed compressed string data and a
 // symbol table owned by Vortex.
-void duckdb_vx_fsst_vector_set(
-    duckdb_vector ffi_vector,
-    const uint64_t *symbols,
-    const uint8_t *symbol_lengths,
-    idx_t symbol_count,
-    idx_t string_block_limit,
-    idx_t count,
-    duckdb_vx_vector_buffer buffer
-);
+void duckdb_vx_fsst_vector_set(duckdb_vector ffi_vector,
+                               const uint64_t *symbols,
+                               const uint8_t *symbol_lengths,
+                               idx_t symbol_count,
+                               idx_t string_block_limit,
+                               idx_t count,
+                               duckdb_vx_vector_buffer buffer);
 
 // Returns whether the vector is currently an FSST vector.
 bool duckdb_vx_vector_is_fsst(duckdb_vector ffi_vector);

vortex-duckdb/cpp/vector.cpp

@@ -109,6 +109,10 @@ extern "C" void duckdb_vx_fsst_vector_set(duckdb_vector ffi_vector,
         decoder->symbol[i] = symbols[i];
     }
 
+    // DuckDB can reuse vector instances across chunk exports. Replace the FSST auxiliary buffer
+    // on each export so heap references to prior compressed byte buffers are dropped instead of
+    // accumulating for the lifetime of the reused vector.
+    vector->SetAuxiliary(make_buffer<VectorFSSTStringBuffer>());
     FSSTVector::RegisterDecoder(*vector, decoder_buffer, string_block_limit);
 
     if (buffer) {

vortex-duckdb/src/exporter/fsst.rs

@@ -15,6 +15,7 @@ use vortex::array::match_each_integer_ptype;
 use vortex::buffer::ByteBuffer;
 use vortex::encodings::fsst::FSST;
 use vortex::encodings::fsst::FSSTArrayExt;
+use vortex::error::vortex_err;
 use vortex::error::VortexResult;
 
 use crate::cpp;
@@ -102,7 +103,7 @@ impl ColumnExporter for FSSTExporter {
                 let start: usize = offsets[offset + row].as_();
                 let end: usize = offsets[offset + row + 1].as_();
                 let value = &bytes[start..end];
-                out[row] = PtrString::new(value);
+                out[row] = PtrString::new(value)?;
             }
         });
 
@@ -125,32 +126,33 @@ struct PtrString {
 }
 
 impl PtrString {
-    fn new(value: &[u8]) -> Self {
-        let length = u32::try_from(value.len()).expect("FSST code length must fit in u32");
+    fn new(value: &[u8]) -> VortexResult<Self> {
+        let length = u32::try_from(value.len())
+            .map_err(|_| vortex_err!("FSST code length {} exceeds u32", value.len()))?;
         if value.len() <= 12 {
             let mut inlined = [0_i8; 12];
             for (dst, src) in inlined.iter_mut().zip(value) {
                 *dst = *src as i8;
             }
-            Self {
+            Ok(Self {
                 value: PtrStringValue {
                     inlined: PtrStringInlined { length, inlined },
                 },
-            }
+            })
         } else {
             let mut prefix = [0_i8; 4];
             for (dst, src) in prefix.iter_mut().zip(value.iter().copied()) {
                 *dst = src as i8;
             }
-            Self {
+            Ok(Self {
                 value: PtrStringValue {
                     pointer: PtrStringPointer {
                         length,
                         prefix,
                         ptr: value.as_ptr().cast_mut().cast::<c_char>(),
                     },
                 },
-            }
+            })
         }
     }
 }