optimizations_roaring_memebership_check

Author: B Vadlamani

datafusion/physical-plan/src/joins/hash_join/exec.rs

@@ -31,7 +31,8 @@ use crate::filter_pushdown::{
     ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase,
     FilterPushdownPropagation,
 };
-use crate::joins::Map;
+use crate::joins::{Map, RoaringMapData};
+use roaring::RoaringBitmap;
 use crate::joins::array_map::ArrayMap;
 use crate::joins::hash_join::inlist_builder::build_struct_inlist_values;
 use crate::joins::hash_join::shared_bounds::{
@@ -1994,7 +1995,7 @@ async fn collect_left_input(
             let batch = concat_batches(&schema, &batches)?;
             let left_values = evaluate_expressions_to_arrays(&on_left, &batch)?;
             let key_col = &left_values[0];
-            let bitmap = match key_col.data_type() {
+            let bitmap: RoaringBitmap = match key_col.data_type() {
                 DataType::UInt32 => key_col
                     .as_any()
                     .downcast_ref::<UInt32Array>()
@@ -2013,7 +2014,13 @@ async fn collect_left_input(
                 _ => return internal_err!("unsupported data type to build bitmap"),
             };
             roaring_map_created_count.add(1);
-            (Map::RoaringMap(bitmap), batch, left_values)
+            let min = bitmap.min().unwrap_or(u32::MAX);
+            let max = bitmap.max().unwrap_or(u32::MIN);
+            (
+                Map::RoaringMap(RoaringMapData { bitmap, min, max }),
+                batch,
+                left_values,
+            )
         } else if let Some((array_map, batch, left_value)) = try_create_array_map(
             &bounds,
             &schema,

datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs

@@ -336,7 +336,7 @@ impl PhysicalExpr for HashTableLookupExpr {
                 let array = map.contain_keys(&join_keys)?;
                 Ok(ColumnarValue::Array(Arc::new(array)))
             }
-            Map::RoaringMap(bitmap) => {
+            Map::RoaringMap(data) => {
                 // For roaring bitmap, check membership directly
                 // Roaring only supports u32 values, so we handle Int32 and UInt32 types
                 if join_keys.len() != 1 {
@@ -346,6 +346,12 @@ impl PhysicalExpr for HashTableLookupExpr {
                     );
                 }
                 let key_col = &join_keys[0];
+                // Bounds were cached at build time so out-of-range probes
+                // short-circuit before the binary-search inside
+                // `RoaringBitmap::contains`.
+                let min = data.min;
+                let max = data.max;
+                let bitmap = &data.bitmap;
                 let contains: BooleanArray = match key_col.data_type() {
                     DataType::Int32 => {
                         let arr = key_col
@@ -356,7 +362,8 @@ impl PhysicalExpr for HashTableLookupExpr {
                             if arr.is_null(i) {
                                 return false;
                             }
-                            bitmap.contains(arr.value(i) as u32)
+                            let v = arr.value(i) as u32;
+                            v >= min && v <= max && bitmap.contains(v)
                         });
                         BooleanArray::new(buffer, None)
                     }
@@ -369,7 +376,8 @@ impl PhysicalExpr for HashTableLookupExpr {
                             if arr.is_null(i) {
                                 return false;
                             }
-                            bitmap.contains(arr.value(i))
+                            let v = arr.value(i);
+                            v >= min && v <= max && bitmap.contains(v)
                         });
                         BooleanArray::new(buffer, None)
                     }

datafusion/physical-plan/src/joins/hash_join/stream.rs

@@ -732,17 +732,26 @@ impl HashJoinStream {
                     next_offset,
                 )
             }
-            Map::RoaringMap(bitmap) => {
+            Map::RoaringMap(data) => {
                 let key_col = &state.values[0];
                 let is_semi = matches!(self.join_type, JoinType::RightSemi);
+                // Bounds were cached at build time so out-of-range probes
+                // short-circuit before the binary-search inside
+                // `RoaringBitmap::contains`.
+                let min = data.min;
+                let max = data.max;
+                let bitmap = &data.bitmap;
                 let mask: BooleanArray = match key_col.data_type() {
                     DataType::Int32 => {
                         let arr = key_col.as_any().downcast_ref::<Int32Array>().unwrap();
                         let buffer = BooleanBuffer::collect_bool(arr.len(), |i| {
                             if arr.is_null(i) {
                                 return !is_semi;
                             }
-                            bitmap.contains(arr.value(i) as u32) == is_semi
+                            let v = arr.value(i) as u32;
+                            let in_set =
+                                v >= min && v <= max && bitmap.contains(v);
+                            in_set == is_semi
                         });
                         BooleanArray::new(buffer, None)
                     }
@@ -752,7 +761,10 @@ impl HashJoinStream {
                             if arr.is_null(i) {
                                 return !is_semi;
                             }
-                            bitmap.contains(arr.value(i)) == is_semi
+                            let v = arr.value(i);
+                            let in_set =
+                                v >= min && v <= max && bitmap.contains(v);
+                            in_set == is_semi
                         });
                         BooleanArray::new(buffer, None)
                     }

datafusion/physical-plan/src/joins/mod.rs

@@ -51,11 +51,20 @@ pub mod join_hash_map;
 use array_map::ArrayMap;
 use utils::JoinHashMapType;
 
+pub struct RoaringMapData {
+    pub bitmap: RoaringBitmap,
+    /// Smallest key in `bitmap`. Cached at build time so probe loops can
+    /// short-circuit out-of-range keys without invoking `RoaringBitmap::contains`.
+    pub min: u32,
+    /// Largest key in `bitmap`. See [`Self::min`].
+    pub max: u32,
+}
+
 pub enum Map {
     HashMap(Box<dyn JoinHashMapType>),
     ArrayMap(ArrayMap),
     // optimized path for single int join keys
-    RoaringMap(RoaringBitmap),
+    RoaringMap(RoaringMapData),
 }
 
 impl Map {
@@ -64,7 +73,7 @@ impl Map {
         match self {
             Map::HashMap(map) => map.len(),
             Map::ArrayMap(array_map) => array_map.num_of_distinct_key(),
-            Map::RoaringMap(bitmap) => bitmap.len() as usize,
+            Map::RoaringMap(data) => data.bitmap.len() as usize,
         }
     }