JingsongLi commented on code in PR #9486:
URL: https://github.com/apache/paimon/pull/9486#discussion_r3893015763
##########
paimon-python/pypaimon/read/datasource/torch_dataset.py:
##########
@@ -152,10 +257,97 @@ def __getitem__(self, index: int):
Returns:
Dictionary containing the row data
"""
- if not self._data:
+ if len(self) == 0:
return None
+ if isinstance(index, slice):
+ return self.__getitems__(range(*index.indices(len(self))))
+ return self.__getitems__([index])[0]
+
+ def __getitems__(self, indices) -> List[dict]:
+ normalized = [self._normalize_index(index) for index in indices]
+ if not normalized:
+ return []
+ if self._row_ids is None:
+ return self._data.take(pa.array(
+ normalized, type=pa.int64())).to_pylist()
+
+ if isinstance(self._row_ids, _RowIdRangeIndex):
+ row_ids = self._row_ids.take(normalized)
+ else:
+ row_ids = self._row_ids.take(pa.array(
+ normalized, type=pa.int64())).to_pylist()
+ ranges = Range.sort_and_merge_overlap(
+ [Range(row_id, row_id) for row_id in set(row_ids)], True)
+ splits = self._select_splits(ranges)
+
+ output_has_row_id = any(
+ field.name == SpecialFields.ROW_ID.name
+ for field in self.table_read.read_type
+ )
+ read_type = list(self.table_read.read_type)
+ if not output_has_row_id:
+ read_type.append(SpecialFields.ROW_ID)
+ batch_read = TableRead(
+ self.table_read.table,
+ self.table_read.predicate,
+ read_type,
+ include_row_kind=self.table_read.include_row_kind,
+ )
+ rows = batch_read.to_arrow(splits).to_pylist()
+ by_row_id = {}
+ for row in rows:
+ row_id = row[SpecialFields.ROW_ID.name]
+ if not output_has_row_id:
+ del row[SpecialFields.ROW_ID.name]
+ by_row_id[row_id] = row
+ missing = set(row_ids) - set(by_row_id)
+ if missing:
+ raise RuntimeError(
+ "Paimon rows disappeared while reading TorchDataset: %s"
+ % sorted(missing)
+ )
+ return [by_row_id[row_id] for row_id in row_ids]
+
+ def _normalize_index(self, index) -> int:
+ index = operator.index(index)
+ if index < 0:
+ index += len(self)
+ if index < 0 or index >= len(self):
+ raise IndexError("TorchDataset index out of range")
+ return index
+
+ def _select_splits(self, ranges) -> List[Split]:
+ selected = []
+ for original in self.splits:
Review Comment:
Verified in f1c6356: split ranges are cached and `_SplitRangeIndex` narrows
each batch to candidate splits. I also compared the index with brute-force
overlap checks across 5,000 randomized cases. This addresses my concern. Please
resolve this thread when convenient. Thanks.
##########
paimon-python/pypaimon/read/datasource/torch_dataset.py:
##########
@@ -152,10 +257,97 @@ def __getitem__(self, index: int):
Returns:
Dictionary containing the row data
"""
- if not self._data:
+ if len(self) == 0:
return None
+ if isinstance(index, slice):
+ return self.__getitems__(range(*index.indices(len(self))))
+ return self.__getitems__([index])[0]
+
+ def __getitems__(self, indices) -> List[dict]:
+ normalized = [self._normalize_index(index) for index in indices]
+ if not normalized:
+ return []
+ if self._row_ids is None:
+ return self._data.take(pa.array(
+ normalized, type=pa.int64())).to_pylist()
+
+ if isinstance(self._row_ids, _RowIdRangeIndex):
+ row_ids = self._row_ids.take(normalized)
+ else:
+ row_ids = self._row_ids.take(pa.array(
+ normalized, type=pa.int64())).to_pylist()
+ ranges = Range.sort_and_merge_overlap(
+ [Range(row_id, row_id) for row_id in set(row_ids)], True)
+ splits = self._select_splits(ranges)
+
+ output_has_row_id = any(
+ field.name == SpecialFields.ROW_ID.name
+ for field in self.table_read.read_type
+ )
+ read_type = list(self.table_read.read_type)
+ if not output_has_row_id:
+ read_type.append(SpecialFields.ROW_ID)
+ batch_read = TableRead(
Review Comment:
Verified in f1c6356: `_ROW_ID` masking now disables lazy routing and falls
back to materialization, with a regression test for the masked-row-ID case.
This addresses my concern. Please resolve this thread when convenient. Thanks.
##########
paimon-python/pypaimon/read/datasource/torch_dataset.py:
##########
@@ -127,11 +169,72 @@ def __init__(self, table_read: TableRead, splits:
List[Split]):
table_read: TableRead instance for reading data
splits: List of splits to read
"""
- arrow_table = table_read.to_arrow(splits)
- if arrow_table is None or arrow_table.num_rows == 0:
- self._data = []
+ self.table_read = table_read
+ self.splits = splits
+ self._data = None
+ self._row_ids = None
+ if self._supports_lazy_row_id_read():
+ self._row_ids = self._compact_row_id_index()
+ if self._row_ids is None:
+ row_id_read = TableRead(
+ table_read.table,
+ table_read.predicate,
+ [SpecialFields.ROW_ID],
+ limit=table_read.limit,
+ )
+ row_id_table = row_id_read.to_arrow(splits)
+ self._row_ids = row_id_table.column(
+ SpecialFields.ROW_ID.name).combine_chunks()
+ if pc.count_distinct(self._row_ids).as_py() != len(
+ self._row_ids):
+ self._row_ids = None
+ self._data = table_read.to_arrow(splits)
else:
- self._data = arrow_table.to_pylist()
+ self._data = table_read.to_arrow(splits)
+
+ def _supports_lazy_row_id_read(self) -> bool:
+ if not self.table_read.table.options.row_tracking_enabled():
Review Comment:
Verified in f1c6356: the lazy path now requires data evolution, and the new
non-data-evolution regression test confirms that this case materializes once.
This addresses my concern. Please resolve this thread when convenient. Thanks.
--
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]
