steFaiz commented on code in PR #8404:
URL: https://github.com/apache/paimon/pull/8404#discussion_r3506055647
##########
paimon-python/pypaimon/deletionvectors/apply_deletion_vector_reader.py:
##########
@@ -15,53 +15,87 @@
# specific language governing permissions and limitations
# under the License.
-from typing import Optional
+from typing import List, Optional
import pyarrow
from pyarrow import RecordBatch
from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader
from pypaimon.read.reader.iface.record_iterator import RecordIterator
from pypaimon.deletionvectors.deletion_vector import DeletionVector
-from pypaimon.utils.roaring_bitmap import RoaringBitmap
from pypaimon.read.reader.iface.record_reader import RecordReader
+class PositionMappedDeletionVector:
+ """
+ Adapts a deletion vector to the returned positions of the current reader.
+ """
+
+ def __init__(
+ self,
+ deletion_vector: DeletionVector,
+ file_offset: int = 0,
+ row_positions: Optional[List[int]] = None,
+ ):
+ self._deletion_vector = deletion_vector
+ self._file_offset = file_offset
+ self._row_positions = row_positions
+
+ def is_deleted(self, position: int) -> bool:
+ return
self._deletion_vector.is_deleted(self._mapped_position(position))
+
+ def _mapped_position(self, position: int) -> int:
+ if self._row_positions is None:
+ return self._file_offset + position
+ if position >= len(self._row_positions):
+ raise ValueError(
+ "Deletion vector row positions are fewer than returned rows."
+ )
+ return self._file_offset + self._row_positions[position]
+
+
class ApplyDeletionVectorReader(RecordBatchReader):
"""
A RecordReader which applies DeletionVector to filter records.
"""
- def __init__(self, reader: RecordReader, deletion_vector: DeletionVector):
+ def __init__(
+ self,
+ reader: RecordReader,
+ deletion_vector,
+ ):
"""
Initialize an ApplyDeletionVectorReader.
Args:
reader: The underlying record reader.
- deletion_vector: The deletion vector to apply.
+ deletion_vector: The deletion vector to apply. It should already
+ be mapped to the returned positions of the underlying reader.
"""
self._reader = reader
self._deletion_vector = deletion_vector
+ self._returned_position = 0
def reader(self) -> RecordReader:
return self._reader
- def deletion_vector(self) -> DeletionVector:
+ def deletion_vector(self):
return self._deletion_vector
def read_arrow_batch(self) -> Optional[RecordBatch]:
self._reader: RecordBatchReader
arrow_batch = self._reader.read_arrow_batch()
if arrow_batch is None:
return None
- # Remove the deleted rows from the batch
- range_bitmap = RoaringBitmap()
- return_batch_pos = self._reader.return_batch_pos()
- range_bitmap.add_range(return_batch_pos - arrow_batch.num_rows,
return_batch_pos - 1)
- intersection_bitmap = RoaringBitmap.remove_all(range_bitmap,
self._deletion_vector.bit_map())
- added_row_list = [x - (return_batch_pos - arrow_batch.num_rows) for x
in
- list(intersection_bitmap)]
- return arrow_batch.take(pyarrow.array(added_row_list,
type=pyarrow.int32()))
+
+ start = self._returned_position
+ end = start + arrow_batch.num_rows
+ self._returned_position = end
+ keep_indices = [
+ i for i, position in enumerate(range(start, end))
+ if not self._deletion_vector.is_deleted(position)
+ ]
+ return arrow_batch.take(pyarrow.array(keep_indices,
type=pyarrow.int32()))
Review Comment:
Thanks! This is a bug! I've wrapped the related code with While loop. Also
added a test.
--
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]
