[python] Keep reading an Avro file after a fully-filtered batch

paimon-python/pypaimon/read/reader/format_avro_reader.py

@@ -56,33 +56,36 @@ def __init__(self, file_io: FileIO, file_path: str, read_fields: List[str], full
             nested_name_paths and any(len(p) > 1 for p in nested_name_paths))
 
     def read_arrow_batch(self) -> Optional[RecordBatch]:
-        pydict_data = {name: [] for name in self._fields}
-        records_in_batch = 0
+        while True:
+            pydict_data = {name: [] for name in self._fields}
+            records_in_batch = 0
 
-        for record in self._avro_reader:
-            if self._has_nested:
-                for col_name, path in zip(self._fields, self._nested_name_paths):
-                    pydict_data[col_name].append(_walk_avro_record(record, path))
-            else:
-                for col_name in self._fields:
-                    pydict_data[col_name].append(record.get(col_name))
-            records_in_batch += 1
-            if records_in_batch >= self._batch_size:
-                break
+            for record in self._avro_reader:
+                if self._has_nested:
+                    for col_name, path in zip(self._fields, self._nested_name_paths):
+                        pydict_data[col_name].append(_walk_avro_record(record, path))
+                else:
+                    for col_name in self._fields:
+                        pydict_data[col_name].append(record.get(col_name))
+                records_in_batch += 1
+                if records_in_batch >= self._batch_size:
+                    break
+
+            # No more records in the file: this is the only real end-of-input.
+            if records_in_batch == 0:
+                return None
+            if self._push_down_predicate is None:
+                return pa.RecordBatch.from_pydict(pydict_data, self._schema)
 
-        if records_in_batch == 0:
-            return None
-        if self._push_down_predicate is None:
-            return pa.RecordBatch.from_pydict(pydict_data, self._schema)
-        else:
             pa_batch = pa.Table.from_pydict(pydict_data, self._schema)
             dataset = ds.InMemoryDataset(pa_batch)
             scanner = dataset.scanner(filter=self._push_down_predicate)
             combine_chunks = scanner.to_table().combine_chunks()
             if combine_chunks.num_rows > 0:
                 return combine_chunks.to_batches()[0]
-            else:
-                return None
+            # This batch matched no rows but the file has more; keep reading rather
+            # than returning None, which the caller treats as end-of-input and would
+            # silently drop every remaining row.
 
     def close(self):
         if self._file:

paimon-python/pypaimon/tests/reader_append_only_test.py

@@ -104,6 +104,39 @@ def test_avro_ao_reader(self):
         actual = self._read_test_table(read_builder).sort_by('user_id')
         self.assertEqual(actual, self.expected)
 
+    def test_avro_ao_reader_filter_keeps_rows_past_first_batch(self):
+        # A filtered Avro read must not stop when a full batch_size (1024) block
+        # matches nothing: the reader returned None there, which the caller reads as
+        # end-of-input, silently dropping every later matching row.
+        schema = Schema.from_pyarrow_schema(
+            self.pa_schema, partition_keys=['dt'], options={'file.format': 'avro'})
+        self.catalog.create_table('default.test_avro_filter_batches', schema, False)
+        table = self.catalog.get_table('default.test_avro_filter_batches')
+
+        n = 3000
+        data = pa.Table.from_pydict({
+            'user_id': list(range(n)),
+            'item_id': [1000 + i for i in range(n)],
+            'behavior': ['a'] * n,
+            'dt': ['p1'] * n,
+        }, schema=self.pa_schema)
+        wb = table.new_batch_write_builder()
+        w, c = wb.new_write(), wb.new_commit()
+        w.write_arrow(data)
+        c.commit(w.prepare_commit())
+        w.close()
+        c.close()
+
+        pb = table.new_read_builder().new_predicate_builder()
+        # >= 2500 over 3000 rows fully filters the first two 1024-row batches, so the
+        # test distinguishes the loop from an implementation that only retries once.
+        read_builder = table.new_read_builder().with_filter(
+            pb.greater_or_equal('user_id', 2500))
+        result = self._read_test_table(read_builder)
+
+        self.assertEqual(
+            sorted(result.column('user_id').to_pylist()), list(range(2500, n)))
+
     def test_lance_ao_reader(self):
         schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'], options={'file.format': 'lance'})
         self.catalog.create_table('default.test_append_only_lance', schema, False)