refactor aggregate unfreeze_from_rows() and fix UDAF fallbacking error (#940)

disable udaf falling-back by default because of possible bugs

Co-authored-by: zhangli20 <zhangli20@kuaishou.com>

Author: richox

native-engine/datafusion-ext-plans/src/agg/acc.rs

@@ -30,7 +30,7 @@ use datafusion::{
 };
 use datafusion_ext_commons::{
     assume,
-    io::{read_bytes_into_vec, read_bytes_slice, read_len, read_scalar, write_len, write_scalar},
+    io::{read_bytes_slice, read_len, read_scalar, write_len, write_scalar},
     unchecked,
 };
 use smallvec::SmallVec;
@@ -50,7 +50,7 @@ pub trait AccColumn: Send {
     fn num_records(&self) -> usize;
     fn mem_used(&self) -> usize;
     fn freeze_to_rows(&self, idx: IdxSelection<'_>, array: &mut [Vec<u8>]) -> Result<()>;
-    fn unfreeze_from_rows(&mut self, array: &[&[u8]], offsets: &mut [usize]) -> Result<()>;
+    fn unfreeze_from_rows(&mut self, cursors: &mut [Cursor<&[u8]>]) -> Result<()>;
     fn spill(&self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()>;
     fn unspill(&mut self, num_rows: usize, r: &mut SpillCompressedReader) -> Result<()>;
 
@@ -442,7 +442,6 @@ impl AccColumn for AccGenericColumn {
                         raw.set_len(new_len);
                     } else {
                         raw.truncate(new_len);
-                        raw.set_len(new_len);
                     }
                 }
                 valids.resize(len, false);
@@ -559,44 +558,38 @@ impl AccColumn for AccGenericColumn {
         Ok(())
     }
 
-    fn unfreeze_from_rows(&mut self, array: &[&[u8]], offsets: &mut [usize]) -> Result<()> {
-        let mut idx = self.num_records();
-        self.resize(idx + array.len());
+    fn unfreeze_from_rows(&mut self, cursors: &mut [Cursor<&[u8]>]) -> Result<()> {
+        assert_eq!(self.num_records(), 0, "expect empty AccColumn");
+        self.resize(cursors.len());
 
         match self {
             &mut AccGenericColumn::Prim {
                 ref mut raw,
                 ref mut valids,
                 prim_size,
             } => {
-                for (data, offset) in array.iter().zip(offsets) {
-                    let mut r = Cursor::new(data);
-                    r.set_position(*offset as u64);
-
-                    let valid = r.read_u8()?;
+                for (idx, cursor) in cursors.iter_mut().enumerate() {
+                    let valid = cursor.read_u8()?;
                     if valid == 1 {
-                        r.read_exact(&mut raw.as_raw_bytes_mut()[prim_size * idx..][..prim_size])?;
+                        cursor.read_exact(
+                            &mut raw.as_raw_bytes_mut()[prim_size * idx..][..prim_size],
+                        )?;
                         valids.set(idx, true);
                     } else {
                         valids.set(idx, false);
                     }
-                    *offset = r.position() as usize;
-                    idx += 1;
                 }
             }
             AccGenericColumn::Bytes {
                 items,
                 heap_mem_used,
             } => {
-                for (data, offset) in array.iter().zip(offsets) {
-                    let mut r = Cursor::new(data);
-                    r.set_position(*offset as u64);
-
-                    let len = read_len(&mut r)?;
+                for (idx, cursor) in cursors.iter_mut().enumerate() {
+                    let len = read_len(cursor)?;
                     if len > 0 {
                         let len = len - 1;
                         let bytes = AccBytes::from_vec({
-                            let vec: Vec<u8> = read_bytes_slice(&mut r, len)?.into();
+                            let vec: Vec<u8> = read_bytes_slice(cursor, len)?.into();
                             vec
                         });
                         if bytes.spilled() {
@@ -606,23 +599,16 @@ impl AccColumn for AccGenericColumn {
                     } else {
                         items[idx] = None;
                     }
-                    *offset = r.position() as usize;
-                    idx += 1;
                 }
             }
             AccGenericColumn::Scalar {
                 items,
                 dt,
                 heap_mem_used,
             } => {
-                for (data, offset) in array.iter().zip(offsets) {
-                    let mut r = Cursor::new(data);
-                    r.set_position(*offset as u64);
-
-                    items[idx] = read_scalar(&mut r, dt, true)?;
+                for (idx, cursor) in cursors.iter_mut().enumerate() {
+                    items[idx] = read_scalar(cursor, dt, true)?;
                     *heap_mem_used += items[idx].size() - size_of::<ScalarValue>();
-                    *offset = r.position() as usize;
-                    idx += 1;
                 }
             }
         }
@@ -678,23 +664,20 @@ impl AccColumn for AccGenericColumn {
     }
 
     fn unspill(&mut self, num_rows: usize, r: &mut SpillCompressedReader) -> Result<()> {
-        let idx = self.num_records();
-        self.resize(idx + num_rows);
+        assert_eq!(self.num_records(), 0, "expect empty AccColumn");
+        self.resize(num_rows);
 
         match self {
             &mut AccGenericColumn::Prim {
                 ref mut raw,
                 ref mut valids,
                 prim_size,
             } => {
-                let mut valid_buf = vec![];
-                let valid_len = (num_rows + 7) / 8;
-                read_bytes_into_vec(r, &mut valid_buf, valid_len)?;
-                let unfreezed_valids = BitVec::<u8>::from_vec(valid_buf);
-                valids.truncate(idx);
-                valids.extend_from_bitslice(unfreezed_valids.as_bitslice());
-
-                for i in idx..idx + num_rows {
+                let mut bits: BitVec<u8> = BitVec::repeat(false, num_rows);
+                r.read_exact(bits.as_raw_mut_slice())?;
+                valids.clear();
+                valids.extend_from_bitslice(bits.as_bitslice());
+                for i in 0..num_rows {
                     if valids[i] {
                         r.read_exact(&mut raw.as_raw_bytes_mut()[prim_size * i..][..prim_size])?;
                     }
@@ -704,7 +687,7 @@ impl AccColumn for AccGenericColumn {
                 items,
                 heap_mem_used,
             } => {
-                for i in idx..idx + num_rows {
+                for i in 0..num_rows {
                     let len = read_len(r)?;
                     if len > 0 {
                         let len = len - 1;
@@ -721,7 +704,7 @@ impl AccColumn for AccGenericColumn {
                 dt,
                 heap_mem_used,
             } => {
-                for i in idx..idx + num_rows {
+                for i in 0..num_rows {
                     items[i] = read_scalar(r, dt, true)?;
                     *heap_mem_used += items[i].size() - size_of::<ScalarValue>();
                 }

native-engine/datafusion-ext-plans/src/agg/agg_ctx.rs

@@ -14,11 +14,12 @@
 
 use std::{
     fmt::{Debug, Formatter},
+    io::Cursor,
     sync::Arc,
 };
 
 use arrow::{
-    array::{Array, ArrayRef, BinaryArray, RecordBatchOptions},
+    array::{ArrayRef, BinaryArray, RecordBatchOptions},
     datatypes::{DataType, Field, Fields, Schema, SchemaRef},
     record_batch::RecordBatch,
     row::{RowConverter, Rows, SortField},
@@ -39,6 +40,7 @@ use crate::{
     agg::{
         acc::AccTable,
         agg::{Agg, IdxSelection},
+        agg_hash_map::AggHashMapKey,
         spark_udaf_wrapper::{AccUDAFBufferRowsColumn, SparkUDAFMemTracker, SparkUDAFWrapper},
         AggExecMode, AggExpr, AggMode, GroupingExpr, AGG_BUF_COLUMN_NAME,
     },
@@ -244,7 +246,9 @@ impl AggContext {
         acc_table: &mut AccTable,
         acc_idx: IdxSelection,
     ) -> Result<()> {
-        let batch_selection = IdxSelection::Range(batch_start_idx, batch_end_idx);
+        // NOTE:
+        // arrow-ffi with sliced batch is buggy in older arrow-java, so we use unsliced
+        // batch with explicit offsets
 
         // partial update
         if self.need_partial_update {
@@ -263,6 +267,7 @@ impl AggContext {
                     input_arrays.push(vec![]);
                 }
             }
+            let batch_selection = IdxSelection::Range(batch_start_idx, batch_end_idx);
             self.partial_update(acc_table, acc_idx, &input_arrays, batch_selection)?;
         }
 
@@ -274,15 +279,21 @@ impl AggContext {
                 let partial_merged_array = as_binary_array(batch.columns().last().unwrap())?;
                 let array = partial_merged_array
                     .iter()
+                    .skip(batch_start_idx)
+                    .take(batch_end_idx - batch_start_idx)
                     .map(|bytes| bytes.unwrap())
                     .collect::<Vec<_>>();
-                let mut offsets = vec![0; partial_merged_array.len()];
+                let mut cursors = array
+                    .iter()
+                    .map(|bytes| Cursor::new(bytes.as_bytes()))
+                    .collect::<Vec<_>>();
 
                 for (agg_idx, _agg) in &self.need_partial_merge_aggs {
                     let acc_col = &mut merging_acc_table.cols_mut()[*agg_idx];
-                    acc_col.unfreeze_from_rows(&array, &mut offsets)?;
+                    acc_col.unfreeze_from_rows(&mut cursors)?;
                 }
             }
+            let batch_selection = IdxSelection::Range(0, batch_end_idx - batch_start_idx);
             self.partial_merge(acc_table, acc_idx, &mut merging_acc_table, batch_selection)?;
         }
         Ok(())
@@ -413,13 +424,6 @@ impl AggContext {
         Ok(vec)
     }
 
-    pub fn unfreeze_acc_table(&self, acc_table: &mut AccTable, data: &[&[u8]]) -> Result<()> {
-        let mut offsets = vec![0; data.len()];
-        for acc_col in acc_table.cols_mut() {
-            acc_col.unfreeze_from_rows(data, &mut offsets)?;
-        }
-        Ok(())
-    }
     pub async fn process_partial_skipped(
         &self,
         batch: RecordBatch,

native-engine/datafusion-ext-plans/src/agg/avg.rs

@@ -15,6 +15,7 @@
 use std::{
     any::Any,
     fmt::{Debug, Formatter},
+    io::Cursor,
     sync::Arc,
 };
 
@@ -209,9 +210,9 @@ impl AccColumn for AccAvgColumn {
         Ok(())
     }
 
-    fn unfreeze_from_rows(&mut self, array: &[&[u8]], offsets: &mut [usize]) -> Result<()> {
-        self.sum.unfreeze_from_rows(array, offsets)?;
-        self.count.unfreeze_from_rows(array, offsets)?;
+    fn unfreeze_from_rows(&mut self, cursors: &mut [Cursor<&[u8]>]) -> Result<()> {
+        self.sum.unfreeze_from_rows(cursors)?;
+        self.count.unfreeze_from_rows(cursors)?;
         Ok(())
     }
 

native-engine/datafusion-ext-plans/src/agg/bloom_filter.rs

@@ -259,21 +259,16 @@ impl AccColumn for AccBloomFilterColumn {
         Ok(())
     }
 
-    fn unfreeze_from_rows(&mut self, array: &[&[u8]], offsets: &mut [usize]) -> Result<()> {
-        let mut idx = self.num_records();
-        self.resize(idx + array.len());
-
-        for (data, offset) in array.iter().zip(offsets) {
-            let mut cursor = Cursor::new(*data);
-            cursor.set_position(*offset as u64);
-
-            if cursor.read_u8()? == 1 {
-                self.bloom_filters[idx] = Some(SparkBloomFilter::read_from(&mut cursor)?);
-            } else {
-                self.bloom_filters[idx] = None;
-            }
-            *offset = cursor.position() as usize;
-            idx += 1;
+    fn unfreeze_from_rows(&mut self, cursors: &mut [Cursor<&[u8]>]) -> Result<()> {
+        assert_eq!(self.num_records(), 0, "expect empty AccColumn");
+        for r in cursors {
+            self.bloom_filters.push({
+                if r.read_u8()? == 1 {
+                    Some(SparkBloomFilter::read_from(r)?)
+                } else {
+                    None
+                }
+            });
         }
         Ok(())
     }
@@ -293,15 +288,15 @@ impl AccColumn for AccBloomFilterColumn {
     }
 
     fn unspill(&mut self, num_rows: usize, r: &mut SpillCompressedReader) -> Result<()> {
-        let idx = self.num_records();
-        self.resize(idx + num_rows);
-
-        for i in idx..idx + num_rows {
-            if r.read_u8()? == 1 {
-                self.bloom_filters[i] = Some(SparkBloomFilter::read_from(r)?);
-            } else {
-                self.bloom_filters[i] = None;
-            }
+        assert_eq!(self.num_records(), 0, "expect empty AccColumn");
+        for _ in 0..num_rows {
+            self.bloom_filters.push({
+                if r.read_u8()? == 1 {
+                    Some(SparkBloomFilter::read_from(r)?)
+                } else {
+                    None
+                }
+            });
         }
         Ok(())
     }

native-engine/datafusion-ext-plans/src/agg/collect.rs

@@ -191,16 +191,12 @@ pub trait AccCollectionColumn: AccColumn + Send + Sync + 'static {
         Ok(())
     }
 
-    fn unfreeze_from_rows(&mut self, array: &[&[u8]], offsets: &mut [usize]) -> Result<()> {
-        let mut idx = self.num_records();
-        self.resize(idx + array.len());
-
-        for (raw, offset) in array.iter().zip(offsets) {
-            let mut cursor = Cursor::new(raw);
-            cursor.set_position(*offset as u64);
-            self.load_raw(idx, &mut cursor)?;
-            *offset = cursor.position() as usize;
-            idx += 1;
+    fn unfreeze_from_rows(&mut self, cursors: &mut [Cursor<&[u8]>]) -> Result<()> {
+        assert_eq!(self.num_records(), 0, "expect empty AccColumn");
+        self.resize(cursors.len());
+
+        for (idx, cursor) in cursors.iter_mut().enumerate() {
+            self.load_raw(idx, cursor)?;
         }
         Ok(())
     }
@@ -298,8 +294,8 @@ impl AccColumn for AccSetColumn {
         AccCollectionColumn::freeze_to_rows(self, idx, array)
     }
 
-    fn unfreeze_from_rows(&mut self, array: &[&[u8]], offsets: &mut [usize]) -> Result<()> {
-        AccCollectionColumn::unfreeze_from_rows(self, array, offsets)
+    fn unfreeze_from_rows(&mut self, cursors: &mut [Cursor<&[u8]>]) -> Result<()> {
+        AccCollectionColumn::unfreeze_from_rows(self, cursors)
     }
 
     fn spill(&self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()> {
@@ -312,12 +308,11 @@ impl AccColumn for AccSetColumn {
     }
 
     fn unspill(&mut self, num_rows: usize, r: &mut SpillCompressedReader) -> Result<()> {
-        let mut idx = self.num_records();
-        self.resize(idx + num_rows);
+        assert_eq!(self.num_records(), 0, "expect empty AccColumn");
+        self.resize(num_rows);
 
-        while idx < self.num_records() {
+        for idx in 0..num_rows {
             self.load_raw(idx, r)?;
-            idx += 1;
         }
         Ok(())
     }
@@ -411,8 +406,8 @@ impl AccColumn for AccListColumn {
         AccCollectionColumn::freeze_to_rows(self, idx, array)
     }
 
-    fn unfreeze_from_rows(&mut self, array: &[&[u8]], offsets: &mut [usize]) -> Result<()> {
-        AccCollectionColumn::unfreeze_from_rows(self, array, offsets)
+    fn unfreeze_from_rows(&mut self, cursors: &mut [Cursor<&[u8]>]) -> Result<()> {
+        AccCollectionColumn::unfreeze_from_rows(self, cursors)
     }
 
     fn spill(&self, idx: IdxSelection<'_>, w: &mut SpillCompressedWriter) -> Result<()> {
@@ -425,12 +420,10 @@ impl AccColumn for AccListColumn {
     }
 
     fn unspill(&mut self, num_rows: usize, r: &mut SpillCompressedReader) -> Result<()> {
-        let mut idx = self.num_records();
-        self.resize(idx + num_rows);
-
-        while idx < self.num_records() {
+        assert_eq!(self.num_records(), 0, "expect empty AccColumn");
+        self.resize(num_rows);
+        for idx in 0..num_rows {
             self.load_raw(idx, r)?;
-            idx += 1;
         }
         Ok(())
     }