[GitHub] [arrow-datafusion] alamb commented on a change in pull request #1841: Implement bitmap_distinct function using roaring bitmap

[GitBox <gi...@apache.org>]

alamb commented on a change in pull request #1841:
URL: https://github.com/apache/arrow-datafusion/pull/1841#discussion_r816834403



##########
File path: datafusion-physical-expr/Cargo.toml
##########
@@ -41,3 +41,4 @@ arrow = { version = "9.0.0", features = ["prettyprint"] }
 paste = "^1.0"
 ahash = { version = "0.7", default-features = false }
 ordered-float = "2.10"
+roaring = "0.8.1"

Review comment:
       What would you think about making this an optional dependency (much like crypto expressions, etc)  as defined on master?
   
   This would let anyone who wants this feature be able to use it, but would not require it for anyone who did not?
   
   https://github.com/apache/arrow-datafusion/blob/7eb3bd8/datafusion-physical-expr/Cargo.toml#L35-L39

##########
File path: datafusion-physical-expr/src/expressions/bitmap_distinct.rs
##########
@@ -0,0 +1,211 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Defines physical expressions that can evaluated at runtime during query execution
+
+use std::any::Any;
+
+use std::fmt::Debug;
+use std::ops::BitOrAssign;
+use std::sync::Arc;
+
+use arrow::array::{
+    Array, ArrayRef, BinaryArray, Int16Array, Int32Array, Int8Array, UInt16Array,
+    UInt32Array, UInt8Array,
+};
+use arrow::datatypes::{DataType, Field};
+use datafusion_common::{DataFusionError, Result, ScalarValue};
+use datafusion_expr::Accumulator;
+use roaring::RoaringBitmap;
+
+use crate::{AggregateExpr, PhysicalExpr};
+
+use super::format_state_name;
+
+/// BITMAP_DISTINCT aggregate expression
+#[derive(Debug)]
+pub struct BitMapDistinct {
+    name: String,
+    input_data_type: DataType,
+    expr: Arc<dyn PhysicalExpr>,
+}
+
+impl BitMapDistinct {
+    /// Create a new BitmapDistinct aggregate function.
+    pub fn new(
+        expr: Arc<dyn PhysicalExpr>,
+        name: impl Into<String>,
+        input_data_type: DataType,
+    ) -> Self {
+        Self {
+            name: name.into(),
+            input_data_type,
+            expr,
+        }
+    }
+}
+
+impl AggregateExpr for BitMapDistinct {
+    /// Return a reference to Any that can be used for downcasting
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    /// the field of the final result of this aggregation.
+    fn field(&self) -> Result<Field> {
+        Ok(Field::new(&self.name, DataType::UInt64, false))
+    }
+
+    fn create_accumulator(&self) -> Result<Box<dyn Accumulator>> {
+        let accumulator: Box<dyn Accumulator> = match &self.input_data_type {
+            DataType::UInt8
+            | DataType::UInt16
+            | DataType::UInt32
+            | DataType::Int8
+            | DataType::Int16
+            | DataType::Int32 => Box::new(BitmapDistinctCountAccumulator::try_new()),
+            other => {
+                return Err(DataFusionError::NotImplemented(format!(
+                    "Support for 'bitmap_distinct' for data type {} is not implemented",
+                    other
+                )))
+            }
+        };
+        Ok(accumulator)
+    }
+
+    fn state_fields(&self) -> Result<Vec<Field>> {
+        Ok(vec![Field::new(
+            &format_state_name(&self.name, "bitmap_registers"),
+            DataType::Binary,
+            false,
+        )])
+    }
+
+    fn expressions(&self) -> Vec<Arc<dyn PhysicalExpr>> {
+        vec![self.expr.clone()]
+    }
+
+    fn name(&self) -> &str {
+        &self.name
+    }
+}
+
+#[derive(Debug)]
+struct BitmapDistinctCountAccumulator {
+    bitmap: roaring::bitmap::RoaringBitmap,
+}
+
+impl BitmapDistinctCountAccumulator {
+    fn try_new() -> Self {
+        Self {
+            bitmap: RoaringBitmap::new(),
+        }
+    }
+}
+
+impl Accumulator for BitmapDistinctCountAccumulator {
+    //state() can be used by physical nodes to aggregate states together and send them over the network/threads, to combine values.
+    fn state(&self) -> Result<Vec<ScalarValue>> {
+        let mut bytes = vec![];
+        self.bitmap.serialize_into(&mut bytes).unwrap();
+        Ok(vec![ScalarValue::Binary(Some(bytes))])
+    }
+
+    fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> {
+        let value = &values[0];
+        if value.is_empty() {
+            return Ok(());
+        }
+        match value.data_type() {
+            DataType::Int8 => {
+                let array = value.as_any().downcast_ref::<Int8Array>().unwrap();
+                for i in 0..array.len() {

Review comment:
       This code doesn't seem to handle nulls (as in don't you have to check `array.is_valid(i)` prior to getting `array.value()`?
   
   A test case for null values would probably be useful 
   
   Maybe you could use an iterator like
   
   ```rust
   for value in array.iter() {
     match value {
       Some(v) => self.bitmap.insert(value as u32);
       None => // do something with NULLs here
     }
   }
   ```




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org