From ae846ca0ddfa19c766f045a2348e46ea20f04c90 Mon Sep 17 00:00:00 2001
From: Jay Han <tunyum94@gmail.com>
Date: Thu, 12 Dec 2024 23:54:28 +0800
Subject: [PATCH 1/2] impl take_fn and filter_fn

---
 encodings/zigzag/src/compute.rs | 76 ++++++++++++++++++++++++++++-----
 1 file changed, 66 insertions(+), 10 deletions(-)

diff --git a/encodings/zigzag/src/compute.rs b/encodings/zigzag/src/compute.rs
index b68b542521..8f26cae349 100644
--- a/encodings/zigzag/src/compute.rs
+++ b/encodings/zigzag/src/compute.rs
@@ -1,4 +1,7 @@
-use vortex_array::compute::{scalar_at, slice, ComputeVTable, ScalarAtFn, SliceFn};
+use vortex_array::compute::{
+    filter, scalar_at, slice, take, ComputeVTable, FilterFn, FilterMask, ScalarAtFn, SliceFn,
+    TakeFn,
+};
 use vortex_array::variants::PrimitiveArrayTrait;
 use vortex_array::{ArrayDType, ArrayData, IntoArrayData};
 use vortex_dtype::match_each_unsigned_integer_ptype;
@@ -9,6 +12,10 @@ use zigzag::{ZigZag as ExternalZigZag, ZigZag};
 use crate::{ZigZagArray, ZigZagEncoding};
 
 impl ComputeVTable for ZigZagEncoding {
+    fn filter_fn(&self) -> Option<&dyn FilterFn<ArrayData>> {
+        Some(self)
+    }
+
     fn scalar_at_fn(&self) -> Option<&dyn ScalarAtFn<ArrayData>> {
         Some(self)
     }
@@ -16,6 +23,10 @@ impl ComputeVTable for ZigZagEncoding {
     fn slice_fn(&self) -> Option<&dyn SliceFn<ArrayData>> {
         Some(self)
     }
+
+    fn take_fn(&self) -> Option<&dyn TakeFn<ArrayData>> {
+        Some(self)
+    }
 }
 
 impl ScalarAtFn<ZigZagArray> for ZigZagEncoding {
@@ -41,6 +52,26 @@ impl ScalarAtFn<ZigZagArray> for ZigZagEncoding {
     }
 }
 
+impl SliceFn<ZigZagArray> for ZigZagEncoding {
+    fn slice(&self, array: &ZigZagArray, start: usize, stop: usize) -> VortexResult<ArrayData> {
+        Ok(ZigZagArray::try_new(slice(array.encoded(), start, stop)?)?.into_array())
+    }
+}
+
+impl TakeFn<ZigZagArray> for ZigZagEncoding {
+    fn take(&self, array: &ZigZagArray, indices: &ArrayData) -> VortexResult<ArrayData> {
+        let encoded = take(array.encoded(), indices)?;
+        Ok(ZigZagArray::try_new(encoded)?.into_array())
+    }
+}
+
+impl FilterFn<ZigZagArray> for ZigZagEncoding {
+    fn filter(&self, array: &ZigZagArray, mask: FilterMask) -> VortexResult<ArrayData> {
+        let encoded = filter(&array.encoded(), mask)?;
+        Ok(ZigZagArray::try_new(encoded)?.into_array())
+    }
+}
+
 trait ZigZagEncoded {
     type Int: ZigZag;
 }
@@ -61,18 +92,14 @@ impl ZigZagEncoded for u64 {
     type Int = i64;
 }
 
-impl SliceFn<ZigZagArray> for ZigZagEncoding {
-    fn slice(&self, array: &ZigZagArray, start: usize, stop: usize) -> VortexResult<ArrayData> {
-        Ok(ZigZagArray::try_new(slice(array.encoded(), start, stop)?)?.into_array())
-    }
-}
-
 #[cfg(test)]
 mod tests {
-    use vortex_array::array::PrimitiveArray;
-    use vortex_array::compute::{scalar_at, search_sorted, SearchResult, SearchSortedSide};
+    use vortex_array::array::{BooleanBuffer, PrimitiveArray};
+    use vortex_array::compute::{
+        filter, scalar_at, search_sorted, take, SearchResult, SearchSortedSide,
+    };
     use vortex_array::validity::Validity;
-    use vortex_array::IntoArrayData;
+    use vortex_array::{IntoArrayData, IntoArrayVariant};
     use vortex_dtype::Nullability;
     use vortex_scalar::Scalar;
 
@@ -100,4 +127,33 @@ mod tests {
             Scalar::primitive(-160, Nullability::Nullable)
         );
     }
+
+    #[test]
+    fn take_zigzag() {
+        let zigzag =
+            ZigZagArray::encode(&PrimitiveArray::from(vec![-189, -160, 1]).into_array()).unwrap();
+        let indices = PrimitiveArray::from(vec![0, 2]).into_array();
+        let actual = take(zigzag, indices).unwrap().into_primitive().unwrap();
+        let expected = ZigZagArray::encode(&PrimitiveArray::from(vec![-189, 1]).into_array())
+            .unwrap()
+            .into_primitive()
+            .unwrap();
+        assert_eq!(actual.into_buffer(), expected.into_buffer());
+    }
+
+    #[test]
+    fn filter_zigzag() {
+        let zigzag =
+            ZigZagArray::encode(&PrimitiveArray::from(vec![-189, -160, 1]).into_array()).unwrap();
+        let filter_mask = BooleanBuffer::from(vec![true, false, true]).into();
+        let actual = filter(&zigzag.into_array(), filter_mask)
+            .unwrap()
+            .into_primitive()
+            .unwrap();
+        let expected = ZigZagArray::encode(&PrimitiveArray::from(vec![-189, 1]).into_array())
+            .unwrap()
+            .into_primitive()
+            .unwrap();
+        assert_eq!(actual.into_buffer(), expected.into_buffer());
+    }
 }

From db61dde2615264627a1820c2641173d04059096b Mon Sep 17 00:00:00 2001
From: Jay Han <tunyum94@gmail.com>
Date: Fri, 13 Dec 2024 09:39:18 +0800
Subject: [PATCH 2/2] change code position

---
 encodings/zigzag/src/compute.rs | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/encodings/zigzag/src/compute.rs b/encodings/zigzag/src/compute.rs
index 8f26cae349..7277749f92 100644
--- a/encodings/zigzag/src/compute.rs
+++ b/encodings/zigzag/src/compute.rs
@@ -29,6 +29,13 @@ impl ComputeVTable for ZigZagEncoding {
     }
 }
 
+impl FilterFn<ZigZagArray> for ZigZagEncoding {
+    fn filter(&self, array: &ZigZagArray, mask: FilterMask) -> VortexResult<ArrayData> {
+        let encoded = filter(&array.encoded(), mask)?;
+        Ok(ZigZagArray::try_new(encoded)?.into_array())
+    }
+}
+
 impl ScalarAtFn<ZigZagArray> for ZigZagEncoding {
     fn scalar_at(&self, array: &ZigZagArray, index: usize) -> VortexResult<Scalar> {
         let scalar = scalar_at(array.encoded(), index)?;
@@ -65,13 +72,6 @@ impl TakeFn<ZigZagArray> for ZigZagEncoding {
     }
 }
 
-impl FilterFn<ZigZagArray> for ZigZagEncoding {
-    fn filter(&self, array: &ZigZagArray, mask: FilterMask) -> VortexResult<ArrayData> {
-        let encoded = filter(&array.encoded(), mask)?;
-        Ok(ZigZagArray::try_new(encoded)?.into_array())
-    }
-}
-
 trait ZigZagEncoded {
     type Int: ZigZag;
 }