eliminate type conversion.

apache · Aug 30, 2024 · 11bf690 · 11bf690
1 parent 787de57
commit 11bf690
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 21 deletions.
diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs
@@ -44,10 +44,10 @@ use datafusion_expr_common::accumulator::Accumulator;
 use datafusion_expr_common::groups_accumulator::{EmitTo, GroupsAccumulator};
 
 pub const MAX_PREALLOC_BLOCK_SIZE: usize = 8192;
-const FLAT_GROUP_INDEX_ID_MASK: u64 = 0;
-const FLAT_GROUP_INDEX_OFFSET_MASK: u64 = u64::MAX;
-const BLOCKED_GROUP_INDEX_ID_MASK: u64 = 0xffffffff00000000;
-const BLOCKED_GROUP_INDEX_OFFSET_MASK: u64 = 0x00000000ffffffff;
+const FLAT_GROUP_INDEX_ID_MASK: usize = 0;
+const FLAT_GROUP_INDEX_OFFSET_MASK: usize = usize::MAX;
+const BLOCKED_GROUP_INDEX_ID_MASK: usize = 0xffffffff00000000;
+const BLOCKED_GROUP_INDEX_OFFSET_MASK: usize = 0x00000000ffffffff;
 
 /// An adapter that implements [`GroupsAccumulator`] for any [`Accumulator`]
 ///
@@ -454,38 +454,38 @@ impl EmitToExt for EmitTo {
 ///
 #[derive(Debug, Clone, Copy)]
 pub struct BlockedGroupIndex {
-    pub block_id: u32,
-    pub block_offset: u64,
+    pub block_id: usize,
+    pub block_offset: usize,
 }
 
 impl BlockedGroupIndex {
     #[inline]
-    pub fn new_from_parts(block_id: u32, block_offset: u64) -> Self {
+    pub fn new_from_parts(block_id: usize, block_offset: usize) -> Self {
         Self {
-            block_id,
-            block_offset,
+            block_id: block_id as usize,
+            block_offset: block_offset as usize,
         }
     }
 
     #[inline]
     pub fn block_id(&self) -> usize {
-        self.block_id as usize
+        self.block_id
     }
 
     #[inline]
     pub fn block_offset(&self) -> usize {
-        self.block_offset as usize
+        self.block_offset
     }
 
     #[inline]
     pub fn as_packed_index(&self) -> usize {
-        (((self.block_id as u64) << 32) | self.block_offset) as usize
+        (self.block_id << 32) | self.block_offset
     }
 }
 
 pub struct BlockedGroupIndexBuilder {
-    block_offset_mask: u64,
-    id_shift: u64,
+    block_offset_mask: usize,
+    id_shift: usize,
 }
 
 impl BlockedGroupIndexBuilder {
@@ -505,8 +505,7 @@ impl BlockedGroupIndexBuilder {
 
     #[inline]
     pub fn build(&self, packed_index: usize) -> BlockedGroupIndex {
-        let packed_index = packed_index as u64;
-        let block_id = (packed_index >> self.id_shift) as u32;
+        let block_id = packed_index >> self.id_shift;
         let block_offset = packed_index & self.block_offset_mask;
 
         BlockedGroupIndex {

diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/accumulate.rs
@@ -919,8 +919,8 @@ mod test {
                     let block_id = *idx / self.block_size;
                     let block_offset = *idx % self.block_size;
                     BlockedGroupIndex::new_from_parts(
-                        block_id as u32,
-                        block_offset as u64,
+                        block_id,
+                        block_offset,
                     )
                     .as_packed_index()
                 })

diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs
@@ -180,8 +180,8 @@ impl GroupValues for GroupValuesRows {
                     cur_blk.push(group_rows.row(row));
 
                     let blocked_index = BlockedGroupIndex::new_from_parts(
-                        blk_id as u32,
-                        blk_offset as u64,
+                        blk_id,
+                        blk_offset,
                     );
                     let group_idx = blocked_index.as_packed_index();
 
@@ -294,7 +294,7 @@ impl GroupValues for GroupValuesRows {
                             // Group index was >= n, shift value down
                             Some(new_blk_id) => {
                                 let new_group_idx = BlockedGroupIndex::new_from_parts(
-                                    new_blk_id as u32,
+                                    new_blk_id,
                                     old_blk_idx.block_offset,
                                 );
                                 bucket.as_mut().1 = new_group_idx.as_packed_index();