2010YOUY01 commented on code in PR #25371: URL: https://github.com/apache/datafusion/pull/25371#discussion_r4037004388
########## datafusion/physical-plan/src/joins/logical_batch.rs: ########## @@ -0,0 +1,748 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`LogicalBatch`]: a logically contiguous batch stored as a sequence of +//! [`RecordBatch`]es. + +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, RecordBatch, UInt32Array, new_empty_array}; +use arrow::compute::{TakeOptions, concat, concat_batches, take}; +use arrow::datatypes::SchemaRef; +use datafusion_common::{Result, exec_datafusion_err, exec_err}; + +/// A logically contiguous batch backed by a sequence of [`RecordBatch`]es +/// that share one schema. Methods on this struct accept global row indices. +/// +/// # Example +/// ```text +/// +/// segment 0 (3 rows) segment 1 (2 rows) segment 2 (4 rows) +/// ┌───┬───┬───┐ ┌───┬───┐ ┌───┬───┬───┬───┐ +/// │ a │ b │ c │ │ d │ e │ │ f │ g │ h │ i │ +/// └───┴───┴───┘ └───┴───┘ └───┴───┴───┴───┘ +/// 0 1 2 3 4 5 6 7 8 ◀── global row index +/// ``` +/// +/// # Motivation +/// +/// Joins (e.g. Nested Loop Join) usually buffer all build-side input, and next concatenating +/// them into a contiguous batch, before the next step. It will 2X the memory usage +/// since fragmented batches and final contiguous batch exist at the same time. This +/// struct avoids concatenation step, and helps reduce memory usage by 2X. +/// +/// Avoiding memory concatenating overhead is not the motivation, since it's usually +/// fast and not a bottleneck in real workloads; at the same time single-batch abstraction +/// help simplify join logic. +/// +/// See issue for details: +/// - <https://github.com/apache/datafusion/issues/23076> +/// +/// # TODO +/// It's named 'logical batch' because it's possible to swap the physical layout +/// and keep the same interface for other usages. For example, segments are aligned +/// at the same size, so it achieves O(1) access speed. +#[derive(Debug, Clone)] +pub(crate) struct LogicalBatch { + schema: SchemaRef, + /// The underlying batches, in row order. Empty batches are dropped on + /// construction, so every segment holds at least one row. + segments: Vec<RecordBatch>, + /// `offsets[i]` is the global index of the first row of `segments[i]`; + /// `offsets[segments.len()]` is the total number of rows. + offsets: Vec<usize>, +} + +impl LogicalBatch { + /// Creates a logical batch from `batches`, which must all have `schema`. + /// + /// # Errors + /// + /// Returns an execution error if a batch has a different schema or the + /// total row count overflows. + pub(crate) fn new(schema: SchemaRef, batches: Vec<RecordBatch>) -> Result<Self> { + if batches.iter().any(|batch| batch.schema() != schema) { + return exec_err!("LogicalBatch input batches must have the same schema"); + } + let segments: Vec<RecordBatch> = batches + .into_iter() + .filter(|batch| batch.num_rows() > 0) + .collect(); + let mut offsets = Vec::with_capacity(segments.len() + 1); + let mut num_rows: usize = 0; + offsets.push(num_rows); + for segment in &segments { + num_rows = num_rows.checked_add(segment.num_rows()).ok_or_else(|| { + exec_datafusion_err!("LogicalBatch total row count exceeds usize::MAX") + })?; + offsets.push(num_rows); + } + Ok(Self { + schema, + segments, + offsets, + }) + } + + /// Creates a logical batch with no rows. + pub(crate) fn new_empty(schema: SchemaRef) -> Self { + Self { + schema, + segments: vec![], + offsets: vec![0], + } + } + + pub(crate) fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + /// Total number of rows across all segments. + pub(crate) fn num_rows(&self) -> usize { + // `offsets` always holds at least the leading 0 + self.offsets[self.offsets.len() - 1] + } + + /// Returns the underlying batches, in row order. + pub(crate) fn into_batches(self) -> Vec<RecordBatch> { + self.segments + } + + /// The row at global index `index`. + /// + /// # Errors + /// + /// Returns an execution error if `index >= num_rows()`. + pub(crate) fn row(&self, index: usize) -> Result<BatchRow<'_>> { + if index >= self.num_rows() { + return exec_err!( + "row index {index} out of bounds for a batch of {} rows", + self.num_rows() + ); + } + let segment = self.segment_of(index, 0); + Ok(BatchRow { + batch: &self.segments[segment], + index: index - self.offsets[segment], + }) + } + + /// Resolves global row indices for [`Self::take_column`]. + /// + /// This plays the role of the index array given to [`take`] on a plain + /// batch: resolve the indices once, then gather as many columns as + /// needed with them. Rows may repeat or appear in any order; gathers + /// preserve that order. + /// + /// # Errors + /// + /// Returns an execution error if any index is `>= num_rows()` or its + /// index within the segment cannot be represented as a `u32`. + pub(crate) fn row_indices( + &self, + rows: impl IntoIterator<Item = usize>, + ) -> Result<RowIndices> { + let rows = rows.into_iter(); + let mut groups = Vec::new(); + let mut indices = Vec::with_capacity(rows.size_hint().0); + // Consecutive rows usually sit in the same segment, so retry the last + // segment before searching + let mut segment = 0; + for row in rows { + if row >= self.num_rows() { + return exec_err!( + "row index {row} out of bounds for a batch of {} rows", + self.num_rows() + ); + } + let next_segment = self.segment_of(row, segment); + // Start a new group whenever the segment changes. Keeping separate + // groups for repeated visits to a segment preserves output order. + if next_segment != segment && !indices.is_empty() { + groups.push(Single { + segment, + indices: UInt32Array::from(std::mem::take(&mut indices)), + }); + } + segment = next_segment; + let index = u32::try_from(row - self.offsets[segment]).map_err(|_| { + exec_datafusion_err!( + "row index {row} within segment {segment} exceeds u32::MAX" + ) + })?; + indices.push(index); + } + + if !indices.is_empty() { + groups.push(Single { + segment, + indices: UInt32Array::from(indices), + }); + } + + Ok(RowIndices(match groups.len() { + 0 => IndicesVariant::Empty, + 1 => IndicesVariant::Single(groups.pop().unwrap()), + _ => IndicesVariant::Multi(groups), + })) + } + + /// Gathers the rows selected by `indices` from column `column`, like + /// [`take`] on the column of a plain batch. + /// + /// # Errors + /// + /// Returns an execution error if the column or resolved row indices + /// are out of bounds for this batch. + pub(crate) fn take_column( + &self, + column: usize, + indices: &RowIndices, + ) -> Result<ArrayRef> { + let field = self.schema.fields().get(column).ok_or_else(|| { + exec_datafusion_err!( + "column index {column} out of bounds for a batch of {} columns", + self.schema.fields().len() + ) + })?; + let take_single = |single: &Single| -> Result<ArrayRef> { + let batch = self.segments.get(single.segment).ok_or_else(|| { + exec_datafusion_err!( + "segment index {} out of bounds for a batch of {} segments", + single.segment, + self.segments.len() + ) + })?; + let values = batch.columns().get(column).ok_or_else(|| { + exec_datafusion_err!( + "column index {column} out of bounds for a batch of {} columns", + batch.num_columns() + ) + })?; + take( Review Comment: Yes, even if it's tuned for different workloads, `batch_size` is always reasonably large for better vectorization. -- 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] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
