refactor: Pre-deserialization of Bitmap

Author: KKould

src/common/io/Cargo.toml

@@ -27,7 +27,7 @@ micromarshal = { workspace = true }
 roaring = { workspace = true }
 scroll = { workspace = true }
 serde = { workspace = true }
-smallvec = { workspace = true }
+smallvec = { workspace = true, features = ["serde"] }
 wkt = { workspace = true }
 
 [dev-dependencies]

src/common/io/src/bitmap.rs

@@ -14,17 +14,24 @@
 
 use std::fmt;
 use std::io;
+use std::io::Read;
+use std::io::Write;
 use std::iter::FromIterator;
 use std::mem;
+use std::mem::size_of;
 use std::ops::BitAndAssign;
 use std::ops::BitOrAssign;
 use std::ops::BitXorAssign;
 use std::ops::SubAssign;
 
+use borsh::BorshDeserialize;
+use borsh::BorshSerialize;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use roaring::treemap::Iter;
 use roaring::RoaringTreemap;
+use serde::Deserialize;
+use serde::Serialize;
 use smallvec::SmallVec;
 
 // https://github.com/ClickHouse/ClickHouse/blob/516a6ed6f8bd8c5f6eed3a10e9037580b2fb6152/src/AggregateFunctions/AggregateFunctionGroupBitmapData.h#L914
@@ -38,7 +45,7 @@ pub const HYBRID_HEADER_LEN: usize = 4;
 type SmallBitmap = SmallVec<[u64; LARGE_THRESHOLD]>;
 
 #[allow(clippy::large_enum_variant)]
-#[derive(Clone, PartialEq)]
+#[derive(Clone, PartialEq, Serialize, Deserialize)]
 pub enum HybridBitmap {
     Small(SmallBitmap),
     Large(RoaringTreemap),
@@ -62,6 +69,15 @@ impl HybridBitmap {
         }
     }
 
+    pub fn memory_size(&self) -> usize {
+        match self {
+            HybridBitmap::Small(set) => {
+                size_of::<SmallBitmap>() + set.capacity() * size_of::<u64>()
+            }
+            HybridBitmap::Large(tree) => tree.serialized_size(),
+        }
+    }
+
     pub fn is_empty(&self) -> bool {
         self.len() == 0
     }
@@ -447,6 +463,33 @@ impl fmt::Debug for HybridBitmap {
     }
 }
 
+impl BorshSerialize for HybridBitmap {
+    fn serialize<W: Write>(&self, writer: &mut W) -> std::io::Result<()> {
+        let mut buf = Vec::new();
+        self.serialize_into(&mut buf)?;
+        let len = u64::try_from(buf.len()).map_err(|_| {
+            io::Error::new(
+                io::ErrorKind::InvalidData,
+                format!("hybrid bitmap serialized size overflow: {}", buf.len()),
+            )
+        })?;
+        writer.write_all(&len.to_le_bytes())?;
+        writer.write_all(&buf)
+    }
+}
+
+impl BorshDeserialize for HybridBitmap {
+    fn deserialize_reader<R: Read>(reader: &mut R) -> std::io::Result<Self> {
+        let mut len_buf = [0u8; 8];
+        reader.read_exact(&mut len_buf)?;
+        let len = u64::from_le_bytes(len_buf) as usize;
+        let mut buf = vec![0u8; len];
+        reader.read_exact(&mut buf)?;
+        deserialize_bitmap(&buf)
+            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))
+    }
+}
+
 pub fn parse_bitmap(buf: &[u8]) -> Result<HybridBitmap> {
     std::str::from_utf8(buf)
         .map_err(|_| ())

src/common/native/src/read/array/binary.rs

@@ -15,13 +15,15 @@
 use std::io::Cursor;
 
 use databend_common_column::binary::BinaryColumn;
+use databend_common_expression::types::bitmap::BitmapColumn;
 use databend_common_expression::types::Bitmap;
 use databend_common_expression::types::Buffer;
 use databend_common_expression::types::GeographyColumn;
 use databend_common_expression::Column;
 use databend_common_expression::TableDataType;
 
 use crate::compression::binary::decompress_binary;
+use crate::error::Error;
 use crate::error::Result;
 use crate::nested::InitNested;
 use crate::nested::NestedState;
@@ -132,25 +134,27 @@ fn try_new_binary_column(
     validity: Option<Bitmap>,
 ) -> Result<Column> {
     let column = BinaryColumn::new(values, offsets);
-    Ok(binary_column_to_column(data_type, column, validity))
+    binary_column_to_column(data_type, column, validity)
 }
 
 fn binary_column_to_column(
     data_type: &TableDataType,
     column: BinaryColumn,
     validity: Option<Bitmap>,
-) -> Column {
+) -> Result<Column> {
     let col = match data_type.remove_nullable() {
         TableDataType::Binary => Column::Binary(column),
-        TableDataType::Bitmap => Column::Bitmap(column),
+        TableDataType::Bitmap => Column::Bitmap(Box::new(
+            BitmapColumn::from_binary(column).map_err(|e| Error::OutOfSpec(e.to_string()))?,
+        )),
         TableDataType::Variant => Column::Variant(column),
         TableDataType::Geometry => Column::Geometry(column),
         TableDataType::Geography => Column::Geography(GeographyColumn(column)),
         _ => unreachable!(),
     };
     if data_type.is_nullable() {
-        col.wrap_nullable(validity)
+        Ok(col.wrap_nullable(validity))
     } else {
-        col
+        Ok(col)
     }
 }

src/common/native/src/write/serialize.rs

@@ -184,11 +184,14 @@ impl<'a, W: Write> ValueVisitor for WriteVisitor<'a, W> {
                 )
             }
 
-            Column::Binary(b)
-            | Column::Bitmap(b)
-            | Column::Variant(b)
-            | Column::Geography(GeographyColumn(b))
-            | Column::Geometry(b) => write_binary::<W>(
+            Column::Binary(b) | Column::Variant(b) | Column::Geometry(b) => write_binary::<W>(
+                self.w,
+                &b,
+                self.validity.clone(),
+                self.write_options.clone(),
+                self.scratch,
+            ),
+            Column::Geography(GeographyColumn(b)) => write_binary::<W>(
                 self.w,
                 &b,
                 self.validity.clone(),

src/query/expression/src/aggregate/group_hash.rs

@@ -19,6 +19,7 @@ use databend_common_column::types::Index;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 
+use crate::types::bitmap::BitmapColumn;
 use crate::types::i256;
 use crate::types::number::Number;
 use crate::types::AccessType;
@@ -47,7 +48,6 @@ use crate::types::StringType;
 use crate::types::TimestampType;
 use crate::types::ValueType;
 use crate::types::VariantType;
-use crate::utils::bitmap::normalize_bitmap_column;
 use crate::visitor::ValueVisitor;
 use crate::with_decimal_mapped_type;
 use crate::with_number_mapped_type;
@@ -167,33 +167,40 @@ impl<const IS_FIRST: bool> ValueVisitor for HashVisitor<'_, IS_FIRST> {
     }
 
     fn visit_binary(&mut self, column: BinaryColumn) -> Result<()> {
-        self.combine_group_hash_string_column::<BinaryType>(&column);
+        self.combine_group_hash_string_column::<BinaryType, _>(&column, |x| Ok(x.agg_hash()))?;
         Ok(())
     }
 
     fn visit_string(&mut self, column: StringColumn) -> Result<()> {
-        self.combine_group_hash_string_column::<StringType>(&column);
+        self.combine_group_hash_string_column::<StringType, _>(&column, |x| {
+            Ok(x.as_bytes().agg_hash())
+        })?;
         Ok(())
     }
 
-    fn visit_bitmap(&mut self, column: BinaryColumn) -> Result<()> {
-        let column = normalize_bitmap_column(&column);
-        self.combine_group_hash_string_column::<BitmapType>(column.as_ref());
+    fn visit_bitmap(&mut self, column: BitmapColumn) -> Result<()> {
+        let mut bytes = Vec::new();
+        self.combine_group_hash_string_column::<BitmapType, _>(&column, |x| {
+            x.serialize_into(&mut bytes).unwrap();
+            let hash = bytes.agg_hash();
+            bytes.clear();
+            Ok(hash)
+        })?;
         Ok(())
     }
 
     fn visit_variant(&mut self, column: BinaryColumn) -> Result<()> {
-        self.combine_group_hash_string_column::<VariantType>(&column);
+        self.combine_group_hash_string_column::<VariantType, _>(&column, |x| Ok(x.agg_hash()))?;
         Ok(())
     }
 
     fn visit_geometry(&mut self, column: BinaryColumn) -> Result<()> {
-        self.combine_group_hash_string_column::<GeometryType>(&column);
+        self.combine_group_hash_string_column::<GeometryType, _>(&column, |x| Ok(x.agg_hash()))?;
         Ok(())
     }
 
     fn visit_geography(&mut self, column: GeographyColumn) -> Result<()> {
-        self.combine_group_hash_string_column::<GeographyType>(&column);
+        self.combine_group_hash_string_column::<GeographyType, _>(&column, |x| Ok(x.0.agg_hash()))?;
         Ok(())
     }
 
@@ -225,20 +232,25 @@ impl<const IS_FIRST: bool> HashVisitor<'_, IS_FIRST> {
         }
     }
 
-    fn combine_group_hash_string_column<T>(&mut self, col: &T::Column)
+    fn combine_group_hash_string_column<T, F>(
+        &mut self,
+        col: &T::Column,
+        mut agg_hash: F,
+    ) -> Result<()>
     where
         T: ValueType,
-        for<'a> T::ScalarRef<'a>: AsRef<[u8]>,
+        for<'a> F: FnMut(&T::ScalarRef<'a>) -> Result<u64>,
     {
         if IS_FIRST {
             for (x, val) in T::iter_column(col).zip(self.values.iter_mut()) {
-                *val = x.as_ref().agg_hash();
+                *val = agg_hash(&x)?;
             }
         } else {
             for (x, val) in T::iter_column(col).zip(self.values.iter_mut()) {
-                *val = (*val).wrapping_mul(NULL_HASH_VAL) ^ x.as_ref().agg_hash();
+                *val = (*val).wrapping_mul(NULL_HASH_VAL) ^ agg_hash(&x)?;
             }
         }
+        Ok(())
     }
 }
 
@@ -359,11 +371,14 @@ where I: Index
         self.visit_binary(column)
     }
 
-    fn visit_bitmap(&mut self, column: crate::types::BinaryColumn) -> Result<()> {
-        let column = normalize_bitmap_column(&column);
+    fn visit_bitmap(&mut self, column: BitmapColumn) -> Result<()> {
+        let mut bytes = Vec::new();
         self.visit_indices(|i| {
-            let value = column.as_ref().index(i.to_usize()).unwrap();
-            value.agg_hash()
+            let value = column.index(i.to_usize()).unwrap();
+            value.serialize_into(&mut bytes).unwrap();
+            let hash = bytes.agg_hash();
+            bytes.clear();
+            hash
         })
     }
 
@@ -429,8 +444,8 @@ where I: Index
 impl<const IS_FIRST: bool, I> IndexHashVisitor<'_, '_, IS_FIRST, I>
 where I: Index
 {
-    fn visit_indices<F>(&mut self, do_hash: F) -> Result<()>
-    where F: Fn(&I) -> u64 {
+    fn visit_indices<F>(&mut self, mut do_hash: F) -> Result<()>
+    where F: FnMut(&I) -> u64 {
         self.visit_indices_update(|i, val| {
             let hash = do_hash(i);
             if IS_FIRST {
@@ -441,8 +456,8 @@ where I: Index
         })
     }
 
-    fn visit_indices_update<F>(&mut self, update: F) -> Result<()>
-    where F: Fn(&I, &mut u64) {
+    fn visit_indices_update<F>(&mut self, mut update: F) -> Result<()>
+    where F: FnMut(&I, &mut u64) {
         for i in self.indices {
             let val = &mut self.target[i.to_usize()];
             update(i, val);

src/query/expression/src/aggregate/payload_flush.rs

@@ -239,7 +239,6 @@ impl Payload {
             DataType::Date => self.flush_type_column::<DateType>(col_offset, state),
             DataType::Binary => Column::Binary(self.flush_binary_column(col_offset, state)),
             DataType::String => Column::String(self.flush_string_column(col_offset, state)),
-            DataType::Bitmap => Column::Bitmap(self.flush_binary_column(col_offset, state)),
             DataType::Variant => Column::Variant(self.flush_binary_column(col_offset, state)),
             DataType::Geometry => Column::Geometry(self.flush_binary_column(col_offset, state)),
             DataType::Nullable(_) => unreachable!(),

src/query/expression/src/aggregate/payload_row.rs

@@ -20,11 +20,14 @@ use databend_common_io::prelude::bincode_serialize_into_buf;
 use super::probe_state::ProbeState;
 use super::row_ptr::RowLayout;
 use super::row_ptr::RowPtr;
+#[cfg(test)]
+use crate::types::bitmap::BitmapColumn;
 use crate::types::decimal::DecimalColumn;
 use crate::types::i256;
 use crate::types::AccessType;
 use crate::types::AnyType;
 use crate::types::BinaryType;
+use crate::types::BitmapType;
 use crate::types::BooleanType;
 use crate::types::DataType;
 use crate::types::DateType;
@@ -34,7 +37,6 @@ use crate::types::NumberColumn;
 use crate::types::NumberType;
 use crate::types::StringType;
 use crate::types::TimestampType;
-use crate::utils::bitmap::normalize_bitmap_column;
 use crate::with_decimal_mapped_type;
 use crate::with_number_mapped_type;
 use crate::Column;
@@ -122,11 +124,12 @@ pub(super) unsafe fn serialize_column_to_rowformat(
             }
         }
         Column::Bitmap(v) => {
-            let normalized = normalize_bitmap_column(v);
-            let col = normalized.as_ref();
+            let mut bytes = Vec::new();
             for &index in select_vector {
-                let data = arena.alloc_slice_copy(unsafe { col.index_unchecked(index) });
-                address[index].write_bytes(offset, data);
+                let map = unsafe { v.index_unchecked(index) };
+                map.serialize_into(bytes.as_mut_slice()).unwrap();
+                address[index].write_bytes(offset, &bytes);
+                bytes.clear();
             }
         }
         Column::Binary(v) | Column::Variant(v) | Column::Geometry(v) => {
@@ -287,15 +290,17 @@ impl ProbeState {
                 },
             ),
             Column::Bitmap(v) => {
-                let normalized = normalize_bitmap_column(v);
-                let col = normalized.as_ref();
+                let mut bytes = Vec::new();
                 self.row_match_column_generic(
                     validity,
                     validity_offset,
                     (count, no_match_count),
                     |idx, row_ptr| unsafe {
-                        let value = BinaryType::index_column_unchecked(col, idx);
-                        row_ptr.is_bytes_eq(col_offset, value)
+                        let value = BitmapType::index_column_unchecked(v, idx);
+                        value.serialize_into(&mut bytes).unwrap();
+                        let is_eq = row_ptr.is_bytes_eq(col_offset, &bytes);
+                        bytes.clear();
+                        is_eq
                     },
                 )
             }
@@ -346,10 +351,10 @@ impl ProbeState {
         validity: Option<&Bitmap>,
         validity_offset: usize,
         (count, mut no_match_count): (usize, usize),
-        compare_fn: F,
+        mut compare_fn: F,
     ) -> (usize, usize)
     where
-        F: Fn(usize, &RowPtr) -> bool,
+        F: FnMut(usize, &RowPtr) -> bool,
     {
         let mut temp = self.get_temp();
         temp.reserve(count);
@@ -451,7 +456,9 @@ mod tests {
         builder.commit_row();
         builder.put_slice(&legacy_bytes);
         builder.commit_row();
-        let column = Column::Bitmap(builder.build());
+        let column = Column::Bitmap(
+            BitmapColumn::from_binary(builder.build()).expect("valid bitmap column"),
+        );
 
         let arena = Bump::new();
         let row_size = rowformat_size(&DataType::Bitmap);