Consolidate filter implementation.

Eventual-Inc · May 17, 2023 · 2728bcf · 2728bcf
1 parent c41933f
commit 2728bcf
Showing 1 changed file with 3 additions and 66 deletions.
diff --git a/src/array/ops/filter.rs b/src/array/ops/filter.rs
@@ -1,81 +1,18 @@
 use crate::{
     array::DataArray,
-    datatypes::{
-        logical::DateArray, BinaryArray, BooleanArray, DaftNumericType, ExtensionArray,
-        FixedSizeListArray, ListArray, NullArray, StructArray, Utf8Array,
-    },
+    datatypes::{logical::DateArray, BooleanArray, DaftArrowBackedType},
     error::DaftResult,
 };
 
 use super::as_arrow::AsArrow;
 
 impl<T> DataArray<T>
 where
-    T: DaftNumericType,
+    T: DaftArrowBackedType,
 {
-    pub fn filter(&self, mask: &BooleanArray) -> DaftResult<Self> {
-        let result = arrow2::compute::filter::filter(self.as_arrow(), mask.as_arrow())?;
-        Self::try_from((self.field.clone(), result))
-    }
-}
-
-impl Utf8Array {
-    pub fn filter(&self, mask: &BooleanArray) -> DaftResult<Self> {
-        let result = arrow2::compute::filter::filter(self.as_arrow(), mask.as_arrow())?;
-        Self::try_from((self.field.clone(), result))
-    }
-}
-
-impl BinaryArray {
-    pub fn filter(&self, mask: &BooleanArray) -> DaftResult<Self> {
-        let result = arrow2::compute::filter::filter(self.as_arrow(), mask.as_arrow())?;
-        Self::try_from((self.field.clone(), result))
-    }
-}
-
-impl BooleanArray {
-    pub fn filter(&self, mask: &BooleanArray) -> DaftResult<Self> {
-        let result = arrow2::compute::filter::filter(self.as_arrow(), mask.as_arrow())?;
-        Self::try_from((self.field.clone(), result))
-    }
-}
-
-impl NullArray {
-    pub fn filter(&self, mask: &BooleanArray) -> DaftResult<Self> {
-        let set_bits = mask.len() - mask.as_arrow().values().unset_bits();
-        Ok(NullArray::full_null(
-            self.name(),
-            self.data_type(),
-            set_bits,
-        ))
-    }
-}
-
-impl ListArray {
-    pub fn filter(&self, mask: &BooleanArray) -> DaftResult<Self> {
-        let result = arrow2::compute::filter::filter(self.as_arrow(), mask.as_arrow())?;
-        Self::try_from((self.field.clone(), result))
-    }
-}
-
-impl FixedSizeListArray {
-    pub fn filter(&self, mask: &BooleanArray) -> DaftResult<Self> {
-        let result = arrow2::compute::filter::filter(self.as_arrow(), mask.as_arrow())?;
-        Self::try_from((self.field.clone(), result))
-    }
-}
-
-impl StructArray {
-    pub fn filter(&self, mask: &BooleanArray) -> DaftResult<Self> {
-        let result = arrow2::compute::filter::filter(self.as_arrow(), mask.as_arrow())?;
-        Self::try_from((self.field.clone(), result))
-    }
-}
-
-impl ExtensionArray {
     pub fn filter(&self, mask: &BooleanArray) -> DaftResult<Self> {
         let result = arrow2::compute::filter::filter(self.data(), mask.as_arrow())?;
-        DataArray::try_from((self.field.clone(), result))
+        Self::try_from((self.field.clone(), result))
     }
 }