From 0e2fcd61b364e0a3f8d071e0e82838a5a1e1b1b4 Mon Sep 17 00:00:00 2001
From: Xiangpeng Hao <haoxiangpeng123@gmail.com>
Date: Sat, 19 Oct 2024 20:30:46 -0500
Subject: [PATCH] implement filter for dict

---
 encodings/dict/src/compute.rs | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/encodings/dict/src/compute.rs b/encodings/dict/src/compute.rs
index 5ca1e2a4b..ccd42f001 100644
--- a/encodings/dict/src/compute.rs
+++ b/encodings/dict/src/compute.rs
@@ -1,5 +1,5 @@
 use vortex::compute::unary::{scalar_at, scalar_at_unchecked, ScalarAtFn};
-use vortex::compute::{slice, take, ArrayCompute, SliceFn, TakeFn};
+use vortex::compute::{filter, slice, take, ArrayCompute, FilterFn, SliceFn, TakeFn};
 use vortex::{Array, IntoArray};
 use vortex_error::{VortexExpect, VortexResult};
 use vortex_scalar::Scalar;
@@ -18,6 +18,10 @@ impl ArrayCompute for DictArray {
     fn take(&self) -> Option<&dyn TakeFn> {
         Some(self)
     }
+
+    fn filter(&self) -> Option<&dyn FilterFn> {
+        Some(self)
+    }
 }
 
 impl ScalarAtFn for DictArray {
@@ -46,6 +50,13 @@ impl TakeFn for DictArray {
     }
 }
 
+impl FilterFn for DictArray {
+    fn filter(&self, predicate: &Array) -> VortexResult<Array> {
+        let codes = filter(self.codes(), predicate)?;
+        Self::try_new(codes, self.values()).map(|a| a.into_array())
+    }
+}
+
 impl SliceFn for DictArray {
     // TODO(robert): Add function to trim the dictionary
     fn slice(&self, start: usize, stop: usize) -> VortexResult<Array> {