Extract limit code

Author: alamb

parquet/src/arrow/arrow_reader/mod.rs

@@ -700,9 +700,10 @@ impl<T: ChunkReader + 'static> ParquetRecordBatchReaderBuilder<T> {
 
         let array_reader = build_array_reader(self.fields.as_deref(), &self.projection, &reader)?;
         let plan = read_plan_builder
-            .with_num_rows(reader.num_rows())
+            .limited(reader.num_rows())
             .with_offset(self.offset)
             .with_limit(self.limit)
+            .build_limited()
             .build()?;
 
         Ok(ParquetRecordBatchReader::new(array_reader, plan))
@@ -899,7 +900,6 @@ impl ParquetRecordBatchReader {
             build_array_reader(levels.levels.as_ref(), &ProjectionMask::all(), row_groups)?;
 
         let read_plan = ReadPlanBuilder::new(batch_size)
-            .with_num_rows(row_groups.num_rows())
             .with_selection(selection)
             .build()?;
 

parquet/src/arrow/arrow_reader/read_plan.rs

@@ -22,24 +22,20 @@ use crate::arrow::array_reader::ArrayReader;
 use crate::arrow::arrow_reader::{
     ArrowPredicate, ParquetRecordBatchReader, RowSelection, RowSelector,
 };
-use crate::errors::ParquetError;
+use crate::errors::{ParquetError, Result};
 use arrow_array::Array;
 use arrow_select::filter::prep_null_mask_filter;
 use std::collections::VecDeque;
 
 /// A builder for [`ReadPlan`]
+///
+/// See also [`LimitedReadPlanBuilder`], for applying limits and offsets
 #[derive(Clone)]
 pub(crate) struct ReadPlanBuilder {
     batch_size: usize,
     /// Current to apply, includes all filters
     selection: Option<RowSelection>,
     // TODO: Cached result of evaluating some columns with the RowSelection
-    /// Total number of rows in the row group before the selection
-    num_rows: Option<usize>,
-    /// Rows to skip before returning any rows
-    offset: Option<usize>,
-    /// Limit on the number of rows to return
-    limit: Option<usize>,
 }
 
 impl ReadPlanBuilder {
@@ -50,9 +46,6 @@ impl ReadPlanBuilder {
         Self {
             batch_size,
             selection: None,
-            num_rows: None,
-            offset: None,
-            limit: None,
         }
     }
 
@@ -67,22 +60,14 @@ impl ReadPlanBuilder {
         self.selection.as_ref()
     }
 
-    /// set the number of rows in the row group
-    pub(crate) fn with_num_rows(mut self, num_rows: usize) -> Self {
-        self.num_rows = Some(num_rows);
-        self
-    }
-
-    /// Set the offset to the given value
-    pub(crate) fn with_offset(mut self, offset: Option<usize>) -> Self {
-        self.offset = offset;
-        self
-    }
-
-    /// Set the limit to the given value
-    pub(crate) fn with_limit(mut self, limit: Option<usize>) -> Self {
-        self.limit = limit;
-        self
+    /// Specify the number of rows in the row group before filtering
+    /// returning a [`LimitedReadPlanBuilder`] that can apply
+    /// offset and limit.
+    ///
+    /// Call [`LimitedReadPlanBuilder::build_limited`] to apply the limits to this
+    /// selection.
+    pub(crate) fn limited(self, row_count: usize) -> LimitedReadPlanBuilder {
+        LimitedReadPlanBuilder::new(self, row_count)
     }
 
     /// Returns true if the current plan selects any rows
@@ -113,7 +98,7 @@ impl ReadPlanBuilder {
         mut self,
         array_reader: Box<dyn ArrayReader>,
         predicate: &mut dyn ArrowPredicate,
-    ) -> crate::errors::Result<Self> {
+    ) -> Result<Self> {
         let reader = ParquetRecordBatchReader::new(array_reader, self.clone().build()?);
         let mut filters = vec![];
         for maybe_batch in reader {
@@ -149,47 +134,9 @@ impl ReadPlanBuilder {
         }
         let Self {
             batch_size,
-            mut selection,
-            num_rows,
-            offset,
-            limit,
+            selection,
         } = self;
 
-        // TODO make this nicer somehow (maybe a different builder returned with Rowcount)
-        if offset.is_some() || limit.is_some() {
-            let Some(row_count) = num_rows else {
-                return Err(general_err!(
-                    "Internal ReadPlanBuilder::build() has limit/offset but called without a row count"
-                ));
-            };
-
-            // If an offset is defined, apply it to the `selection`
-            if let Some(offset) = offset {
-                selection = Some(match row_count.checked_sub(offset) {
-                    None => RowSelection::from(vec![]),
-                    Some(remaining) => selection
-                        .map(|selection| selection.offset(offset))
-                        .unwrap_or_else(|| {
-                            RowSelection::from(vec![
-                                RowSelector::skip(offset),
-                                RowSelector::select(remaining),
-                            ])
-                        }),
-                });
-            }
-
-            // If a limit is defined, apply it to the final `selection`
-            if let Some(limit) = limit {
-                selection = Some(
-                    selection
-                        .map(|selection| selection.limit(limit))
-                        .unwrap_or_else(|| {
-                            RowSelection::from(vec![RowSelector::select(limit.min(row_count))])
-                        }),
-                );
-            }
-        }
-
         let selection = selection.map(|s| s.trim().into());
 
         Ok(ReadPlan {
@@ -199,6 +146,90 @@ impl ReadPlanBuilder {
     }
 }
 
+/// A builder for [`ReadPlan`] that applies a limit and offset to the read plan
+///
+/// See [`ReadPlanBuilder::limited`] `
+pub(crate) struct LimitedReadPlanBuilder {
+    /// The underlying read plan builder
+    inner: ReadPlanBuilder,
+    /// Total number of rows in the row group before the selection, limit or
+    /// offset are applied
+    row_count: usize,
+    /// The offset to apply to the read plan
+    offset: Option<usize>,
+    /// The limit to apply to the read plan
+    limit: Option<usize>,
+}
+
+impl LimitedReadPlanBuilder {
+    /// Create a new `LimitedReadPlanBuilder` from the existing builder and number of rows
+    fn new(inner: ReadPlanBuilder, row_count: usize) -> Self {
+        Self {
+            inner,
+            row_count,
+            offset: None,
+            limit: None,
+        }
+    }
+
+    /// Set the offset to apply to the read plan
+    pub(crate) fn with_offset(mut self, offset: Option<usize>) -> Self {
+        self.offset = offset;
+        self
+    }
+
+    /// Set the limit to apply to the read plan
+    pub(crate) fn with_limit(mut self, limit: Option<usize>) -> Self {
+        self.limit = limit;
+        self
+    }
+
+    /// Finalize apply the offset and limit and return the underlying builder
+    pub(crate) fn build_limited(self) -> ReadPlanBuilder {
+        let Self {
+            mut inner,
+            row_count,
+            offset,
+            limit,
+        } = self;
+
+        // If selection is empty, truncate (needed??)
+        if !inner.selects_any() {
+            inner.selection = Some(RowSelection::from(vec![]));
+        }
+
+        // If an offset is defined, apply it to the `selection`
+        if let Some(offset) = offset {
+            inner.selection = Some(match row_count.checked_sub(offset) {
+                None => RowSelection::from(vec![]),
+                Some(remaining) => inner
+                    .selection
+                    .map(|selection| selection.offset(offset))
+                    .unwrap_or_else(|| {
+                        RowSelection::from(vec![
+                            RowSelector::skip(offset),
+                            RowSelector::select(remaining),
+                        ])
+                    }),
+            });
+        }
+
+        // If a limit is defined, apply it to the final `selection`
+        if let Some(limit) = limit {
+            inner.selection = Some(
+                inner
+                    .selection
+                    .map(|selection| selection.limit(limit))
+                    .unwrap_or_else(|| {
+                        RowSelection::from(vec![RowSelector::select(limit.min(row_count))])
+                    }),
+            );
+        }
+
+        inner
+    }
+}
+
 /// Describes what rows to read from a Parquet Row Group
 /// including based on [`RowSelection`] and [`RowFilter`].
 pub(crate) struct ReadPlan {

parquet/src/arrow/async_reader/mod.rs

@@ -638,9 +638,10 @@ where
         let plan = plan_builder
             // TODO remove this clone (builder selection is used below to fetch pages)
             .clone()
-            .with_num_rows(row_group.row_count)
+            .limited(row_group.row_count)
             .with_offset(self.offset)
             .with_limit(self.limit)
+            .build_limited()
             .build()?;
 
         let rows_after = plan.num_rows_selected().unwrap_or(row_group.row_count);