Rich-T-kid commented on code in PR #11174:
URL: https://github.com/apache/arrow-rs/pull/11174#discussion_r4202721363
##########
arrow-select/src/coalesce.rs:
##########
@@ -2893,4 +3084,5 @@ mod tests {
"draining must release the accounted memory"
);
}
+
Review Comment:
remove
##########
arrow-select/src/coalesce.rs:
##########
@@ -809,6 +974,32 @@ trait InProgressArray: std::fmt::Debug + Send + Sync {
}
}
+ /// Copy rows at positions `indices` from `source` into the in-progress
+ /// array. Fallback path: builds a `UInt32Array` from `indices` without
+ /// cloning the slice, then calls [`take`]. Types with a fast path
+ /// (primitive, byte-view) override this to scatter directly from `source`.
+ fn copy_rows_by_indices(
+ &mut self,
+ source: &ArrayRef,
+ indices: &[u32],
+ ) -> Result<(), ArrowError> {
+ if indices.is_empty() {
+ return Ok(());
+ }
+ // Avoid cloning `indices` into a fresh Vec: wrap the existing slice in
+ // a `ScalarBuffer` (zero-copy via `Buffer::from_slice_ref` when the
+ // slice outlives the kernel call, which it does here).
+ use arrow_buffer::{Buffer, ScalarBuffer};
+ let buf = Buffer::from_slice_ref(indices);
Review Comment:
remove
##########
arrow-select/src/coalesce.rs:
##########
@@ -721,6 +723,169 @@ impl BatchCoalescer {
}
}
+/// Row buffer used by [`PartitionCoalescer`].
+#[derive(Debug)]
+struct BufferCoalescer {
+ schema: SchemaRef,
+ target_batch_size: usize,
+ in_progress_arrays: Vec<Box<dyn InProgressArray>>,
+ buffered_rows: usize,
+ completed: VecDeque<RecordBatch>,
+}
+
+impl BufferCoalescer {
+ fn new(schema: SchemaRef, target_batch_size: usize) -> Self {
+ let in_progress_arrays = schema
+ .fields()
+ .iter()
+ .map(|field| create_in_progress_array(field.data_type(),
target_batch_size))
+ .collect();
+ Self {
+ schema,
+ target_batch_size,
+ in_progress_arrays,
+ buffered_rows: 0,
+ completed: VecDeque::with_capacity(1),
+ }
+ }
+
+ fn push_rows(
+ &mut self,
+ batch: &RecordBatch,
+ rows: &[u32],
+ ) -> Result<(), ArrowError> {
+ if rows.is_empty() {
+ return Ok(());
+ }
+ let columns = batch.columns();
+ if columns.len() != self.in_progress_arrays.len() {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "Batch has {} columns but BufferCoalescer expects {}",
+ columns.len(),
+ self.in_progress_arrays.len()
+ )));
+ }
+
+ let mut row_offset = 0;
+ let mut remaining = rows.len();
+ while remaining > self.target_batch_size - self.buffered_rows {
+ let room = self.target_batch_size - self.buffered_rows;
+ let chunk = &rows[row_offset..row_offset + room];
+ for (in_progress, column) in
self.in_progress_arrays.iter_mut().zip(columns) {
+ in_progress.copy_rows_by_indices(column, chunk)?;
+ }
+ self.buffered_rows += room;
+ row_offset += room;
+ remaining -= room;
+ self.flush()?;
+ }
+ if remaining > 0 {
+ let chunk = &rows[row_offset..row_offset + remaining];
+ for (in_progress, column) in
self.in_progress_arrays.iter_mut().zip(columns) {
+ in_progress.copy_rows_by_indices(column, chunk)?;
+ }
+ self.buffered_rows += remaining;
+ }
+ Ok(())
+ }
+
+ fn next_completed_batch(&mut self) -> Option<RecordBatch> {
+ self.completed.pop_front()
+ }
+
+ fn finish(&mut self) -> Result<(), ArrowError> {
+ self.flush()
+ }
+
+ fn flush(&mut self) -> Result<(), ArrowError> {
+ if self.buffered_rows == 0 {
+ return Ok(());
+ }
+ let arrays = self
+ .in_progress_arrays
+ .iter_mut()
+ .map(|in_progress| in_progress.finish())
+ .collect::<Result<Vec<_>, ArrowError>>()?;
+ let batch = unsafe {
+ RecordBatch::new_unchecked(Arc::clone(&self.schema), arrays,
self.buffered_rows)
+ };
+ self.buffered_rows = 0;
+ self.completed.push_back(batch);
+ Ok(())
+ }
+}
+
+/// Single-partition coalescer: filters an input batch down to rows whose
+/// partition-id equals `my_partition`, buffering them into target-sized
+/// batches. Each output partition in a repartition operator owns one of these.
+#[derive(Debug)]
+pub struct PartitionCoalescer {
+ my_partition: u64,
+ /// Reused scratch of row indices where `pid == my_partition`.
+ bucket: Vec<u32>,
+ coalescer: BufferCoalescer,
+}
+
+impl PartitionCoalescer {
+ /// Create a coalescer that buffers rows tagged with `my_partition`.
+ pub fn new(schema: SchemaRef, target_batch_size: usize, my_partition:
usize) -> Self {
+ Self {
+ my_partition: my_partition as u64,
+ bucket: Vec::new(),
+ coalescer: BufferCoalescer::new(schema, target_batch_size),
+ }
+ }
+
+ /// Append rows from `batch` whose partition-id equals `my_partition`.
+ pub fn push(
+ &mut self,
+ batch: &RecordBatch,
+ partition_ids: &UInt64Array,
+ ) -> Result<(), ArrowError> {
+ let pids = partition_ids.values();
+ self.bucket.clear();
+ self.bucket.reserve(pids.len());
+ let my = self.my_partition;
+ unsafe {
+ let base = self.bucket.as_mut_ptr();
+ let mut out = base;
+ // Branchless conditional store.
+ for (row, &pid) in pids.iter().enumerate() {
+ *out = row as u32;
+ out = out.add((pid == my) as usize);
+ }
+ self.bucket.set_len(out.offset_from(base) as usize);
+ }
+ if !self.bucket.is_empty() {
+ self.coalescer.push_rows(batch, &self.bucket)?;
+ }
+ Ok(())
+ }
+
+ /// Append a pre-bucketed slice of row indices.
Review Comment:
better doc comment here
##########
arrow-select/src/coalesce.rs:
##########
@@ -721,6 +723,169 @@ impl BatchCoalescer {
}
}
+/// Row buffer used by [`PartitionCoalescer`].
+#[derive(Debug)]
+struct BufferCoalescer {
+ schema: SchemaRef,
+ target_batch_size: usize,
+ in_progress_arrays: Vec<Box<dyn InProgressArray>>,
+ buffered_rows: usize,
+ completed: VecDeque<RecordBatch>,
+}
+
+impl BufferCoalescer {
+ fn new(schema: SchemaRef, target_batch_size: usize) -> Self {
+ let in_progress_arrays = schema
+ .fields()
+ .iter()
+ .map(|field| create_in_progress_array(field.data_type(),
target_batch_size))
+ .collect();
+ Self {
+ schema,
+ target_batch_size,
+ in_progress_arrays,
+ buffered_rows: 0,
+ completed: VecDeque::with_capacity(1),
+ }
+ }
+
+ fn push_rows(
+ &mut self,
+ batch: &RecordBatch,
+ rows: &[u32],
+ ) -> Result<(), ArrowError> {
+ if rows.is_empty() {
+ return Ok(());
+ }
+ let columns = batch.columns();
+ if columns.len() != self.in_progress_arrays.len() {
+ return Err(ArrowError::InvalidArgumentError(format!(
+ "Batch has {} columns but BufferCoalescer expects {}",
+ columns.len(),
+ self.in_progress_arrays.len()
+ )));
+ }
+
+ let mut row_offset = 0;
+ let mut remaining = rows.len();
+ while remaining > self.target_batch_size - self.buffered_rows {
+ let room = self.target_batch_size - self.buffered_rows;
+ let chunk = &rows[row_offset..row_offset + room];
+ for (in_progress, column) in
self.in_progress_arrays.iter_mut().zip(columns) {
+ in_progress.copy_rows_by_indices(column, chunk)?;
+ }
+ self.buffered_rows += room;
+ row_offset += room;
+ remaining -= room;
+ self.flush()?;
+ }
+ if remaining > 0 {
+ let chunk = &rows[row_offset..row_offset + remaining];
+ for (in_progress, column) in
self.in_progress_arrays.iter_mut().zip(columns) {
+ in_progress.copy_rows_by_indices(column, chunk)?;
+ }
+ self.buffered_rows += remaining;
+ }
+ Ok(())
+ }
+
+ fn next_completed_batch(&mut self) -> Option<RecordBatch> {
+ self.completed.pop_front()
+ }
+
+ fn finish(&mut self) -> Result<(), ArrowError> {
+ self.flush()
+ }
+
+ fn flush(&mut self) -> Result<(), ArrowError> {
+ if self.buffered_rows == 0 {
+ return Ok(());
+ }
+ let arrays = self
+ .in_progress_arrays
+ .iter_mut()
+ .map(|in_progress| in_progress.finish())
+ .collect::<Result<Vec<_>, ArrowError>>()?;
+ let batch = unsafe {
+ RecordBatch::new_unchecked(Arc::clone(&self.schema), arrays,
self.buffered_rows)
+ };
+ self.buffered_rows = 0;
+ self.completed.push_back(batch);
+ Ok(())
+ }
+}
+
+/// Single-partition coalescer: filters an input batch down to rows whose
+/// partition-id equals `my_partition`, buffering them into target-sized
+/// batches. Each output partition in a repartition operator owns one of these.
+#[derive(Debug)]
+pub struct PartitionCoalescer {
+ my_partition: u64,
+ /// Reused scratch of row indices where `pid == my_partition`.
+ bucket: Vec<u32>,
+ coalescer: BufferCoalescer,
+}
+
+impl PartitionCoalescer {
+ /// Create a coalescer that buffers rows tagged with `my_partition`.
+ pub fn new(schema: SchemaRef, target_batch_size: usize, my_partition:
usize) -> Self {
+ Self {
+ my_partition: my_partition as u64,
+ bucket: Vec::new(),
+ coalescer: BufferCoalescer::new(schema, target_batch_size),
+ }
+ }
+
+ /// Append rows from `batch` whose partition-id equals `my_partition`.
Review Comment:
better doc comment here. maybe an example would nice
--
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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
