JingsongLi commented on code in PR #9486:
URL: https://github.com/apache/paimon/pull/9486#discussion_r3891813462
##########
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:
[P1] Please gate lazy row-ID reads on data evolution, or teach
`RawFileSplitRead` to honor `IndexedSplit`. `row-tracking.enabled=true` is
valid without data evolution, but that path selects `RawFileSplitRead`, which
ignores the `IndexedSplit.row_ranges` created below. I reproduced this with a
10-row table: `dataset[3]` returned the correct row, but the batch read
physically returned all 10 rows; the equivalent data-evolution table returned
only one. With shuffled mini-batches this can reread whole target-sized splits
per batch and be much more expensive than the previous one-time
materialization. Please either return `False` here when
`data_evolution_enabled()` is false, or add row-range pushdown to
`RawFileSplitRead`, plus a non-data-evolution regression test that asserts the
physical rows read.
##########
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:
[P2] Please avoid rescanning every split for every DataLoader batch. This
loop is on the `__getitems__` hot path and rebuilds merged file ranges for all
splits each time, making routing O(number of batches × number of splits). In a
simple benchmark with 10,000 one-file splits, `_select_splits` alone took about
9.5 ms per batch before any I/O. Please cache each split’s merged ranges and
build a row-ID interval index once during construction, then binary-search only
the intersecting splits here.
--
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]
