fix(row selection): off-by-one error

apache · Aug 16, 2024 · 1ec6d9a · 1ec6d9a
1 parent 8f74869
commit 1ec6d9a
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/crates/iceberg/src/arrow/reader.rs b/crates/iceberg/src/arrow/reader.rs
@@ -441,9 +441,6 @@ impl ArrowReader {
                 // skip row groups that aren't present in selected_row_groups
                 if idx == selected_row_groups[selected_row_groups_idx] {
                     selected_row_groups_idx += 1;
-                    if selected_row_groups_idx == selected_row_groups.len() {
-                        break;
-                    }
                 } else {
                     continue;
                 }
@@ -459,6 +456,12 @@ impl ArrowReader {
             )?;
 
             results.push(selections_for_page);
+
+            if let Some(selected_row_groups) = selected_row_groups {
+                if selected_row_groups_idx == selected_row_groups.len() {
+                    break;
+                }
+            }
         }
 
         Ok(results.into_iter().flatten().collect::<Vec<_>>().into())