This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new f06ebd6e3c [python][torch] Make map-style reads lazy (#9486)
f06ebd6e3c is described below
commit f06ebd6e3c0d2977830342cbda8e423f1fc342e5
Author: XiaoHongbo <[email protected]>
AuthorDate: Mon Aug 31 16:58:12 2026 +0800
[python][torch] Make map-style reads lazy (#9486)
---
docs/docs/pypaimon/pytorch.md | 5 +-
.../pypaimon/filesystem/caching_file_io.py | 10 +-
.../pypaimon/read/datasource/torch_dataset.py | 270 ++++++++++++++++++++-
paimon-python/pypaimon/read/table_read.py | 4 +-
.../pypaimon/tests/caching_file_io_test.py | 16 ++
paimon-python/pypaimon/tests/torch_read_test.py | 235 ++++++++++++++++++
6 files changed, 525 insertions(+), 15 deletions(-)
diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md
index 8961d76434..af9189bbae 100644
--- a/docs/docs/pypaimon/pytorch.md
+++ b/docs/docs/pypaimon/pytorch.md
@@ -54,8 +54,9 @@ for batch_idx, batch_data in enumerate(dataloader):
# {'user_id': tensor([7, 8]), 'behavior': ['g', 'h']}
```
-When the `streaming` parameter is true, it will iteratively read;
-when it is false, it will read the full amount of data into memory.
+When the `streaming` parameter is true, it will iteratively read. When it is
+false, eligible data-evolution reads fetch each DataLoader batch lazily by row
+ID; other reads retain an Arrow table in memory for map-style access.
**`prefetch_concurrency`** (default: 1): In streaming row mode, controls
reader threads per DataLoader worker. It has no effect in non-streaming mode.
diff --git a/paimon-python/pypaimon/filesystem/caching_file_io.py
b/paimon-python/pypaimon/filesystem/caching_file_io.py
index 8aa64e227f..8b10166461 100644
--- a/paimon-python/pypaimon/filesystem/caching_file_io.py
+++ b/paimon-python/pypaimon/filesystem/caching_file_io.py
@@ -341,6 +341,11 @@ class CachingFileIO(FileIO):
else:
self._whitelist = whitelist
+ def __getstate__(self):
+ state = self.__dict__.copy()
+ state['_cache'] = None
+ return state
+
# Fallback caps when local-cache.max-size is unset (memory shares the
heap).
_DEFAULT_MEMORY_CACHE_MAX_SIZE = 256 * 1024 * 1024
_DEFAULT_DISK_CACHE_MAX_SIZE = 10 * 1024 * 1024 * 1024
@@ -456,7 +461,10 @@ class CachingFileIO(FileIO):
return self._delegate.write_row(*args, **kwargs)
def __getattr__(self, name):
- return getattr(self._delegate, name)
+ delegate = self.__dict__.get('_delegate')
+ if delegate is None:
+ raise AttributeError(name)
+ return getattr(delegate, name)
def close(self):
self._delegate.close()
diff --git a/paimon-python/pypaimon/read/datasource/torch_dataset.py
b/paimon-python/pypaimon/read/datasource/torch_dataset.py
index daf435887b..de4bb2cace 100644
--- a/paimon-python/pypaimon/read/datasource/torch_dataset.py
+++ b/paimon-python/pypaimon/read/datasource/torch_dataset.py
@@ -18,6 +18,8 @@
"""
Module to read a Paimon table into PyTorch Dataset.
"""
+import bisect
+import operator
import os
import queue
import random
@@ -26,13 +28,18 @@ import warnings
from typing import Any, Callable, Iterator, List, Optional
import pyarrow as pa
+import pyarrow.compute as pc
import torch
from torch.utils.data import Dataset, IterableDataset
+from pypaimon.globalindex.indexed_split import IndexedSplit
+from pypaimon.read.query_auth_split import QueryAuthSplit
from pypaimon.read.reader.concat_batch_reader import (
_MAX_ARROW_OFFSET, _batch_offset_usage)
-from pypaimon.read.split import Split
+from pypaimon.read.split import DataSplit, Split
from pypaimon.read.table_read import TableRead
+from pypaimon.table.special_fields import SpecialFields
+from pypaimon.utils.range import Range
def _share_epoch_with_torch_workers(value):
@@ -111,12 +118,75 @@ def _balanced_slice(values: List[Any], shard_id: int,
shard_count: int):
return values[start:start + size]
+class _RowIdRangeIndex:
+
+ def __init__(self, ranges, limit=None):
+ self.ranges = []
+ self.ends = []
+ remaining = limit
+ size = 0
+ for row_range in ranges:
+ count = row_range.count()
+ if remaining is not None:
+ if remaining <= 0:
+ break
+ count = min(count, remaining)
+ remaining -= count
+ if count == 0:
+ continue
+ self.ranges.append(Range(
+ row_range.from_, row_range.from_ + count - 1))
+ size += count
+ self.ends.append(size)
+
+ def __len__(self):
+ return self.ends[-1] if self.ends else 0
+
+ def take(self, indices):
+ row_ids = []
+ for index in indices:
+ range_index = bisect.bisect_right(self.ends, index)
+ previous_end = self.ends[range_index - 1] if range_index else 0
+ row_ids.append(
+ self.ranges[range_index].from_ + index - previous_end)
+ return row_ids
+
+
+class _SplitRangeIndex:
+
+ def __init__(self, ranges_by_split):
+ self.intervals = sorted(
+ (row_range.from_, row_range.to, split_index)
+ for split_index, ranges in enumerate(ranges_by_split)
+ for row_range in ranges
+ )
+ self.starts = [interval[0] for interval in self.intervals]
+ self.max_ends = []
+ max_end = -1
+ for _, end, _ in self.intervals:
+ max_end = max(max_end, end)
+ self.max_ends.append(max_end)
+
+ def find(self, ranges):
+ split_indices = set()
+ for row_range in ranges:
+ right = bisect.bisect_right(self.starts, row_range.to)
+ left = bisect.bisect_left(
+ self.max_ends, row_range.from_, 0, right)
+ for position in range(left, right):
+ _, end, split_index = self.intervals[position]
+ if end >= row_range.from_:
+ split_indices.add(split_index)
+ return sorted(split_indices)
+
+
class TorchDataset(Dataset):
"""
- PyTorch Dataset implementation for reading Paimon table data.
+ Map-style PyTorch Dataset for Paimon table data.
- This class enables Paimon table data to be used directly with PyTorch's
- training pipeline, allowing for efficient data loading and batching.
+ Eligible data-evolution reads are fetched lazily by DataLoader batch.
+ Other reads retain their Arrow representation instead of expanding all
+ rows into Python objects.
"""
def __init__(self, table_read: TableRead, splits: List[Split]):
@@ -127,11 +197,104 @@ class TorchDataset(Dataset):
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
+ self._split_ranges = None
+ self._split_range_index = None
+ if self._supports_lazy_row_id_read():
+ self._split_ranges = [
+ self._row_ranges_for_split(split) for split in splits
+ ]
+ self._split_range_index = _SplitRangeIndex(self._split_ranges)
+ 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._materialize()
else:
- self._data = arrow_table.to_pylist()
+ self._materialize()
+
+ def _supports_lazy_row_id_read(self) -> bool:
+ if not self.table_read.table.options.row_tracking_enabled():
+ return False
+ if not self.table_read.table.options.data_evolution_enabled():
+ return False
+ if self.table_read.include_row_kind:
+ return False
+ if self.table_read.nested_name_paths:
+ return False
+ if any(self._row_id_is_masked(split) for split in self.splits):
+ return False
+ return all(self._supports_indexed_split(split) for split in
self.splits)
+
+ @staticmethod
+ def _row_id_is_masked(split) -> bool:
+ if not isinstance(split, QueryAuthSplit):
+ return False
+ masking = getattr(split.auth_result, "column_masking", None)
+ return bool(masking and SpecialFields.ROW_ID.name in masking)
+
+ @staticmethod
+ def _supports_indexed_split(split) -> bool:
+ if isinstance(split, QueryAuthSplit):
+ split = split.split
+ if isinstance(split, IndexedSplit):
+ return split.scores() is None
+ return isinstance(split, DataSplit)
+
+ def _compact_row_id_index(self):
+ if self.table_read.predicate is not None:
+ return None
+ ranges = []
+ for original, split_ranges in zip(
+ self.splits, self._split_ranges):
+ split = original
+ if isinstance(split, QueryAuthSplit):
+ if getattr(split.auth_result, "filter", None):
+ return None
+ split = split.split
+ deletion_files = split.data_deletion_files or []
+ if any(deletion is not None for deletion in deletion_files):
+ return None
+ ranges.extend(split_ranges)
+ merged = Range.sort_and_merge_overlap(ranges, True)
+ if sum(r.count() for r in ranges) != sum(
+ r.count() for r in merged):
+ return None
+ return _RowIdRangeIndex(ranges, self.table_read.limit)
+
+ @staticmethod
+ def _row_ranges_for_split(split):
+ if isinstance(split, QueryAuthSplit):
+ split = split.split
+ if isinstance(split, IndexedSplit):
+ ranges = split.row_ranges()
+ else:
+ ranges = [
+ data_file.row_id_range()
+ for data_file in split.files
+ if data_file.first_row_id is not None
+ ]
+ return Range.sort_and_merge_overlap(ranges, True)
+
+ def _materialize(self):
+ self._row_ids = None
+ self._data = self.table_read.to_arrow(self.splits)
+ self.table_read = None
+ self.splits = None
+ self._split_ranges = None
+ self._split_range_index = None
def __len__(self) -> int:
"""
@@ -140,7 +303,9 @@ class TorchDataset(Dataset):
Returns:
Total number of rows across all splits
"""
- return len(self._data)
+ if self._row_ids is not None:
+ return len(self._row_ids)
+ return 0 if self._data is None else self._data.num_rows
def __getitem__(self, index: int):
"""
@@ -152,10 +317,93 @@ class TorchDataset(Dataset):
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()
- return self._data[index]
+ 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 = []
+ split_indices = self._split_range_index.find(ranges)
+ for split_index in split_indices:
+ original = self.splits[split_index]
+ auth_result = None
+ split = original
+ if isinstance(split, QueryAuthSplit):
+ auth_result = split.auth_result
+ split = split.split
+
+ if isinstance(split, IndexedSplit):
+ split = split.data_split()
+ allowed = Range.and_(
+ ranges, self._split_ranges[split_index])
+ if not allowed:
+ continue
+
+ indexed = IndexedSplit(
+ split,
+ allowed,
+ exact_merged_row_count=sum(r.count() for r in allowed),
+ )
+ selected.append(
+ QueryAuthSplit(indexed, auth_result)
+ if auth_result is not None else indexed
+ )
+ return selected
class _BaseTorchIterDataset(IterableDataset):
diff --git a/paimon-python/pypaimon/read/table_read.py
b/paimon-python/pypaimon/read/table_read.py
index cfe13fd75b..78b2d26f41 100644
--- a/paimon-python/pypaimon/read/table_read.py
+++ b/paimon-python/pypaimon/read/table_read.py
@@ -669,7 +669,9 @@ class TableRead:
Args:
splits: Splits to read.
- streaming: Whether to stream data.
+ streaming: Whether to return an iterable dataset. Non-streaming
+ eligible data-evolution reads fetch map-style batches lazily
+ by row ID.
prefetch_concurrency: Reader threads per DataLoader worker in row
format.
batch_format: ``"row"``, ``"pyarrow"``, or ``"torch"``. Batch
diff --git a/paimon-python/pypaimon/tests/caching_file_io_test.py
b/paimon-python/pypaimon/tests/caching_file_io_test.py
index 12288ef120..940a4f4ab2 100644
--- a/paimon-python/pypaimon/tests/caching_file_io_test.py
+++ b/paimon-python/pypaimon/tests/caching_file_io_test.py
@@ -20,6 +20,7 @@
import io
import os
+import pickle
import shutil
import tempfile
import threading
@@ -385,6 +386,21 @@ class CachingFileIOTest(unittest.TestCase):
caching_io.new_output_stream("/out")
delegate.new_output_stream.assert_called_once_with("/out")
+ def test_pickle_drops_cache(self):
+ from pypaimon.filesystem.local_file_io import LocalFileIO
+ from pypaimon.utils.file_type import FileType
+
+ cache = LocalDiskCacheManager(
+ self.cache_dir, 2 ** 63 - 1, block_size=64)
+ caching_io = CachingFileIO(
+ LocalFileIO(), cache, {FileType.META, FileType.DATA})
+
+ restored = pickle.loads(pickle.dumps(caching_io))
+
+ self.assertIsNone(restored._cache)
+ self.assertEqual({FileType.META, FileType.DATA}, restored._whitelist)
+ self.assertIsInstance(restored._delegate, LocalFileIO)
+
def test_to_filesystem_path_forwarded_to_delegate(self):
delegate = MagicMock()
delegate.to_filesystem_path.side_effect = lambda p:
p.replace("oss://", "")
diff --git a/paimon-python/pypaimon/tests/torch_read_test.py
b/paimon-python/pypaimon/tests/torch_read_test.py
index 809fe582ac..0dc269110d 100644
--- a/paimon-python/pypaimon/tests/torch_read_test.py
+++ b/paimon-python/pypaimon/tests/torch_read_test.py
@@ -33,14 +33,18 @@ import torch
from torch.utils.data import DataLoader
from pypaimon import CatalogFactory, Schema
+from pypaimon.catalog.table_query_auth import TableQueryAuthResult
from pypaimon.multimodal.table import MultimodalTable
from pypaimon.read.datasource.torch_dataset import (
+ _SplitRangeIndex,
TorchIterDataset,
TorchShuffledIterDataset,
_resolve_distributed_context,
)
+from pypaimon.read.table_read import TableRead
from pypaimon.table.file_store_table import FileStoreTable
+from pypaimon.utils.range import Range
def _collect_spawned_worker_splits(dataset, output):
@@ -500,8 +504,239 @@ class TorchReadTest(unittest.TestCase):
self.assertEqual(sorted_behaviors, expected_behaviors,
f"Behaviors mismatch. Expected {expected_behaviors},
got {sorted_behaviors}")
+ if not is_streaming:
+ self.assertIsInstance(dataset._data, pa.Table)
+
print(f"✓ Test passed: Successfully read {len(all_user_ids)} rows with
correct data")
+ def test_non_streaming_row_tracking_reads_batches_lazily(self):
+ schema = Schema.from_pyarrow_schema(
+ self.pa_schema,
+ partition_keys=['dt'],
+ options={
+ 'data-evolution.enabled': 'true',
+ 'row-tracking.enabled': 'true',
+ },
+ )
+ identifier = 'default.test_torch_lazy_row_tracking'
+ self.catalog.create_table(identifier, schema, False)
+ table = self.catalog.get_table(identifier)
+ self._write_test_table(table)
+
+ read_builder = table.new_read_builder().with_projection(
+ ['user_id', 'behavior'])
+ splits = read_builder.new_scan().plan().splits()
+ self.assertGreater(len(splits), 1)
+ table_read = read_builder.new_read()
+ expected = table_read.to_arrow(splits).to_pylist()
+
+ with patch.object(
+ TableRead, 'to_arrow', side_effect=AssertionError(
+ 'dataset construction must not read table data')):
+ dataset = table_read.to_torch(splits, streaming=False)
+
+ self.assertIsNone(dataset._data)
+ self.assertEqual(len(expected), len(dataset))
+ indices = [len(dataset) - 1, 1, 1, 0]
+ self.assertEqual(
+ [expected[index] for index in indices],
+ dataset.__getitems__(indices),
+ )
+ self.assertEqual(expected[-1], dataset[-1])
+ restored = pickle.loads(pickle.dumps(dataset))
+ self.assertEqual(expected[1], restored[1])
+
+ row_id_builder = table.new_read_builder().with_projection(
+ ['user_id', '_ROW_ID'])
+ row_id_splits = row_id_builder.new_scan().plan().splits()
+ row_id_read = row_id_builder.new_read()
+ expected_row_id = row_id_read.to_arrow(
+ row_id_splits).to_pylist()[0]
+ row_id_dataset = row_id_read.to_torch(
+ row_id_splits, streaming=False)
+ self.assertEqual(expected_row_id, row_id_dataset[0])
+
+ loader = DataLoader(
+ dataset,
+ batch_size=3,
+ num_workers=2,
+ shuffle=True,
+ generator=torch.Generator().manual_seed(7),
+ )
+ actual_ids = []
+ for batch in loader:
+ actual_ids.extend(batch['user_id'].tolist())
+ self.assertEqual(
+ sorted(row['user_id'] for row in expected),
+ sorted(actual_ids),
+ )
+
+ def
test_non_streaming_row_tracking_without_data_evolution_materializes(self):
+ schema = Schema.from_pyarrow_schema(
+ self.pa_schema,
+ options={'row-tracking.enabled': 'true'},
+ )
+ identifier = 'default.test_torch_row_tracking_without_data_evolution'
+ self.catalog.create_table(identifier, schema, False)
+ table = self.catalog.get_table(identifier)
+ self._write_test_table(table)
+
+ read_builder = table.new_read_builder().with_projection(
+ ['user_id', 'behavior'])
+ splits = read_builder.new_scan().plan().splits()
+ table_read = read_builder.new_read()
+ expected = table_read.to_arrow(
+ splits, parallelism=1).to_pylist()
+ with patch.object(
+ table_read, 'to_arrow', wraps=table_read.to_arrow) as read:
+ dataset = table_read.to_torch(splits, streaming=False)
+ read.assert_called_once_with(splits)
+
+ self.assertIsInstance(dataset._data, pa.Table)
+ self.assertIsNone(dataset.table_read)
+ self.assertIsNone(dataset.splits)
+ with patch.object(
+ TableRead, 'to_arrow', side_effect=AssertionError(
+ 'materialized dataset must not read another batch')):
+ self.assertEqual(expected, dataset[:])
+
+ def test_non_streaming_row_id_masking_materializes(self):
+ schema = Schema.from_pyarrow_schema(
+ self.pa_schema,
+ options={
+ 'data-evolution.enabled': 'true',
+ 'row-tracking.enabled': 'true',
+ },
+ )
+ identifier = 'default.test_torch_masked_row_id'
+ self.catalog.create_table(identifier, schema, False)
+ table = self.catalog.get_table(identifier)
+ self._write_test_table(table)
+ auth = TableQueryAuthResult(
+ filter=None,
+ column_masking={
+ '_ROW_ID': json.dumps({'name': 'NULL'}),
+ },
+ )
+ table.catalog_environment.table_query_auth = (
+ lambda options, table_identifier: lambda select: auth
+ )
+
+ read_builder = table.new_read_builder().with_projection(
+ ['user_id', 'behavior'])
+ splits = read_builder.new_scan().plan().splits()
+ table_read = read_builder.new_read()
+ expected = table_read.to_arrow(
+ splits, parallelism=1).to_pylist()
+ dataset = table_read.to_torch(splits, streaming=False)
+
+ self.assertIsInstance(dataset._data, pa.Table)
+ self.assertEqual(expected, dataset[:])
+
+ def test_non_streaming_dataset_with_cache_is_pickleable(self):
+ with tempfile.TemporaryDirectory() as tempdir:
+ catalog = CatalogFactory.create({
+ 'warehouse': os.path.join(tempdir, 'warehouse'),
+ 'local-cache.enabled': 'true',
+ 'local-cache.whitelist': 'meta,global-index,data',
+ })
+ catalog.create_database('default', True)
+ for suffix, options, lazy in [
+ ('arrow', {}, False),
+ ('lazy', {
+ 'data-evolution.enabled': 'true',
+ 'row-tracking.enabled': 'true',
+ }, True),
+ ]:
+ with self.subTest(suffix=suffix):
+ identifier = 'default.test_torch_cache_' + suffix
+ schema = Schema.from_pyarrow_schema(
+ self.pa_schema, options=options)
+ catalog.create_table(identifier, schema, False)
+ table = catalog.get_table(identifier)
+ self._write_test_table(table)
+ read_builder = table.new_read_builder().with_projection(
+ ['user_id', 'behavior'])
+ splits = read_builder.new_scan().plan().splits()
+ dataset = read_builder.new_read().to_torch(
+ splits, streaming=False)
+ self.assertEqual(lazy, dataset._data is None)
+
+ restored = pickle.loads(pickle.dumps(dataset))
+ self.assertEqual(dataset[0], restored[0])
+
+ def test_split_range_index(self):
+ index = _SplitRangeIndex([
+ [Range(split * 10, split * 10 + 9)]
+ for split in range(10000)
+ ])
+
+ self.assertEqual([5000], index.find([Range(50003, 50005)]))
+ self.assertEqual(
+ [0, 9999], index.find([Range(0, 0), Range(99999, 99999)]))
+
+ def test_non_streaming_row_tracking_preserves_filter_and_limit(self):
+ schema = Schema.from_pyarrow_schema(
+ self.pa_schema,
+ options={
+ 'data-evolution.enabled': 'true',
+ 'row-tracking.enabled': 'true',
+ },
+ )
+ identifier = 'default.test_torch_lazy_filter_limit'
+ self.catalog.create_table(identifier, schema, False)
+ table = self.catalog.get_table(identifier)
+ self._write_test_table(table)
+
+ read_builder = table.new_read_builder().with_projection(
+ ['user_id', 'behavior'])
+ predicate = read_builder.new_predicate_builder().greater_than(
+ 'user_id', 2)
+ read_builder.with_filter(predicate).with_limit(3)
+ splits = read_builder.new_scan().plan().splits()
+ table_read = read_builder.new_read()
+ expected = table_read.to_arrow(
+ splits, parallelism=1).to_pylist()
+ dataset = table_read.to_torch(splits, streaming=False)
+
+ self.assertIsNone(dataset._data)
+ self.assertEqual(3, len(dataset))
+ self.assertEqual(
+ expected,
+ dataset.__getitems__(range(len(dataset))),
+ )
+
+ def test_non_streaming_row_tracking_respects_deletion_vectors(self):
+ schema = Schema.from_pyarrow_schema(
+ self.pa_schema,
+ options={
+ 'data-evolution.enabled': 'true',
+ 'row-tracking.enabled': 'true',
+ 'deletion-vectors.enabled': 'true',
+ },
+ )
+ identifier = 'default.test_torch_lazy_deletion_vector'
+ self.catalog.create_table(identifier, schema, False)
+ table = self.catalog.get_table(identifier)
+ self._write_test_table(table)
+
+ write_builder = table.new_batch_write_builder()
+ messages = write_builder.new_update().delete_by_row_id([2])
+ table_commit = write_builder.new_commit()
+ table_commit.commit(messages)
+ table_commit.close()
+
+ read_builder = table.new_read_builder().with_projection(
+ ['user_id', 'behavior'])
+ splits = read_builder.new_scan().plan().splits()
+ table_read = read_builder.new_read()
+ expected = table_read.to_arrow(splits).to_pylist()
+ dataset = table_read.to_torch(splits, streaming=False)
+
+ self.assertIsNone(dataset._data)
+ self.assertIsInstance(dataset._row_ids, pa.Array)
+ self.assertEqual(expected, dataset[:])
+
def test_torch_streaming_prefetch_concurrency(self):
schema = Schema.from_pyarrow_schema(self.pa_schema,
partition_keys=['user_id'])
self.catalog.create_table('default.test_torch_prefetch_concurrency',
schema, False)
