Repository: spiraldb/vortex
Pull request #2962: perf[duckdb]: Partially canonicalize chunk(2, [dict, dict]) and cache the values 
Author: login: joseph-isaacs
Created: 2025-04-10T09:41:14Z | Merged: 2025-04-10T16:07:34Z



Diff:
diff --git a/vortex-duckdb/src/convert/array/chunked.rs b/vortex-duckdb/src/convert/array/chunked.rs
new file mode 100644
index 00000000000..8d07d839461
--- /dev/null
+++ b/vortex-duckdb/src/convert/array/chunked.rs
@@ -0,0 +1,87 @@
+use duckdb::vtab::arrow::WritableVector;
+use vortex_array::arrays::ChunkedArray;
+use vortex_array::{Array, IntoArray};
+use vortex_dict::DictArray;
+use vortex_error::VortexResult;
+
+use crate::convert::array::to_duckdb;
+use crate::{ConversionCache, ToDuckDB};
+
+impl ToDuckDB for ChunkedArray {
+    fn to_duckdb(
+        &self,
+        chunk: &mut dyn WritableVector,
+        cache: &mut ConversionCache,
+    ) -> VortexResult<()> {
+        let chunks = self.chunks();
+        if chunks.len() == 1 {
+            return to_duckdb(&self.chunks()[0], chunk, cache);
+        } else if chunks.len() == 2 {
+            // It is common that a 2k split can contain a chunked array containing a pair of dictionaries
+            // We this is special cased and handled without a usual canonical, once there is a way to
+            // pre-canonicalize or cache, then this can be removed
+            if let (Some(dict1), Some(dict2)) = (
+                chunks[0].as_any().downcast_ref::<DictArray>(),
+                chunks[1].as_any().downcast_ref::<DictArray>(),
+            ) {
+                let canon_values1 = cache.cached_array(dict1.values())?;
+                let canon_values2 = cache.cached_array(dict2.values())?;
+
+                let chunked = ChunkedArray::try_new(
+                    vec![
+                        DictArray::try_new(dict1.codes().clone(), canon_values1)?.into_array(),
+                        DictArray::try_new(dict2.codes().clone(), canon_values2)?.into_array(),
+                    ],
+                    self.dtype().clone(),
+                )?;
+                return to_duckdb(&chunked.to_canonical()?.into_array(), chunk, cache);
+            }
+        }
+        to_duckdb(&self.to_canonical()?.into_array(), chunk, cache)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use duckdb::core::{DataChunkHandle, LogicalTypeHandle, LogicalTypeId};
+    use vortex_array::arrays::{ChunkedArray, StructArray};
+    use vortex_array::compute::slice;
+    use vortex_array::{Array, IntoArray};
+    use vortex_buffer::buffer;
+    use vortex_dict::DictArray;
+
+    use crate::{ConversionCache, to_duckdb_chunk};
+
+    #[test]
+    fn chunked_of_dict_to_duckdb() {
+        let dict1 = DictArray::try_new(
+            buffer![0u32, 1, 2, 2].into_array(),
+            buffer![0, 1, 2, 3, 4, 5, 6, 7].into_array(),
+        )
+        .unwrap();
+        let dict2 = DictArray::try_new(
+            buffer![0u32, 1, 2, 2].into_array(),
+            buffer![0, 1, 2, 3, 4, 5, 6, 7].into_array(),
+        )
+        .unwrap();
+        let dtype = dict1.dtype().clone();
+        let chunk =
+            ChunkedArray::try_new(vec![dict1.into_array(), dict2.into_array()], dtype).unwrap();
+
+        let sliced = slice(&chunk, 2, 7).unwrap();
+
+        let struct_ = StructArray::from_fields(&[("a", sliced)]).unwrap();
+        let mut cache = ConversionCache::default();
+        let mut data_chunk =
+            DataChunkHandle::new(&[LogicalTypeHandle::from(LogicalTypeId::Integer)]);
+
+        to_duckdb_chunk(&struct_, &mut data_chunk, &mut cache).unwrap();
+
+        assert_eq!(
+            format!("{:?}", data_chunk),
+            r#"Chunk - [1 Columns]
+- FLAT INTEGER: 5 = [ 2, 2, 0, 1, 2]
+"#
+        )
+    }
+}
diff --git a/vortex-duckdb/src/convert/array/mod.rs b/vortex-duckdb/src/convert/array/mod.rs
index 4e3ef78a6a2..74f47895223 100644
--- a/vortex-duckdb/src/convert/array/mod.rs
+++ b/vortex-duckdb/src/convert/array/mod.rs
@@ -1,7 +1,10 @@
+mod chunked;
 mod data_chunk_adaptor;
 mod run_end;
 mod varbinview;
 
+use std::sync::Arc;
+
 use arrow_array::ArrayRef as ArrowArrayRef;
 pub use data_chunk_adaptor::NamedDataChunk;
 use duckdb::core::{DataChunkHandle, FlatVector, SelectionVector};
@@ -18,7 +21,7 @@ use vortex_array::compute::{take, to_arrow_preferred};
 use vortex_array::validity::Validity;
 use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::vtable::EncodingVTable;
-use vortex_array::{Array, ArrayRef, ArrayStatistics, IntoArray, ToCanonical};
+use vortex_array::{Array, ArrayRef, ArrayStatistics, Canonical, IntoArray, ToCanonical};
 use vortex_dict::{DictArray, DictEncoding};
 use vortex_dtype::{NativePType, match_each_integer_ptype};
 use vortex_error::{VortexExpect, VortexResult, vortex_err};
@@ -32,7 +35,8 @@ use crate::{DUCKDB_STANDARD_VECTOR_SIZE, ToDuckDBType};
 
 #[derive(Default)]
 pub struct ConversionCache {
-    pub values_cache: HashMap<usize, FlatVector>,
+    pub values_cache: HashMap<usize, (ArrayRef, FlatVector)>,
+    pub canonical_cache: HashMap<usize, (ArrayRef, Canonical)>,
     // A value which must be unique for a given duckdb pipeline.
     pub instance_id: u64,
 }
@@ -44,6 +48,39 @@ impl ConversionCache {
             ..Self::default()
         }
     }
+
+    fn insert_cached_array(
+        &mut self,
+        arr_value: usize,
+        array: &ArrayRef,
+    ) -> VortexResult<ArrayRef> {
+        let canon = array.to_canonical()?;
+        self.canonical_cache
+            .insert(arr_value, (array.clone(), canon));
+        Ok(self
+            .canonical_cache
+            .get(&arr_value)
+            .vortex_expect("just added")
+            .1
+            .clone()
+            .into_array())
+    }
+
+    pub fn cached_array(&mut self, array: &ArrayRef) -> VortexResult<ArrayRef> {
+        let arr_value = Arc::as_ptr(array).addr();
+
+        let entry = self.canonical_cache.get(&arr_value);
+        match entry {
+            None => self.insert_cached_array(arr_value, array),
+            Some((cached_array_ref, cached_canonical)) => {
+                if Arc::ptr_eq(cached_array_ref, array) {
+                    Ok(cached_canonical.clone().into_array())
+                } else {
+                    self.insert_cached_array(arr_value, array)
+                }
+            }
+        }
+    }
 }
 
 pub trait ToDuckDB {
@@ -117,19 +154,26 @@ fn try_to_duckdb(
     }
 }
 
-impl ToDuckDB for ChunkedArray {
-    fn to_duckdb(
-        &self,
-        chunk: &mut dyn WritableVector,
-        cache: &mut ConversionCache,
-    ) -> VortexResult<()> {
-        // TODO(joe): support multi-chunk arrays without canonical.
-        if self.chunks().len() > 1 {
-            to_duckdb(&self.to_canonical()?.into_array(), chunk, cache)
-        } else {
-            to_duckdb(&self.chunks()[0], chunk, cache)
-        }
-    }
+fn create_and_insert_duckdb_dict_value_array_into_cache(
+    values: &ArrayRef,
+    cache: &mut ConversionCache,
+    value_ptr: usize,
+) -> VortexResult<FlatVector> {
+    let mut value_vector = FlatVector::allocate_new_vector_with_capacity(
+        values.dtype().to_duckdb_type()?,
+        values.len(),
+    );
+    let cached_array = cache.cached_array(values)?;
+    to_duckdb(&cached_array, &mut value_vector, cache)?;
+    cache
+        .values_cache
+        .insert(value_ptr, (values.clone(), value_vector));
+    Ok(cache
+        .values_cache
+        .get(&value_ptr)
+        .vortex_expect("just added")
+        .1
+        .clone())
 }
 
 impl ToDuckDB for DictArray {
@@ -138,18 +182,20 @@ impl ToDuckDB for DictArray {
         chunk: &mut dyn WritableVector,
         cache: &mut ConversionCache,
     ) -> VortexResult<()> {
+        let values = self.values();
+
         // Note you can only have nullable values (not codes/selection vectors),
         // so we cannot assign a selection vector.
         if !self.codes().all_valid()? {
-            let values = take(self.values(), self.codes())?;
+            let values = take(values, self.codes())?;
             return to_duckdb(&values, chunk, cache);
         };
 
-        let value_ptr = (self.values().as_ref() as *const dyn Array as *const ()) as usize;
+        let value_ptr = Arc::as_ptr(values).addr();
 
         let mut vector: FlatVector = if self.values().len() <= DUCKDB_STANDARD_VECTOR_SIZE {
             // If the values fit into a single vector, put the values in the pre-allocated vector.
-            to_duckdb(self.values(), chunk, cache)?;
+            to_duckdb(values, chunk, cache)?;
             chunk.flat_vector()
         } else {
             // If the values don't fit allocated a larger vector and that the data chunk vector
@@ -157,26 +203,25 @@ impl ToDuckDB for DictArray {
             let entry = cache.values_cache.get(&value_ptr);
             let value_vector = match entry {
                 None => {
-                    let mut value_vector = FlatVector::allocate_new_vector_with_capacity(
-                        self.values().dtype().to_duckdb_type()?,
-                        self.values().len(),
-                    );
-                    to_duckdb(self.values(), &mut value_vector, cache)?;
-                    cache.values_cache.insert(value_ptr, value_vector);
-                    cache
-                        .values_cache
-                        .get(&value_ptr)
-                        .vortex_expect("just added")
+                    create_and_insert_duckdb_dict_value_array_into_cache(values, cache, value_ptr)?
+                }
+                Some((cached_array_ref, entry)) => {
+                    if Arc::ptr_eq(cached_array_ref, values) {
+                        entry.clone()
+                    } else {
+                        create_and_insert_duckdb_dict_value_array_into_cache(
+                            values, cache, value_ptr,
+                        )?
+                    }
                 }
-                Some(entry) => entry,
             };
 
             let mut vector = chunk.flat_vector();
-            vector.reference(value_vector);
+            vector.reference(&value_vector);
             vector
         };
         let sel = selection_vector_from_array(self.codes().to_primitive()?);
-        vector.slice(self.values().len() as u64, sel);
+        vector.slice(values.len() as u64, sel);
         vector.set_dictionary_id(format!("{}-{}", cache.instance_id, value_ptr));
         Ok(())
     }


Reviews:
user: login: AdamGS
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T09:57:39Z
user: login: AdamGS
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T09:58:28Z
user: login: 0ax1
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T09:59:27Z
user: login: joseph-isaacs
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T10:12:19Z
user: login: joseph-isaacs
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T10:12:59Z
user: login: joseph-isaacs
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T10:37:47Z
user: login: AdamGS
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T10:42:53Z
user: login: AdamGS
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T10:43:22Z
user: login: joseph-isaacs
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T13:14:28Z
user: login: onursatici
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T14:12:54Z
user: login: joseph-isaacs
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T14:17:35Z
user: login: 0ax1
body: lgtm but would be good if @AdamGS could also take a final look
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-10T16:07:33Z

Review comments:
diff_hunk: @@ -33,6 +36,7 @@ use crate::{DUCKDB_STANDARD_VECTOR_SIZE, ToDuckDBType};
 #[derive(Default)]
 pub struct ConversionCache {
     pub values_cache: HashMap<usize, FlatVector>,
+    pub canonical_cache: HashMap<usize, Canonical>,
path: vortex-duckdb/src/convert/array/mod.rs
user: login: AdamGS
body: what's the lifetime of the `ConvesionCache`? 
created_at: 2025-04-10T09:57:39Z
author_association: CONTRIBUTOR
diff_hunk: @@ -44,6 +48,23 @@ impl ConversionCache {
             ..Self::default()
         }
     }
+
+    pub fn cached_array(&mut self, array: &ArrayRef) -> VortexResult<ArrayRef> {
+        let arr_value = Arc::as_ptr(array).addr();
path: vortex-duckdb/src/convert/array/mod.rs
user: login: AdamGS
body: also using the ptr seems error prone and can cause some pretty horrible bugs
created_at: 2025-04-10T09:58:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -33,6 +36,7 @@ use crate::{DUCKDB_STANDARD_VECTOR_SIZE, ToDuckDBType};
 #[derive(Default)]
 pub struct ConversionCache {
     pub values_cache: HashMap<usize, FlatVector>,
+    pub canonical_cache: HashMap<usize, Canonical>,
path: vortex-duckdb/src/convert/array/mod.rs
user: login: 0ax1
body: along those lines, should we limit the number of entries, have some cache eviction policy?
created_at: 2025-04-10T09:59:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -33,6 +36,7 @@ use crate::{DUCKDB_STANDARD_VECTOR_SIZE, ToDuckDBType};
 #[derive(Default)]
 pub struct ConversionCache {
     pub values_cache: HashMap<usize, FlatVector>,
+    pub canonical_cache: HashMap<usize, Canonical>,
path: vortex-duckdb/src/convert/array/mod.rs
user: login: joseph-isaacs
body: The scope of the chunks returned by vortex after splitting 
created_at: 2025-04-10T10:12:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -44,6 +48,23 @@ impl ConversionCache {
             ..Self::default()
         }
     }
+
+    pub fn cached_array(&mut self, array: &ArrayRef) -> VortexResult<ArrayRef> {
+        let arr_value = Arc::as_ptr(array).addr();
path: vortex-duckdb/src/convert/array/mod.rs
user: login: joseph-isaacs
body: Any other ideas. I hate this too
created_at: 2025-04-10T10:12:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -33,6 +36,7 @@ use crate::{DUCKDB_STANDARD_VECTOR_SIZE, ToDuckDBType};
 #[derive(Default)]
 pub struct ConversionCache {
     pub values_cache: HashMap<usize, FlatVector>,
+    pub canonical_cache: HashMap<usize, Canonical>,
path: vortex-duckdb/src/convert/array/mod.rs
user: login: joseph-isaacs
body: Happy to limit, what is your exact concern here?
created_at: 2025-04-10T10:37:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -44,6 +48,23 @@ impl ConversionCache {
             ..Self::default()
         }
     }
+
+    pub fn cached_array(&mut self, array: &ArrayRef) -> VortexResult<ArrayRef> {
+        let arr_value = Arc::as_ptr(array).addr();
path: vortex-duckdb/src/convert/array/mod.rs
user: login: AdamGS
body: my two ideas are
1. you can do something similar to what @onursatici did (IDK if it ended up merging), that its tied to the lifetime of the array and can't accidentally return the wrong thing 
2. You can key the cache by column + row offest/range
created_at: 2025-04-10T10:42:53Z
author_association: CONTRIBUTOR
diff_hunk: @@ -33,6 +36,7 @@ use crate::{DUCKDB_STANDARD_VECTOR_SIZE, ToDuckDBType};
 #[derive(Default)]
 pub struct ConversionCache {
     pub values_cache: HashMap<usize, FlatVector>,
+    pub canonical_cache: HashMap<usize, Canonical>,
path: vortex-duckdb/src/convert/array/mod.rs
user: login: AdamGS
body: if `ConvesionCache` lives a long time, it can consume a ton of memory for data no one is reading
created_at: 2025-04-10T10:43:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -44,6 +48,23 @@ impl ConversionCache {
             ..Self::default()
         }
     }
+
+    pub fn cached_array(&mut self, array: &ArrayRef) -> VortexResult<ArrayRef> {
+        let arr_value = Arc::as_ptr(array).addr();
path: vortex-duckdb/src/convert/array/mod.rs
user: login: joseph-isaacs
body: Went for usize -> (Arc, Value) where we use the value ptr to get to arc and then get that using Arc::ptr_eq
created_at: 2025-04-10T13:14:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -44,6 +48,23 @@ impl ConversionCache {
             ..Self::default()
         }
     }
+
+    pub fn cached_array(&mut self, array: &ArrayRef) -> VortexResult<ArrayRef> {
+        let arr_value = Arc::as_ptr(array).addr();
path: vortex-duckdb/src/convert/array/mod.rs
user: login: onursatici
body: Is the idea here because you have the arc itself as a value in the cache you guarantee that it won't be dropped as long as it is in the cache? This way you know the same address the arc is pointing to won't be reused to allocate a new array
created_at: 2025-04-10T14:12:54Z
author_association: CONTRIBUTOR
diff_hunk: @@ -44,6 +48,23 @@ impl ConversionCache {
             ..Self::default()
         }
     }
+
+    pub fn cached_array(&mut self, array: &ArrayRef) -> VortexResult<ArrayRef> {
+        let arr_value = Arc::as_ptr(array).addr();
path: vortex-duckdb/src/convert/array/mod.rs
user: login: joseph-isaacs
body: exactly

created_at: 2025-04-10T14:17:35Z
author_association: CONTRIBUTOR

PR comments:
user: login: codspeed-hq[bot]
created_at: 2025-04-10T09:47:44Z
body: <!-- __CODSPEED_PERFORMANCE_REPORT_COMMENT__ -->
<!-- __CODSPEED_INSTRUMENTATION_PERFORMANCE_REPORT_COMMENT__ -->

## [CodSpeed Performance Report](https://codspeed.io/spiraldb/vortex/branches/ji%2Fchunk-dict-dict)

### Merging #2962 will **degrade performances by 11.5%**

<sub>Comparing <code>ji/chunk-dict-dict</code> (51c7030) with <code>develop</code> (756c9f6)</sub>



### Summary

`❌ 1` regressions  
`✅ 810` untouched benchmarks  


> :warning: _Please fix the performance issues or [acknowledge them on CodSpeed](https://codspeed.io/spiraldb/vortex/branches/ji%2Fchunk-dict-dict)._

### Benchmarks breakdown

|     | Benchmark | `BASE` | `HEAD` | Change |
| --- | --------- | ----------------------- | ------------------- | ------ |
| ❌ | `` encode_specialized `` | 1.1 ms | 1.2 ms | -11.5% |

author_association: NONE
user: login: joseph-isaacs
created_at: 2025-04-10T16:27:58Z
body: Looks like this had auto merge (next time if you approve maybe turn this off), @AdamGS If you have any comments I am happy to follow up
author_association: CONTRIBUTOR