Move Selection logic into ReadPlan builder

parquet/src/arrow/arrow_reader/mod.rs

@@ -809,53 +809,43 @@ impl ParquetRecordBatchReader {
     /// simplify error handling with `?`
     fn next_inner(&mut self) -> Result<Option<RecordBatch>> {
         let mut read_records = 0;
-        let batch_size = self.batch_size();
-        match self.read_plan.selection_mut() {
-            Some(selection) => {
-                while read_records < batch_size && !selection.is_empty() {
-                    let front = selection.pop_front().unwrap();
+
+        match &mut self.read_plan {
+            ReadPlan::All { batch_size } => {
+                self.array_reader.read_records(*batch_size)?;
+            }
+            ReadPlan::Subset { iterator } => {
+                let batch_size = iterator.batch_size();
+
+                while read_records < batch_size {
+                    let Some(front) = self.read_plan.next() else {
+                        break;
+                    };
+
                     if front.skip {
                         let skipped = self.array_reader.skip_records(front.row_count)?;
 
                         if skipped != front.row_count {
                             return Err(general_err!(
-                                "failed to skip rows, expected {}, got {}",
+                                "Internal Error: failed to skip rows, expected {}, got {}",
                                 front.row_count,
                                 skipped
                             ));
                         }
-                        continue;
-                    }
-
-                    //Currently, when RowSelectors with row_count = 0 are included then its interpreted as end of reader.
-                    //Fix is to skip such entries. See https://github.com/apache/arrow-rs/issues/2669
-                    if front.row_count == 0 {
-                        continue;
-                    }
-
-                    // try to read record
-                    let need_read = batch_size - read_records;
-                    let to_read = match front.row_count.checked_sub(need_read) {
-                        Some(remaining) if remaining != 0 => {
-                            // if page row count less than batch_size we must set batch size to page row count.
-                            // add check avoid dead loop
-                            selection.push_front(RowSelector::select(remaining));
-                            need_read
+                    } else {
+                        let read = self.array_reader.read_records(front.row_count)?;
+                        if read == 0 {
+                            break;
                         }
-                        _ => front.row_count,
-                    };
-                    match self.array_reader.read_records(to_read)? {
-                        0 => break,
-                        rec => read_records += rec,
+
+                        read_records += read
                     };
                 }
             }
-            None => {
-                self.array_reader.read_records(batch_size)?;
-            }
-        };
+        }
 
         let array = self.array_reader.consume_batch()?;
+
         let struct_array = array.as_struct_opt().ok_or_else(|| {
             ArrowError::ParquetError("Struct array reader should return struct array".to_string())
         })?;
@@ -928,9 +918,13 @@ impl ParquetRecordBatchReader {
         }
     }
 
+    #[cfg(test)]
     #[inline(always)]
     pub(crate) fn batch_size(&self) -> usize {
-        self.read_plan.batch_size()
+        match &self.read_plan {
+            ReadPlan::All { batch_size } => *batch_size,
+            ReadPlan::Subset { iterator } => iterator.batch_size(),
+        }
     }
 }
 
@@ -3910,7 +3904,7 @@ mod tests {
             .build()
             .unwrap();
         assert_ne!(1024, num_rows);
-        assert_eq!(reader.read_plan.batch_size(), num_rows as usize);
+        assert_eq!(reader.batch_size(), num_rows as usize);
     }
 
     #[test]