From 8372b000365e796e33d7f22bf4fd468dfd439ffc Mon Sep 17 00:00:00 2001
From: Antoine Dupuis <1048669+mnshdw@users.noreply.github.com>
Date: Fri, 28 Aug 2026 17:21:12 +0200
Subject: [PATCH] fix(dict): clear encoder state in DictEncoder::reset

Fixes #9692

Signed-off-by: Antoine Dupuis <1048669+mnshdw@users.noreply.github.com>
---
 vortex-array/src/builders/dict/bytes.rs     | 31 +++++++++++++++++++++
 vortex-array/src/builders/dict/primitive.rs | 28 ++++++++++++++++++-
 2 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/vortex-array/src/builders/dict/bytes.rs b/vortex-array/src/builders/dict/bytes.rs
index 361735536ed..be05adf4100 100644
--- a/vortex-array/src/builders/dict/bytes.rs
+++ b/vortex-array/src/builders/dict/bytes.rs
@@ -286,6 +286,10 @@ impl<Code: UnsignedPType> DictEncoder for BytesDictBuilder<Code> {
     }
 
     fn reset(&mut self) -> ArrayRef {
+        if let Some(lookup) = self.lookup.as_mut() {
+            lookup.clear();
+        }
+        self.null_code = OnceCell::new();
         let views = mem::take(&mut self.views).freeze();
         let buffer = mem::take(&mut self.values).freeze();
         let value_nulls = mem::take(&mut self.values_nulls).freeze();
@@ -325,8 +329,11 @@ mod test {
     use crate::arrays::VarBinViewArray;
     use crate::arrays::dict::DictArraySlotsExt;
     use crate::arrays::varbinview::BinaryView;
+    use crate::assert_arrays_eq;
     use crate::buffer::BufferHandle;
+    use crate::builders::dict::UNCONSTRAINED;
     use crate::builders::dict::dict_encode;
+    use crate::builders::dict::dict_encoder;
     use crate::dtype::DType;
     use crate::dtype::Nullability;
     use crate::validity::Validity;
@@ -430,4 +437,28 @@ mod test {
         assert_eq!(codes.as_slice::<u8>(), &[0, 0, 1, 1, 0, 1, 0, 1]);
         Ok(())
     }
+
+    #[test]
+    fn reset_clears_dict() -> VortexResult<()> {
+        let mut ctx = SESSION.create_execution_ctx();
+        let first = VarBinViewArray::from_iter_str(["one", "two"]).into_array();
+        let mut encoder = dict_encoder(&first, &UNCONSTRAINED);
+
+        assert_arrays_eq!(
+            encoder.encode(&first, &mut ctx)?,
+            PrimitiveArray::from_iter([0u64, 1]),
+            &mut ctx
+        );
+        assert_arrays_eq!(encoder.reset(), first, &mut ctx);
+
+        let second = VarBinViewArray::from_iter_str(["one", "three"]).into_array();
+        assert_arrays_eq!(
+            encoder.encode(&second, &mut ctx)?,
+            PrimitiveArray::from_iter([0u64, 1]),
+            &mut ctx
+        );
+        assert_arrays_eq!(encoder.reset(), second, &mut ctx);
+
+        Ok(())
+    }
 }
diff --git a/vortex-array/src/builders/dict/primitive.rs b/vortex-array/src/builders/dict/primitive.rs
index dea2932367b..4c8d8ae6856 100644
--- a/vortex-array/src/builders/dict/primitive.rs
+++ b/vortex-array/src/builders/dict/primitive.rs
@@ -183,8 +183,10 @@ where
     }
 
     fn reset(&mut self) -> ArrayRef {
+        self.lookup.clear();
+        self.null_code = OnceCell::new();
         PrimitiveArray::new(
-            self.values.clone(),
+            mem::take(&mut self.values),
             Validity::from_bit_buffer(mem::take(&mut self.values_nulls).freeze(), self.nullability),
         )
         .into_array()
@@ -206,7 +208,9 @@ mod test {
     use crate::VortexSessionExecute;
     use crate::arrays::dict::DictArraySlotsExt;
     use crate::assert_arrays_eq;
+    use crate::builders::dict::UNCONSTRAINED;
     use crate::builders::dict::dict_encode;
+    use crate::builders::dict::dict_encoder;
     use crate::builders::dict::primitive::PrimitiveArray;
 
     static SESSION: LazyLock<VortexSession> = LazyLock::new(crate::array_session);
@@ -247,4 +251,26 @@ mod test {
             PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array();
         assert_arrays_eq!(dict.values(), expected_values, &mut ctx);
     }
+
+    #[test]
+    fn reset_clears_dict() {
+        let mut ctx = SESSION.create_execution_ctx();
+        let first = PrimitiveArray::from_option_iter([Some(1i32), None, Some(3)]).into_array();
+        let mut encoder = dict_encoder(&first, &UNCONSTRAINED);
+
+        assert_arrays_eq!(
+            encoder.encode(&first, &mut ctx).unwrap(),
+            buffer![0u32, 1, 2].into_array(),
+            &mut ctx
+        );
+        assert_arrays_eq!(encoder.reset(), first, &mut ctx);
+
+        let second = PrimitiveArray::from_option_iter([Some(4i32), Some(5)]).into_array();
+        assert_arrays_eq!(
+            encoder.encode(&second, &mut ctx).unwrap(),
+            buffer![0u32, 1].into_array(),
+            &mut ctx
+        );
+        assert_arrays_eq!(encoder.reset(), second, &mut ctx);
+    }
 }