Google Git Google Git
Sign in
apache/paimon/refs/heads/codex/append-unknown-row-count/./paimon-python/pypaimon/tests/torch_read_test.py
blob: 9316b4b7485d5e5fd9c9e8dfb03fb28fe7e4a6e6 [file]
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
import json
import multiprocessing
import os
import pickle
import random
import shutil
import subprocess
import sys
import tempfile
import unittest
from types import SimpleNamespace
from unittest.mock import patch
import pyarrow as pa
from parameterized import parameterized
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,
TorchBatchIterDataset,
TorchShuffledBatchIterDataset,
TorchShuffledIterDataset,
_resolve_distributed_context,
_shuffle_record_batches,
)
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):
os.environ.pop("RANK", None)
os.environ.pop("WORLD_SIZE", None)
output.put(dataset._worker_splits(None))
class _EmptyTableRead:
read_type = []
def __init__(self, limit=None):
self.limit = limit
def _output_arrow_schema(self):
return pa.schema([])
class TorchDistributedShardingTest(unittest.TestCase):
@staticmethod
def _table_read(limit=None):
return _EmptyTableRead(limit)
@staticmethod
def _worker(worker_id, num_workers):
return SimpleNamespace(id=worker_id, num_workers=num_workers)
def _dataset(
self,
splits,
limit=None,
dataset_type=TorchIterDataset,
**kwargs
):
return dataset_type(
self._table_read(limit),
splits,
auto_detect_rank=True,
**kwargs
)
def _assignments(self, split_count, world_size, num_workers):
splits = list(range(split_count))
assignments = {}
for rank in range(world_size):
dataset = self._dataset(splits)
for worker_id in range(num_workers):
with patch(
"pypaimon.read.datasource.torch_dataset."
"_resolve_distributed_context",
return_value=(rank, world_size),
):
assignments[(rank, worker_id)] = dataset._worker_splits(
self._worker(worker_id, num_workers)
)
return assignments
def assertCompleteNonOverlapping(self, assignments, expected):
assigned = [
split for splits in assignments.values() for split in splits
]
self.assertCountEqual(assigned, expected)
self.assertEqual(len(assigned), len(set(assigned)))
@parameterized.expand([
("single", 7, 1, 1, [7]),
("workers", 10, 1, 3, [4, 3, 3]),
("ranks", 10, 3, 1, [4, 3, 3]),
("rank_workers", 17, 3, 2, [3, 3, 3, 3, 3, 2]),
("uneven", 11, 2, 2, [3, 3, 3, 2]),
("sparse", 3, 2, 3, [1, 1, 0, 1, 0, 0]),
])
def test_balanced_assignments(
self, _, split_count, world_size, num_workers, expected_sizes
):
assignments = self._assignments(
split_count, world_size, num_workers
)
self.assertCompleteNonOverlapping(
assignments, list(range(split_count))
)
self.assertEqual(
[len(splits) for splits in assignments.values()], expected_sizes
)
def test_binding_limit_rejects_distributed_sharding(self):
splits = [SimpleNamespace(row_count=10) for _ in range(4)]
dataset = self._dataset(splits, limit=5)
with patch(
"pypaimon.read.datasource.torch_dataset."
"_resolve_distributed_context",
return_value=(0, 2),
), self.assertRaisesRegex(ValueError, "limit is not supported"):
dataset._worker_splits(None)
def test_zero_limit_returns_no_splits(self):
dataset = self._dataset([SimpleNamespace(row_count=10)], limit=0)
with patch(
"pypaimon.read.datasource.torch_dataset."
"_resolve_distributed_context",
return_value=(1, 2),
):
self.assertEqual(dataset._worker_splits(None), [])
def test_initialized_distributed_context_precedes_environment(self):
with patch.dict(
os.environ, {"RANK": "4", "WORLD_SIZE": "5"}, clear=True
), patch.object(
torch.distributed, "is_available", return_value=True
), patch.object(
torch.distributed, "is_initialized", return_value=True
), patch.object(
torch.distributed, "get_rank", return_value=1
), patch.object(
torch.distributed, "get_world_size", return_value=3
):
context = _resolve_distributed_context(True)
self.assertEqual(context, (1, 3))
def test_explicit_context_ignores_global_process_group(self):
with patch.object(
torch.distributed, "is_available", return_value=True
), patch.object(
torch.distributed, "is_initialized", return_value=True
), patch.object(
torch.distributed, "get_rank", return_value=3
), patch.object(
torch.distributed, "get_world_size", return_value=4
):
dataset = TorchIterDataset(
self._table_read(),
list(range(8)),
sharding_rank=1,
sharding_world_size=2,
)
assigned = dataset._worker_splits(None)
self.assertEqual((dataset.rank, dataset.world_size), (1, 2))
self.assertEqual(assigned, [4, 5, 6, 7])
def test_context_is_resolved_after_dataset_construction(self):
dataset = TorchIterDataset(
self._table_read(), list(range(8)), auto_detect_rank=True
)
with patch.dict(
os.environ, {"RANK": "2", "WORLD_SIZE": "4"}, clear=True
), patch.object(
torch.distributed, "is_available", return_value=True
), patch.object(
torch.distributed, "is_initialized", return_value=False
):
context = _resolve_distributed_context(True)
assigned = dataset._worker_splits(None)
self.assertEqual(context, (2, 4))
self.assertEqual(assigned, [4, 5])
def test_constructor_context_is_preserved_in_spawned_worker(self):
with patch(
"pypaimon.read.datasource.torch_dataset."
"_resolve_distributed_context",
return_value=(0, 1),
):
dataset = TorchIterDataset(
self._table_read(), list(range(8)), auto_detect_rank=True
)
context = multiprocessing.get_context("spawn")
output = context.Queue()
process = context.Process(
target=_collect_spawned_worker_splits,
args=(dataset, output),
)
with patch(
"pypaimon.read.datasource.torch_dataset."
"_resolve_distributed_context",
return_value=(1, 2),
):
process.start()
process.join(30)
if process.is_alive():
process.terminate()
process.join()
self.assertEqual(process.exitcode, 0)
self.assertEqual(output.get(timeout=5), [4, 5, 6, 7])
output.close()
def test_same_process_uses_latest_context(self):
with patch(
"pypaimon.read.datasource.torch_dataset."
"_resolve_distributed_context",
return_value=(1, 2),
):
dataset = TorchIterDataset(
self._table_read(), list(range(8)), auto_detect_rank=True
)
with patch(
"pypaimon.read.datasource.torch_dataset."
"_resolve_distributed_context",
return_value=(0, 1),
):
self.assertEqual(dataset._worker_splits(None), list(range(8)))
def test_auto_falls_back_to_single_process(self):
with patch.dict(os.environ, {}, clear=True), patch.object(
torch.distributed, "is_available", return_value=False
):
context = _resolve_distributed_context(True)
self.assertEqual(context, (0, 1))
def test_disabled_preserves_worker_sharding(self):
splits = list(range(8))
with patch.dict(
os.environ, {"RANK": "1", "WORLD_SIZE": "2"}, clear=True
), patch.object(
torch.distributed, "is_available", return_value=True
), patch.object(
torch.distributed, "is_initialized", return_value=True
):
dataset = TorchIterDataset(
self._table_read(), splits, auto_detect_rank=False
)
assigned = dataset._worker_splits(self._worker(1, 2))
self.assertEqual(assigned, list(range(4, 8)))
def test_shuffled_dataset_is_reproducible_and_rank_local(self):
splits = list(range(20))
datasets = [
self._dataset(
splits,
dataset_type=TorchShuffledIterDataset,
seed=17,
buffer_size=20,
)
for _ in range(2)
]
local_splits = []
for rank, dataset in enumerate(datasets):
with patch(
"pypaimon.read.datasource.torch_dataset."
"_resolve_distributed_context",
return_value=(rank, 2),
):
local_splits.append(dataset._worker_splits(None))
self.assertTrue(set(local_splits[0]).isdisjoint(local_splits[1]))
self.assertCountEqual(local_splits[0] + local_splits[1], splits)
restored = pickle.loads(pickle.dumps(datasets[1]))
self.assertTrue(restored.auto_detect_rank)
rows = [{"id": value} for value in range(20)]
with patch(
"pypaimon.read.datasource.torch_dataset."
"_resolve_distributed_context",
return_value=(0, 2),
):
first = list(datasets[0]._iter_buffer_shuffled_rows(iter(rows), 0))
repeat = list(datasets[0]._iter_buffer_shuffled_rows(iter(rows), 0))
other_worker = list(
datasets[0]._iter_buffer_shuffled_rows(iter(rows), 1)
)
with patch(
"pypaimon.read.datasource.torch_dataset."
"_resolve_distributed_context",
return_value=(1, 2),
):
other_rank = list(
datasets[1]._iter_buffer_shuffled_rows(iter(rows), 0)
)
self.assertEqual(first, repeat)
self.assertNotEqual(first, other_rank)
self.assertNotEqual(first, other_worker)
datasets[0].set_epoch(1)
with patch(
"pypaimon.read.datasource.torch_dataset."
"_resolve_distributed_context",
return_value=(0, 2),
):
next_epoch = list(
datasets[0]._iter_buffer_shuffled_rows(iter(rows), 0)
)
self.assertNotEqual(first, next_epoch)
def test_invalid_distributed_context(self):
with self.assertRaisesRegex(ValueError, "auto_detect_rank"):
TorchIterDataset(
self._table_read(), [], auto_detect_rank="auto"
)
with self.assertRaisesRegex(ValueError, "must be set together"):
TorchIterDataset(
self._table_read(), [], sharding_rank=0
)
with self.assertRaisesRegex(ValueError, "cannot be combined"):
TorchIterDataset(
self._table_read(),
[],
auto_detect_rank=True,
sharding_rank=0,
sharding_world_size=1,
)
with patch.dict(os.environ, {"RANK": "one"}, clear=True), patch.object(
torch.distributed, "is_available", return_value=False
), self.assertRaisesRegex(ValueError, "must be set together"):
_resolve_distributed_context(True)
with patch.dict(
os.environ, {"RANK": "one", "WORLD_SIZE": "2"}, clear=True
), patch.object(
torch.distributed, "is_available", return_value=False
), self.assertRaisesRegex(ValueError, "must be integers"):
_resolve_distributed_context(True)
for environment, message in [
({"RANK": "0", "WORLD_SIZE": "0"}, "greater than 0"),
({"RANK": "2", "WORLD_SIZE": "2"}, "0 <= rank"),
]:
with self.subTest(environment=environment), patch.dict(
os.environ, environment, clear=True
), patch.object(
torch.distributed, "is_available", return_value=False
), self.assertRaisesRegex(ValueError, message):
_resolve_distributed_context(True)
@unittest.skipUnless(
torch.distributed.is_available(), "torch.distributed is unavailable"
)
def test_torchrun_rank_and_worker_sharding(self):
script = os.path.join(
os.path.dirname(__file__), "torch_distributed_sharding_worker.py"
)
python_root = os.path.abspath(
os.path.join(os.path.dirname(__file__), "..", "..")
)
with tempfile.TemporaryDirectory() as output_dir:
env = os.environ.copy()
env["PYTHONPATH"] = os.pathsep.join(
filter(None, [python_root, env.get("PYTHONPATH")])
)
process = subprocess.run(
[
sys.executable,
"-m",
"torch.distributed.run",
"--standalone",
"--nproc-per-node=2",
script,
output_dir,
],
env=env,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True,
timeout=180,
)
self.assertEqual(
process.returncode,
0,
"torchrun failed:\n%s\n%s" % (
process.stdout, process.stderr
),
)
rows = []
for rank in range(2):
with open(
os.path.join(output_dir, "rank-%d.json" % rank),
encoding="utf-8",
) as result_file:
rows.extend(json.load(result_file))
split_ids = [row["split_id"] for row in rows]
self.assertCountEqual(split_ids, list(range(11)))
self.assertEqual(len(split_ids), len(set(split_ids)))
assignments = {}
for row in rows:
assignments.setdefault(
(row["rank"], row["worker"]), []
).append(row["split_id"])
self.assertEqual(
{key: sorted(values) for key, values in assignments.items()},
{
(0, 0): [0, 1, 2],
(0, 1): [3, 4, 5],
(1, 0): [6, 7, 8],
(1, 1): [9, 10],
},
)
class TorchBatchShuffleTest(unittest.TestCase):
@parameterized.expand([1, 3, 8, 50])
def test_columnar_shuffle_matches_existing_row_shuffle(self, buffer_size):
rows = [{"id": i, "nested": None if i % 4 == 0 else {"values": [i, None]},
"text": None if i % 3 == 0 else str(i)} for i in range(25)]
table = pa.Table.from_pylist(rows)
source = table.to_batches(max_chunksize=4)
source.insert(1, source[0].slice(0, 0))
actual = pa.Table.from_batches(list(_shuffle_record_batches(
iter(source), buffer_size, random.Random(17))))
legacy = TorchShuffledIterDataset(
_EmptyTableRead(), [], seed=17, buffer_size=buffer_size)
expected = list(legacy._iter_buffer_shuffled_rows(iter(rows), 0))
self.assertEqual(actual.schema, table.schema)
self.assertEqual(actual.to_pylist(), expected)
def test_shuffle_does_not_collect_whole_input(self):
consumed = []
def source():
for start in range(0, 300, 3):
consumed.extend(range(start, start + 3))
yield pa.record_batch([pa.array(range(start, start + 3))], names=["id"])
batches = _shuffle_record_batches(source(), 5, random.Random(0))
self.assertEqual(next(batches).num_rows, 5)
self.assertLessEqual(len(consumed), 12)
batches.close()
def test_shuffle_drains_before_arrow_offset_overflow(self):
table = pa.table({"text": [str(i).zfill(4) for i in range(12)]})
with patch("pypaimon.read.datasource.torch_dataset._MAX_ARROW_OFFSET", 12):
batches = list(_shuffle_record_batches(
iter(table.to_batches(max_chunksize=3)), 3, random.Random(3)))
self.assertEqual(sorted(row["text"] for b in batches for row in b.to_pylist()),
table["text"].to_pylist())
self.assertTrue(all(sum(len(v) for v in b["text"].to_pylist()) <= 12 for b in batches))
@parameterized.expand(["complete", "close", "read_error", "convert_error"])
def test_interleaved_readers_are_bounded_and_closed(self, mode):
opened, closed, active = [], [], set()
maximum = [0]
def read(dataset, splits):
split = splits[0]
opened.append(split)
active.add(split)
maximum[0] = max(maximum[0], len(active))
try:
for i in range(10):
if mode == "read_error" and split == 1 and i == 1:
raise RuntimeError("read failed")
yield pa.record_batch([pa.array([split * 10 + i])], names=["id"])
finally:
active.remove(split)
closed.append(split)
def convert(batch):
if mode == "convert_error":
raise RuntimeError("conversion failed")
return batch
dataset = TorchShuffledBatchIterDataset(
_EmptyTableRead(), list(range(8)),
buffer_size=3, max_buffer_input_splits=2, batch_size=2,
batch_format="torch", to_tensor_fn=convert)
with patch.object(TorchBatchIterDataset, "_arrow_batches_for_splits", read):
batches = iter(dataset)
if mode == "close":
next(batches)
batches.close()
elif mode.endswith("error"):
with self.assertRaisesRegex(RuntimeError, "failed"):
list(batches)
else:
rows = [row["id"] for batch in batches for row in batch.to_pylist()]
self.assertEqual(sorted(rows), list(range(80)))
self.assertFalse(active)
self.assertCountEqual(opened, closed)
self.assertEqual(maximum[0], 2)
def test_empty_input(self):
self.assertEqual(list(_shuffle_record_batches(iter([]), 5, random.Random(0))), [])
class TorchReadTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.tempdir = tempfile.mkdtemp()
cls.warehouse = os.path.join(cls.tempdir, 'warehouse')
cls.catalog = CatalogFactory.create({
'warehouse': cls.warehouse
})
cls.catalog.create_database('default', True)
cls.pa_schema = pa.schema([
('user_id', pa.int32()),
('item_id', pa.int64()),
('behavior', pa.string()),
('dt', pa.string())
])
cls.pk_pa_schema = pa.schema([
pa.field('user_id', pa.int32(), nullable=False),
('item_id', pa.int64()),
pa.field('behavior', pa.string(), nullable=False),
('dt', pa.string())
])
cls.expected = pa.Table.from_pydict({
'user_id': [1, 2, 3, 4, 5, 6, 7, 8],
'item_id': [1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008],
'behavior': ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'],
'dt': ['p1', 'p1', 'p2', 'p1', 'p2', 'p1', 'p2', 'p2'],
}, schema=cls.pa_schema)
@classmethod
def tearDownClass(cls):
shutil.rmtree(cls.tempdir, ignore_errors=True)
def test_nested_projection_uses_output_names_and_duplicate_values(self):
schema = pa.schema([
('id', pa.int32()), ('payload', pa.struct([('value', pa.int32()), ('unused', pa.string())]))])
self.catalog.create_table('default.torch_canonical_read_type', Schema.from_pyarrow_schema(
schema, options={'file.format': 'parquet'}), False)
table = self.catalog.get_table('default.torch_canonical_read_type')
builder = table.new_batch_write_builder()
writer, commit = builder.new_write(), builder.new_commit()
try:
writer.write_arrow(pa.Table.from_pylist([
{'id': 1, 'payload': {'value': 10, 'unused': 'x'}},
{'id': 2, 'payload': None}], schema=schema))
commit.commit(writer.prepare_commit())
finally:
writer.close()
commit.close()
builder = table.new_read_builder().with_projection(['payload.value', 'id', 'id'])
read, splits = builder.new_read(), builder.new_scan().plan().splits()
self.assertEqual([field.name for field in read.read_type], ['payload', 'id'])
expected = [{'payload_value': 10, 'id': 1, 'id__0': 1},
{'payload_value': None, 'id': 2, 'id__0': 2}]
for streaming in (False, True):
with self.subTest(streaming=streaming):
dataset = read.to_torch(splits, streaming=streaming)
self.assertEqual(sorted(list(dataset), key=lambda row: row['id']), expected)
dataset = read.to_torch(splits, streaming=True, batch_format='pyarrow')
actual = pa.Table.from_batches(list(dataset)).sort_by('id')
self.assertEqual(actual.to_pylist(), expected)
@parameterized.expand([True, False])
def test_torch_read(self, is_streaming: bool = False):
schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['user_id'])
self.catalog.create_table(f'default.test_torch_read_{str(is_streaming)}', schema, False)
table = self.catalog.get_table(f'default.test_torch_read_{str(is_streaming)}')
self._write_test_table(table)
read_builder = table.new_read_builder().with_projection(['user_id', 'behavior'])
table_scan = read_builder.new_scan()
table_read = read_builder.new_read()
splits = table_scan.plan().splits()
dataset = table_read.to_torch(splits, streaming=is_streaming)
dataloader = DataLoader(
dataset,
batch_size=2,
num_workers=2,
shuffle=False
)
# Collect all data from dataloader
all_user_ids = []
all_behaviors = []
for batch_idx, batch_data in enumerate(dataloader):
user_ids = batch_data['user_id'].tolist()
behaviors = batch_data['behavior']
all_user_ids.extend(user_ids)
all_behaviors.extend(behaviors)
# Sort by user_id for comparison
sorted_data = sorted(zip(all_user_ids, all_behaviors), key=lambda x: x[0])
sorted_user_ids = [x[0] for x in sorted_data]
sorted_behaviors = [x[1] for x in sorted_data]
# Expected data (sorted by user_id)
expected_user_ids = [1, 2, 3, 4, 5, 6, 7, 8]
expected_behaviors = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h']
# Verify results
self.assertEqual(sorted_user_ids, expected_user_ids,
f"User IDs mismatch. Expected {expected_user_ids}, got {sorted_user_ids}")
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)
table = self.catalog.get_table('default.test_torch_prefetch_concurrency')
self._write_test_table(table)
read_builder = table.new_read_builder().with_projection(['user_id', 'behavior'])
table_scan = read_builder.new_scan()
table_read = read_builder.new_read()
splits = table_scan.plan().splits()
self.assertGreater(len(splits), 0, "Need at least one split to test prefetch")
dataset = table_read.to_torch(splits, streaming=True, prefetch_concurrency=4)
dataloader = DataLoader(
dataset,
batch_size=2,
num_workers=0,
shuffle=False
)
all_user_ids = []
all_behaviors = []
for batch_data in dataloader:
all_user_ids.extend(batch_data['user_id'].tolist())
all_behaviors.extend(batch_data['behavior'])
sorted_data = sorted(zip(all_user_ids, all_behaviors), key=lambda x: x[0])
sorted_user_ids = [x[0] for x in sorted_data]
sorted_behaviors = [x[1] for x in sorted_data]
expected_user_ids = [1, 2, 3, 4, 5, 6, 7, 8]
expected_behaviors = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h']
self.assertEqual(len(all_user_ids), 8, "Should read 8 rows with prefetch_concurrency")
self.assertEqual(sorted_user_ids, expected_user_ids)
self.assertEqual(sorted_behaviors, expected_behaviors)
def test_torch_streaming_pyarrow_batches(self):
schema = Schema.from_pyarrow_schema(
self.pa_schema, partition_keys=['user_id']
)
self.catalog.create_table(
'default.test_torch_pyarrow_batches', schema, False
)
table = self.catalog.get_table(
'default.test_torch_pyarrow_batches'
)
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=True,
batch_format='pyarrow',
batch_size=3,
)
dataloader = DataLoader(
dataset,
batch_size=None,
num_workers=2,
shuffle=False,
)
batches = list(dataloader)
self.assertTrue(batches)
self.assertTrue(
all(isinstance(batch, pa.RecordBatch) for batch in batches)
)
self.assertTrue(all(0 < batch.num_rows <= 3 for batch in batches))
result = pa.Table.from_batches(batches).sort_by('user_id').to_pydict()
self.assertEqual(result['user_id'], list(range(1, 9)))
self.assertEqual(result['behavior'], list('abcdefgh'))
def test_torch_streaming_tensor_batches(self):
schema = Schema.from_pyarrow_schema(
self.pa_schema, partition_keys=['user_id']
)
self.catalog.create_table(
'default.test_torch_tensor_batches', schema, False
)
table = self.catalog.get_table(
'default.test_torch_tensor_batches'
)
self._write_test_table(table)
read_builder = table.new_read_builder().with_projection(
['user_id', 'item_id']
)
splits = read_builder.new_scan().plan().splits()
dataset = read_builder.new_read().to_torch(
splits,
streaming=True,
batch_format='torch',
batch_size=3,
)
batches = list(dataset)
self.assertEqual([len(batch['user_id']) for batch in batches], [3, 3, 2])
self.assertTrue(
all(batch['user_id'].dtype == torch.int32 for batch in batches)
)
self.assertTrue(
all(batch['item_id'].dtype == torch.int64 for batch in batches)
)
user_ids = torch.cat(
[batch['user_id'] for batch in batches]
).sort().values.tolist()
self.assertEqual(user_ids, list(range(1, 9)))
def test_torch_streaming_batches_respect_limit(self):
schema = Schema.from_pyarrow_schema(
self.pa_schema, partition_keys=['user_id']
)
self.catalog.create_table(
'default.test_torch_batch_limit', schema, False
)
table = self.catalog.get_table('default.test_torch_batch_limit')
self._write_test_table(table)
read_builder = table.new_read_builder().with_projection(
['user_id']
).with_limit(5)
splits = read_builder.new_scan().plan().splits()
dataset = read_builder.new_read().to_torch(
splits,
streaming=True,
batch_format='pyarrow',
batch_size=3,
)
batches = list(dataset)
self.assertEqual([batch.num_rows for batch in batches], [3, 2])
def test_torch_streaming_batches_respect_limit_with_workers(self):
schema = Schema.from_pyarrow_schema(
self.pa_schema, partition_keys=['user_id']
)
self.catalog.create_table(
'default.test_torch_batch_worker_limit', schema, False
)
table = self.catalog.get_table(
'default.test_torch_batch_worker_limit'
)
self._write_test_table(table)
predicate = (
table.new_read_builder().new_predicate_builder()
.greater_than('item_id', 0)
)
read_builder = (
table.new_read_builder()
.with_filter(predicate)
.with_projection(['user_id'])
.with_limit(5)
)
splits = read_builder.new_scan().plan().splits()
self.assertGreater(len(splits), 1)
dataset = read_builder.new_read().to_torch(
splits,
streaming=True,
batch_format='pyarrow',
batch_size=3,
)
self.assertEqual(
dataset._worker_splits(SimpleNamespace(id=1, num_workers=2)),
[],
)
batches = list(DataLoader(
dataset, batch_size=None, num_workers=2
))
user_ids = [
value
for batch in batches
for value in batch.column('user_id').to_pylist()
]
self.assertEqual(len(user_ids), 5)
self.assertEqual(len(set(user_ids)), 5)
def test_non_binding_limit_preserves_worker_splits(self):
schema = Schema.from_pyarrow_schema(
self.pa_schema, partition_keys=['user_id']
)
self.catalog.create_table(
'default.test_torch_non_binding_limit', schema, False
)
table = self.catalog.get_table(
'default.test_torch_non_binding_limit'
)
self._write_test_table(table)
read_builder = table.new_read_builder().with_limit(1000)
splits = read_builder.new_scan().plan().splits()
self.assertGreater(len(splits), 1)
table_read = read_builder.new_read()
for batch_format in ['row', 'pyarrow']:
dataset = table_read.to_torch(
splits,
streaming=True,
batch_format=batch_format,
)
assigned = [
dataset._worker_splits(
SimpleNamespace(id=worker_id, num_workers=2)
)
for worker_id in range(2)
]
self.assertTrue(all(assigned))
self.assertCountEqual(
[id(split) for group in assigned for split in group],
[id(split) for split in splits],
)
def test_non_binding_limit_uses_merged_row_counts(self):
from pypaimon.read.datasource.torch_dataset import TorchIterDataset
table_read = _EmptyTableRead(limit=8)
splits = [
SimpleNamespace(row_count=10, merged_row_count=lambda: 4),
SimpleNamespace(row_count=10, merged_row_count=lambda: 4),
]
dataset = TorchIterDataset(table_read, splits)
assigned = [
dataset._worker_splits(
SimpleNamespace(id=worker_id, num_workers=2)
)
for worker_id in range(2)
]
self.assertTrue(all(assigned))
self.assertCountEqual(
[id(split) for group in assigned for split in group],
[id(split) for split in splits],
)
def test_torch_batch_sizing_respects_arrow_offset_limit(self):
from pypaimon.read.datasource.torch_dataset import (
_sized_record_batches)
batches = iter([
pa.record_batch([pa.array(['aaaa'])], names=['value']),
pa.record_batch([pa.array(['bbbb'])], names=['value']),
])
with patch(
'pypaimon.read.datasource.torch_dataset._MAX_ARROW_OFFSET', 4
):
actual = list(_sized_record_batches(batches, batch_size=2))
self.assertEqual(
[batch.column('value').to_pylist() for batch in actual],
[['aaaa'], ['bbbb']],
)
def test_default_tensor_converter_supports_fixed_size_list(self):
from pypaimon.read.datasource.torch_dataset import _default_to_tensor
values = pa.array([1, 2, 3, 4, 5, 6], type=pa.int32())
features = pa.FixedSizeListArray.from_arrays(values, 3)
batch = pa.RecordBatch.from_arrays([features], ['features'])
result = _default_to_tensor(batch)
self.assertEqual(result['features'].dtype, torch.int32)
self.assertEqual(result['features'].tolist(), [[1, 2, 3], [4, 5, 6]])
def test_torch_streaming_custom_tensor_conversion(self):
schema = Schema.from_pyarrow_schema(self.pa_schema)
self.catalog.create_table(
'default.test_torch_custom_tensor_batch', schema, False
)
table = self.catalog.get_table(
'default.test_torch_custom_tensor_batch'
)
self._write_test_table(table)
read_builder = table.new_read_builder().with_projection(
['user_id', 'behavior']
)
splits = read_builder.new_scan().plan().splits()
def to_tensor(batch):
return {
'user_id': torch.from_numpy(
batch.column('user_id').to_numpy(zero_copy_only=False)
),
'behavior': batch.column('behavior').to_pylist(),
}
dataset = read_builder.new_read().to_torch(
splits,
streaming=True,
batch_format='torch',
batch_size=5,
to_tensor_fn=to_tensor,
)
batches = list(dataset)
self.assertEqual([len(batch['user_id']) for batch in batches], [5, 3])
self.assertEqual(
sorted(value for batch in batches for value in batch['behavior']),
list('abcdefgh'),
)
default_dataset = read_builder.new_read().to_torch(
splits,
streaming=True,
batch_format='torch',
)
with self.assertRaisesRegex(ValueError, "batch_format='pyarrow'"):
next(iter(default_dataset))
def test_torch_batch_options_validation(self):
schema = Schema.from_pyarrow_schema(self.pa_schema)
self.catalog.create_table(
'default.test_torch_batch_validation', schema, False
)
table = self.catalog.get_table(
'default.test_torch_batch_validation'
)
self._write_test_table(table)
read_builder = table.new_read_builder().with_projection(['user_id'])
splits = read_builder.new_scan().plan().splits()
table_read = read_builder.new_read()
with self.assertRaisesRegex(ValueError, 'batch_format must be one of'):
table_read.to_torch(
splits, streaming=True, batch_format='numpy'
)
with self.assertRaisesRegex(ValueError, 'requires streaming=True'):
table_read.to_torch(splits, batch_format='pyarrow')
with self.assertRaisesRegex(ValueError, 'requires streaming=True'):
table_read.to_torch(splits, auto_detect_rank=True)
with self.assertRaisesRegex(ValueError, 'requires streaming=True'):
table_read.to_torch(
splits,
sharding_rank=0,
sharding_world_size=1,
)
with self.assertRaisesRegex(ValueError, 'batch_size must be'):
table_read.to_torch(
splits,
streaming=True,
batch_format='torch',
batch_size=0,
)
with self.assertRaisesRegex(ValueError, 'batch_size requires'):
table_read.to_torch(splits, streaming=True, batch_size=2)
with self.assertRaisesRegex(ValueError, 'buffer_size'):
table_read.to_torch(
splits,
streaming=True,
batch_format='torch',
shuffle=True,
buffer_size=0,
)
for invalid in [0, -1, 1.9, True, 2]:
with self.subTest(prefetch_concurrency=invalid):
with self.assertRaisesRegex(
ValueError, 'prefetch_concurrency'
):
table_read.to_torch(
splits,
streaming=True,
batch_format='pyarrow',
prefetch_concurrency=invalid,
)
def test_torch_distributed_sharding_public_api(self):
schema = Schema.from_pyarrow_schema(
self.pa_schema, partition_keys=['user_id']
)
self.catalog.create_table(
'default.test_torch_distributed_api', schema, False
)
table = self.catalog.get_table(
'default.test_torch_distributed_api'
)
self._write_test_table(table)
read_builder = table.new_read_builder().with_projection(['user_id'])
splits = read_builder.new_scan().plan().splits()
table_read = read_builder.new_read()
with patch(
"pypaimon.read.datasource.torch_dataset."
"_resolve_distributed_context",
return_value=(1, 2),
):
datasets = [
table_read.to_torch(
splits,
streaming=True,
batch_format=batch_format,
shuffle=shuffle,
auto_detect_rank=True,
)
for batch_format, shuffle in [
('row', False),
('row', True),
('pyarrow', False),
('pyarrow', True),
]
]
expected = splits[(len(splits) + 1) // 2:]
for dataset in datasets:
self.assertTrue(dataset.auto_detect_rank)
self.assertEqual(dataset._worker_splits(None), expected)
pre_sharded = splits[::2]
with patch.dict(
os.environ, {"RANK": "1", "WORLD_SIZE": "2"}, clear=True
):
dataset = table_read.to_torch(pre_sharded, streaming=True)
self.assertFalse(dataset.auto_detect_rank)
self.assertEqual(dataset._worker_splits(None), pre_sharded)
dataset = table_read.to_torch(
splits,
streaming=True,
sharding_rank=1,
sharding_world_size=2,
)
self.assertFalse(dataset.auto_detect_rank)
self.assertEqual(
dataset._worker_splits(None),
splits[(len(splits) + 1) // 2:],
)
self.assertIsNotNone(table_read.to_torch(splits))
def test_blob_torch_read(self):
"""Test end-to-end blob functionality using blob descriptors."""
import random
from pypaimon import Schema
from pypaimon.table.row.blob import BlobDescriptor
# Create schema with blob column
pa_schema = pa.schema([
('id', pa.int32()),
('picture', pa.large_binary()),
])
schema = Schema.from_pyarrow_schema(
pa_schema,
options={
'row-tracking.enabled': 'true',
'data-evolution.enabled': 'true',
'blob-as-descriptor': 'true'
}
)
# Create table
self.catalog.create_table('default.test_blob_torch_read', schema, False)
table: FileStoreTable = self.catalog.get_table('default.test_blob_torch_read')
# Create test blob data (1MB)
blob_data = bytearray(1024 * 1024)
random.seed(42) # For reproducible tests
for i in range(len(blob_data)):
blob_data[i] = random.randint(0, 255)
blob_data = bytes(blob_data)
# Create external blob file
external_blob_path = os.path.join(self.tempdir, 'external_blob')
with open(external_blob_path, 'wb') as f:
f.write(blob_data)
# Create blob descriptor pointing to external file
blob_descriptor = BlobDescriptor(external_blob_path, 0, len(blob_data))
# Create test data with blob descriptor
test_data = pa.Table.from_pydict({
'id': [1],
'picture': [blob_descriptor.serialize()]
}, schema=pa_schema)
# Write data using table API
write_builder = table.new_batch_write_builder()
writer = write_builder.new_write()
writer.write_arrow(test_data)
# Commit the data
commit_messages = writer.prepare_commit()
commit = write_builder.new_commit()
commit.commit(commit_messages)
# Read data back
read_builder = table.new_read_builder()
table_scan = read_builder.new_scan()
table_read = read_builder.new_read()
result = table_read.to_torch(table_scan.plan().splits())
dataloader = DataLoader(
result,
batch_size=1,
num_workers=0,
shuffle=False
)
# Collect and verify data
all_ids = []
all_pictures = []
for batch_idx, batch_data in enumerate(dataloader):
ids = batch_data['id'].tolist()
pictures = batch_data['picture']
all_ids.extend(ids)
all_pictures.extend(pictures)
# Verify results
self.assertEqual(len(all_ids), 1, "Should have exactly 1 row")
self.assertEqual(all_ids[0], 1, "ID should be 1")
# Verify blob descriptor
picture_bytes = all_pictures[0]
self.assertIsInstance(picture_bytes, bytes, "Picture should be bytes")
# Deserialize and verify blob descriptor
from pypaimon.table.row.blob import BlobDescriptor
read_blob_descriptor = BlobDescriptor.deserialize(picture_bytes)
self.assertEqual(read_blob_descriptor.length, len(blob_data),
f"Blob length mismatch. Expected {len(blob_data)}, got {read_blob_descriptor.length}")
self.assertGreaterEqual(read_blob_descriptor.offset, 0, "Offset should be non-negative")
# Read and verify blob content
from pypaimon.common.uri_reader import UriReaderFactory
from pypaimon.common.options.config import CatalogOptions
from pypaimon.table.row.blob import Blob
catalog_options = {CatalogOptions.WAREHOUSE.key(): self.warehouse}
uri_reader_factory = UriReaderFactory(catalog_options)
uri_reader = uri_reader_factory.create(read_blob_descriptor.uri)
blob = Blob.from_descriptor(uri_reader, read_blob_descriptor)
# Verify blob data matches original
read_blob_data = blob.to_data()
self.assertEqual(len(read_blob_data), len(blob_data),
f"Blob data length mismatch. Expected {len(blob_data)}, got {len(read_blob_data)}")
self.assertEqual(read_blob_data, blob_data, "Blob data content should match original")
print(f"✓ Blob torch read test passed: Successfully read and verified {len(blob_data)} bytes of blob data")
def test_video_frame_rows_through_streaming_dataloader(self):
from pypaimon.table.row.blob import Blob, VideoFrameDescriptor
pa_schema = pa.schema([
('episode_id', pa.int64()),
('video', pa.large_binary()),
])
schema = Schema.from_pyarrow_schema(
pa_schema,
options={
'row-tracking.enabled': 'true',
'data-evolution.enabled': 'true',
'video-frame-field': 'video',
# ScanQuery.to_torch must override this read setting.
'blob-as-descriptor': 'false',
},
)
identifier = 'default.test_shared_video_torch_read'
self.catalog.create_table(identifier, schema, False)
raw_table = self.catalog.get_table(identifier)
table = MultimodalTable(self.catalog, identifier, raw_table)
video_path = os.path.join(self.tempdir, 'shared-video.mp4')
video_bytes = b'one-physical-video'
with open(video_path, 'wb') as output:
output.write(video_bytes)
video = Blob.from_local(video_path)
table.add_video(video, [{'episode_id': 7} for _ in range(8)])
dataset = table.scan().select([
'episode_id', 'video'
]).to_torch(streaming=True)
rows = []
for batch in DataLoader(
dataset, batch_size=2, num_workers=2, shuffle=False):
rows.extend(zip(
batch['episode_id'].tolist(),
batch['video'],
))
descriptors = [
VideoFrameDescriptor.deserialize(row[1])
for row in rows
]
descriptors.sort(key=lambda value: value.frame_index)
self.assertEqual(list(range(8)), [d.frame_index for d in descriptors])
self.assertTrue(all(
value.payload_descriptor == descriptors[0].payload_descriptor
for value in descriptors
))
self.assertTrue(descriptors[0].uri.endswith('.video'))
self.assertEqual(len(video_bytes), descriptors[0].length)
def test_torch_read_pk_table(self):
"""Test torch read with primary key table."""
# Create PK table with user_id as primary key and behavior as partition key
schema = Schema.from_pyarrow_schema(
self.pa_schema,
primary_keys=['user_id', 'behavior'],
partition_keys=['behavior'],
options={'bucket': 2}
)
self.catalog.create_table('default.test_pk_table', schema, False)
table = self.catalog.get_table('default.test_pk_table')
self._write_test_table(table)
read_builder = table.new_read_builder().with_projection(['user_id', 'behavior'])
table_scan = read_builder.new_scan()
table_read = read_builder.new_read()
splits = table_scan.plan().splits()
dataset = table_read.to_torch(splits, streaming=True)
dataloader = DataLoader(
dataset,
batch_size=2,
num_workers=3,
shuffle=False
)
# Collect all data from dataloader
all_user_ids = []
all_behaviors = []
for batch_idx, batch_data in enumerate(dataloader):
user_ids = batch_data['user_id'].tolist()
behaviors = batch_data['behavior']
all_user_ids.extend(user_ids)
all_behaviors.extend(behaviors)
# Sort by user_id for comparison
sorted_data = sorted(zip(all_user_ids, all_behaviors), key=lambda x: x[0])
sorted_user_ids = [x[0] for x in sorted_data]
sorted_behaviors = [x[1] for x in sorted_data]
# Expected data (sorted by user_id)
expected_user_ids = [1, 2, 3, 4, 5, 6, 7, 8]
expected_behaviors = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h']
# Verify results
self.assertEqual(sorted_user_ids, expected_user_ids,
f"User IDs mismatch. Expected {expected_user_ids}, got {sorted_user_ids}")
self.assertEqual(sorted_behaviors, expected_behaviors,
f"Behaviors mismatch. Expected {expected_behaviors}, got {sorted_behaviors}")
print(f"✓ PK table test passed: Successfully read {len(all_user_ids)} rows with correct data")
def test_torch_read_large_append_table(self):
"""Test torch read with large data volume on append-only table."""
# Create append-only table
schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt'])
self.catalog.create_table('default.test_large_append', schema, False)
table = self.catalog.get_table('default.test_large_append')
# Write large amount of data
write_builder = table.new_batch_write_builder()
total_rows = 100000
batch_size = 10000
num_batches = total_rows // batch_size
print(f"\n{'=' * 60}")
print(f"Writing {total_rows} rows to append-only table...")
print(f"{'=' * 60}")
for batch_idx in range(num_batches):
table_write = write_builder.new_write()
table_commit = write_builder.new_commit()
start_id = batch_idx * batch_size + 1
end_id = start_id + batch_size
data = {
'user_id': list(range(start_id, end_id)),
'item_id': [1000 + i for i in range(start_id, end_id)],
'behavior': [chr(ord('a') + (i % 26)) for i in range(batch_size)],
'dt': [f'p{i % 4}' for i in range(batch_size)],
}
pa_table = pa.Table.from_pydict(data, schema=self.pa_schema)
table_write.write_arrow(pa_table)
table_commit.commit(table_write.prepare_commit())
table_write.close()
table_commit.close()
if (batch_idx + 1) % 2 == 0:
print(f" Written {(batch_idx + 1) * batch_size} rows...")
# Read data using torch
print(f"\nReading {total_rows} rows using Torch DataLoader...")
read_builder = table.new_read_builder().with_projection(['user_id', 'behavior'])
table_scan = read_builder.new_scan()
table_read = read_builder.new_read()
splits = table_scan.plan().splits()
print(f"Total splits: {len(splits)}")
dataset = table_read.to_torch(splits, streaming=True)
dataloader = DataLoader(
dataset,
batch_size=1000,
num_workers=4,
shuffle=False
)
# Collect all data
all_user_ids = []
batch_count = 0
for batch_idx, batch_data in enumerate(dataloader):
batch_count += 1
user_ids = batch_data['user_id'].tolist()
all_user_ids.extend(user_ids)
if (batch_idx + 1) % 20 == 0:
print(f" Read {len(all_user_ids)} rows...")
all_user_ids.sort()
# Verify data
self.assertEqual(len(all_user_ids), total_rows,
f"Row count mismatch. Expected {total_rows}, got {len(all_user_ids)}")
self.assertEqual(all_user_ids, list(range(1, total_rows + 1)),
f"Row count mismatch. Expected {total_rows}, got {len(all_user_ids)}")
print(f"\n{'=' * 60}")
print("✓ Large append table test passed!")
print(f" Total rows: {total_rows}")
print(f" Total batches: {batch_count}")
print(f"{'=' * 60}\n")
def test_torch_read_large_pk_table(self):
"""Test torch read with large data volume on primary key table."""
# Create PK table
large_pk_pa_schema = pa.schema([
pa.field('user_id', pa.int32(), nullable=False),
('item_id', pa.int64()),
('behavior', pa.string()),
('dt', pa.string())
])
schema = Schema.from_pyarrow_schema(
self.pa_schema,
primary_keys=['user_id'],
partition_keys=['dt'],
options={'bucket': '4'}
)
self.catalog.create_table('default.test_large_pk', schema, False)
table = self.catalog.get_table('default.test_large_pk')
# Write large amount of data
write_builder = table.new_batch_write_builder()
total_rows = 100000
batch_size = 10000
num_batches = total_rows // batch_size
print(f"\n{'=' * 60}")
print(f"Writing {total_rows} rows to PK table...")
print(f"{'=' * 60}")
for batch_idx in range(num_batches):
table_write = write_builder.new_write()
table_commit = write_builder.new_commit()
start_id = batch_idx * batch_size + 1
end_id = start_id + batch_size
data = {
'user_id': list(range(start_id, end_id)),
'item_id': [1000 + i for i in range(start_id, end_id)],
'behavior': [chr(ord('a') + (i % 26)) for i in range(batch_size)],
'dt': [f'p{i % 4}' for i in range(batch_size)],
}
pa_table = pa.Table.from_pydict(data, schema=large_pk_pa_schema)
table_write.write_arrow(pa_table)
table_commit.commit(table_write.prepare_commit())
table_write.close()
table_commit.close()
if (batch_idx + 1) % 2 == 0:
print(f" Written {(batch_idx + 1) * batch_size} rows...")
# Read data using torch
print(f"\nReading {total_rows} rows using Torch DataLoader...")
read_builder = table.new_read_builder()
table_scan = read_builder.new_scan()
table_read = read_builder.new_read()
splits = table_scan.plan().splits()
print(f"Total splits: {len(splits)}")
dataset = table_read.to_torch(splits, streaming=True)
dataloader = DataLoader(
dataset,
batch_size=1000,
num_workers=8,
shuffle=False
)
# Collect all data
all_user_ids = []
batch_count = 0
for batch_idx, batch_data in enumerate(dataloader):
batch_count += 1
user_ids = batch_data['user_id'].tolist()
all_user_ids.extend(user_ids)
if (batch_idx + 1) % 20 == 0:
print(f" Read {len(all_user_ids)} rows...")
all_user_ids.sort()
# Verify data
self.assertEqual(len(all_user_ids), total_rows,
f"Row count mismatch. Expected {total_rows}, got {len(all_user_ids)}")
self.assertEqual(all_user_ids, list(range(1, total_rows + 1)),
f"Row count mismatch. Expected {total_rows}, got {len(all_user_ids)}")
print(f"\n{'=' * 60}")
print("✓ Large PK table test passed!")
print(f" Total rows: {total_rows}")
print(f" Total batches: {batch_count}")
print(" Primary key uniqueness: ✓")
print(f"{'=' * 60}\n")
def test_torch_read_with_predicate(self):
"""Test torch read with predicate filtering."""
schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['user_id'])
self.catalog.create_table('default.test_predicate', schema, False)
table = self.catalog.get_table('default.test_predicate')
self._write_test_table(table)
# Test case 1: Filter by user_id > 4
print(f"\n{'=' * 60}")
print("Test Case 1: user_id > 4")
print(f"{'=' * 60}")
predicate_builder = table.new_read_builder().new_predicate_builder()
predicate = predicate_builder.greater_than('user_id', 4)
read_builder = table.new_read_builder().with_filter(predicate)
table_scan = read_builder.new_scan()
table_read = read_builder.new_read()
splits = table_scan.plan().splits()
dataset = table_read.to_torch(splits, streaming=True)
dataloader = DataLoader(
dataset,
batch_size=2,
num_workers=0,
shuffle=False
)
all_user_ids = []
for batch_idx, batch_data in enumerate(dataloader):
user_ids = batch_data['user_id'].tolist()
all_user_ids.extend(user_ids)
all_user_ids.sort()
expected_user_ids = [5, 6, 7, 8]
self.assertEqual(all_user_ids, expected_user_ids,
f"User IDs mismatch. Expected {expected_user_ids}, got {all_user_ids}")
print(f"✓ Filtered {len(all_user_ids)} rows: {all_user_ids}")
# Test case 2: Filter by user_id <= 3
print(f"\n{'=' * 60}")
print("Test Case 2: user_id <= 3")
print(f"{'=' * 60}")
predicate = predicate_builder.less_or_equal('user_id', 3)
read_builder = table.new_read_builder().with_filter(predicate)
table_scan = read_builder.new_scan()
table_read = read_builder.new_read()
splits = table_scan.plan().splits()
dataset = table_read.to_torch(splits, streaming=True)
dataloader = DataLoader(
dataset,
batch_size=2,
num_workers=0,
shuffle=False
)
all_user_ids = []
for batch_idx, batch_data in enumerate(dataloader):
user_ids = batch_data['user_id'].tolist()
all_user_ids.extend(user_ids)
all_user_ids.sort()
expected_user_ids = [1, 2, 3]
self.assertEqual(all_user_ids, expected_user_ids,
f"User IDs mismatch. Expected {expected_user_ids}, got {all_user_ids}")
print(f"✓ Filtered {len(all_user_ids)} rows: {all_user_ids}")
# Test case 3: Filter by behavior = 'a'
print(f"\n{'=' * 60}")
print("Test Case 3: behavior = 'a'")
print(f"{'=' * 60}")
predicate = predicate_builder.equal('behavior', 'a')
read_builder = table.new_read_builder().with_filter(predicate)
table_scan = read_builder.new_scan()
table_read = read_builder.new_read()
splits = table_scan.plan().splits()
dataset = table_read.to_torch(splits, streaming=True)
dataloader = DataLoader(
dataset,
batch_size=2,
num_workers=0,
shuffle=False
)
all_user_ids = []
all_behaviors = []
for batch_idx, batch_data in enumerate(dataloader):
user_ids = batch_data['user_id'].tolist()
behaviors = batch_data['behavior']
all_user_ids.extend(user_ids)
all_behaviors.extend(behaviors)
expected_user_ids = [1]
expected_behaviors = ['a']
self.assertEqual(all_user_ids, expected_user_ids,
f"User IDs mismatch. Expected {expected_user_ids}, got {all_user_ids}")
self.assertEqual(all_behaviors, expected_behaviors,
f"Behaviors mismatch. Expected {expected_behaviors}, got {all_behaviors}")
print(f"✓ Filtered {len(all_user_ids)} rows: user_ids={all_user_ids}, behaviors={all_behaviors}")
# Test case 4: Filter by user_id IN (2, 4, 6)
print(f"\n{'=' * 60}")
print("Test Case 4: user_id IN (2, 4, 6)")
print(f"{'=' * 60}")
predicate = predicate_builder.is_in('user_id', [2, 4, 6])
read_builder = table.new_read_builder().with_filter(predicate)
table_scan = read_builder.new_scan()
table_read = read_builder.new_read()
splits = table_scan.plan().splits()
dataset = table_read.to_torch(splits, streaming=True)
dataloader = DataLoader(
dataset,
batch_size=2,
num_workers=0,
shuffle=False
)
all_user_ids = []
for batch_idx, batch_data in enumerate(dataloader):
user_ids = batch_data['user_id'].tolist()
all_user_ids.extend(user_ids)
all_user_ids.sort()
expected_user_ids = [2, 4, 6]
self.assertEqual(all_user_ids, expected_user_ids,
f"User IDs mismatch. Expected {expected_user_ids}, got {all_user_ids}")
print(f"✓ Filtered {len(all_user_ids)} rows: {all_user_ids}")
# Test case 5: Combined filter (user_id > 2 AND user_id < 7)
print(f"\n{'=' * 60}")
print("Test Case 5: user_id > 2 AND user_id < 7")
print(f"{'=' * 60}")
predicate1 = predicate_builder.greater_than('user_id', 2)
predicate2 = predicate_builder.less_than('user_id', 7)
combined_predicate = predicate_builder.and_predicates([predicate1, predicate2])
read_builder = table.new_read_builder().with_filter(combined_predicate)
table_scan = read_builder.new_scan()
table_read = read_builder.new_read()
splits = table_scan.plan().splits()
dataset = table_read.to_torch(splits, streaming=True)
dataloader = DataLoader(
dataset,
batch_size=2,
num_workers=0,
shuffle=False
)
all_user_ids = []
for batch_idx, batch_data in enumerate(dataloader):
user_ids = batch_data['user_id'].tolist()
all_user_ids.extend(user_ids)
all_user_ids.sort()
expected_user_ids = [3, 4, 5, 6]
self.assertEqual(all_user_ids, expected_user_ids,
f"User IDs mismatch. Expected {expected_user_ids}, got {all_user_ids}")
print(f"✓ Filtered {len(all_user_ids)} rows: {all_user_ids}")
print(f"\n{'=' * 60}")
print("✓ All predicate test cases passed!")
print(f"{'=' * 60}\n")
@parameterized.expand([(1, None), (7, None), (7, 3), (11, 17)])
def test_torch_batch_shuffle_formats_and_epochs(self, buffer_size, batch_size):
table = self._create_shuffle_append_table(
'default.batch_shuffle_%s_%s' % (buffer_size, batch_size))
builder = table.new_read_builder().with_projection(['user_id', 'item_id'])
splits = builder.new_scan().plan().splits()
read = builder.new_read()
outputs = []
for batch_format in ['pyarrow', 'torch']:
dataset = read.to_torch(
splits, streaming=True, shuffle=True, batch_format=batch_format,
batch_size=batch_size, seed=23, buffer_size=buffer_size,
max_buffer_input_splits=3)
def collect():
batches = list(dataset)
if batch_format == 'pyarrow':
return pa.Table.from_batches(batches).to_pydict(), [b.num_rows for b in batches]
return {name: torch.cat([b[name] for b in batches]).tolist()
for name in ['user_id', 'item_id']}, [len(b['user_id']) for b in batches]
with patch.object(TableRead, 'to_arrow', side_effect=AssertionError('whole table')), \
patch.object(TableRead, 'to_iterator', side_effect=AssertionError('row conversion')):
data, sizes = collect()
self.assertEqual(data, collect()[0])
self.assertEqual(sorted(data['user_id']), list(range(80)))
self.assertEqual(data['item_id'], [1000 + i for i in data['user_id']])
self.assertTrue(all(size <= (batch_size or buffer_size) for size in sizes))
if batch_size:
self.assertTrue(all(size == batch_size for size in sizes[:-1]))
outputs.append(data)
dataset.set_epoch(1)
if buffer_size > 1:
self.assertNotEqual(data, collect()[0])
dataset.set_epoch(0)
self.assertEqual(data, collect()[0])
self.assertEqual(outputs[0], outputs[1])
def test_torch_batch_shuffle_filter_limit_and_empty(self):
table = self._create_shuffle_append_table('default.batch_shuffle_limit')
builder = table.new_read_builder().with_projection(['user_id'])
predicate = builder.new_predicate_builder().greater_or_equal('user_id', 10)
builder.with_filter(predicate).with_limit(9)
splits = builder.new_scan().plan().splits()
read = builder.new_read()
expected = read.to_arrow(splits)['user_id'].to_pylist()
dataset = read.to_torch(splits, streaming=True, shuffle=True, batch_format='pyarrow',
batch_size=4, buffer_size=3, max_buffer_input_splits=3)
actual = [row['user_id'] for batch in dataset for row in batch.to_pylist()]
self.assertCountEqual(actual, expected)
for limit in [0, 9]:
read.limit = limit
empty = read.to_torch([], streaming=True, shuffle=True, batch_format='pyarrow', buffer_size=3)
self.assertEqual(list(empty), [])
def test_torch_batch_shuffle_persistent_spawn_workers(self):
table = self._create_shuffle_append_table('default.batch_shuffle_spawn', partition_keys=['dt'])
builder = table.new_read_builder().with_projection(['user_id'])
splits = builder.new_scan().plan().splits()
self.assertGreaterEqual(len(splits), 4)
all_ids = []
for rank in range(2):
dataset = builder.new_read().to_torch(
splits, streaming=True, shuffle=True, batch_format='torch',
batch_size=4, buffer_size=7, seed=23,
sharding_rank=rank, sharding_world_size=2)
loader = DataLoader(dataset, batch_size=None, num_workers=2,
persistent_workers=True, multiprocessing_context='spawn')
try:
epoch0 = self._collect_torch_user_ids_from_dataloader(loader)
self.assertEqual(epoch0, self._collect_torch_user_ids_from_dataloader(loader))
dataset.set_epoch(1)
epoch1 = self._collect_torch_user_ids_from_dataloader(loader)
self.assertCountEqual(epoch0, epoch1)
self.assertNotEqual(epoch0, epoch1)
all_ids.extend(epoch0)
finally:
if loader._iterator is not None:
loader._iterator._shutdown_workers()
self.assertEqual(sorted(all_ids), list(range(80)))
def test_torch_batch_shuffle_vectors_deletions_and_snapshot(self):
import pypaimon.multimodal as pm
schema = pa.schema([('id', pa.int64()), ('vector', pa.list_(pa.float32(), 2))])
table = pm.connect(options={'warehouse': self.warehouse}).create_table(
'batch_shuffle_vectors', schema=schema, options={'file.format': 'parquet'})
table.add(pa.table({'id': list(range(8)), 'vector': [[i, i + 1] for i in range(8)]}, schema=schema))
snapshot = table.raw_table.snapshot_manager().get_latest_snapshot().id
table.delete('id = 2')
for snapshot_id, expected in [(snapshot, list(range(1, 8))), (None, [1, 3, 4, 5, 6, 7])]:
dataset = table.scan(snapshot_id=snapshot_id).where('id >= 1').to_torch(
streaming=True, batch_format='torch', batch_size=3, shuffle=True,
seed=7, buffer_size=4, max_buffer_input_splits=1)
batches = list(dataset)
ids = torch.cat([batch['id'] for batch in batches]).tolist()
vectors = torch.cat([batch['vector'] for batch in batches])
self.assertEqual(vectors.dtype, torch.float32)
self.assertEqual(vectors.tolist(), [[i, i + 1] for i in ids])
self.assertEqual(sorted(ids), expected)
def test_torch_streaming_shuffle_single_worker(self):
table = self._create_shuffle_append_table('default.test_torch_shuffle_single')
read_builder = table.new_read_builder().with_projection(['user_id'])
table_read = read_builder.new_read()
splits = read_builder.new_scan().plan().splits()
expected = list(range(80))
for max_buffer_input_splits in [1, 3]:
with self.subTest(max_buffer_input_splits=max_buffer_input_splits):
dataset = table_read.to_torch(
splits,
streaming=True,
shuffle=True,
seed=17,
buffer_size=7,
max_buffer_input_splits=max_buffer_input_splits,
)
ids = self._collect_torch_user_ids(dataset, num_workers=0)
self.assertEqual(sorted(ids), expected)
self.assertNotEqual(ids, expected)
def test_torch_streaming_shuffle_seed_and_epoch(self):
table = self._create_shuffle_append_table('default.test_torch_shuffle_epoch')
read_builder = table.new_read_builder().with_projection(['user_id'])
table_read = read_builder.new_read()
splits = read_builder.new_scan().plan().splits()
dataset = table_read.to_torch(
splits,
streaming=True,
shuffle=True,
seed=23,
buffer_size=11,
max_buffer_input_splits=4,
)
epoch0 = self._collect_torch_user_ids(dataset, num_workers=0)
epoch0_again = self._collect_torch_user_ids(dataset, num_workers=0)
self.assertEqual(epoch0, epoch0_again)
dataset.set_epoch(1)
epoch1 = self._collect_torch_user_ids(dataset, num_workers=0)
self.assertEqual(sorted(epoch1), list(range(80)))
self.assertNotEqual(epoch0, epoch1)
dataset.set_epoch(0)
self.assertEqual(epoch0, self._collect_torch_user_ids(dataset, num_workers=0))
other_seed_dataset = table_read.to_torch(
splits,
streaming=True,
shuffle=True,
seed=24,
buffer_size=11,
max_buffer_input_splits=4,
)
self.assertNotEqual(
epoch0,
self._collect_torch_user_ids(other_seed_dataset, num_workers=0),
)
def test_torch_streaming_shuffle_epoch_with_persistent_workers(self):
table = self._create_shuffle_append_table('default.test_torch_shuffle_persistent_epoch')
read_builder = table.new_read_builder().with_projection(['user_id'])
table_read = read_builder.new_read()
splits = read_builder.new_scan().plan().splits()
dataset = table_read.to_torch(
splits,
streaming=True,
shuffle=True,
seed=23,
buffer_size=11,
max_buffer_input_splits=4,
)
dataloader = DataLoader(
dataset,
batch_size=8,
num_workers=2,
persistent_workers=True,
shuffle=False,
)
epoch0 = self._collect_torch_user_ids_from_dataloader(dataloader)
self.assertEqual(epoch0, self._collect_torch_user_ids_from_dataloader(dataloader))
dataset.set_epoch(1)
epoch1 = self._collect_torch_user_ids_from_dataloader(dataloader)
self.assertEqual(sorted(epoch1), list(range(80)))
self.assertNotEqual(epoch0, epoch1)
def test_torch_streaming_shuffle_multi_worker(self):
table = self._create_shuffle_append_table('default.test_torch_shuffle_multi')
read_builder = table.new_read_builder().with_projection(['user_id'])
table_read = read_builder.new_read()
splits = read_builder.new_scan() \
.with_chunk_shuffle(seed=31, chunk_size=5) \
.plan() \
.splits()
dataset = table_read.to_torch(
splits,
streaming=True,
shuffle=True,
seed=31,
buffer_size=13,
max_buffer_input_splits=4,
)
ids = self._collect_torch_user_ids(dataset, num_workers=2)
expected = list(range(80))
self.assertEqual(len(ids), len(expected))
self.assertEqual(sorted(ids), expected)
def test_torch_streaming_shuffle_rejects_non_streaming(self):
table = self._create_shuffle_append_table('default.test_torch_shuffle_non_streaming')
read_builder = table.new_read_builder()
table_read = read_builder.new_read()
splits = read_builder.new_scan().plan().splits()
with self.assertRaisesRegex(ValueError, "streaming=True"):
table_read.to_torch(splits, streaming=False, shuffle=True)
def test_torch_streaming_shuffle_accepts_pk_table_splits(self):
pa_schema = pa.schema([
pa.field('user_id', pa.int32(), nullable=False),
('item_id', pa.int64()),
('behavior', pa.string()),
('dt', pa.string())
])
schema = Schema.from_pyarrow_schema(
pa_schema,
primary_keys=['user_id'],
options={'bucket': '1'},
)
self.catalog.create_table('default.test_torch_shuffle_pk', schema, False)
table = self.catalog.get_table('default.test_torch_shuffle_pk')
self._write_test_table(table)
read_builder = table.new_read_builder().with_projection(['user_id'])
splits = read_builder.new_scan().plan().splits()
dataset = read_builder.new_read().to_torch(
splits,
streaming=True,
shuffle=True,
seed=7,
buffer_size=3,
)
ids = self._collect_torch_user_ids(dataset, num_workers=0)
self.assertEqual(sorted(ids), [1, 2, 3, 4, 5, 6, 7, 8])
def test_torch_streaming_shuffle_rejects_invalid_dataset_options(self):
table = self._create_shuffle_append_table('default.test_torch_shuffle_invalid_options')
read_builder = table.new_read_builder().with_projection(['user_id'])
table_read = read_builder.new_read()
splits = read_builder.new_scan().plan().splits()
with self.assertRaisesRegex(ValueError, "prefetch_concurrency"):
table_read.to_torch(
splits,
streaming=True,
shuffle=True,
prefetch_concurrency=2,
)
with self.assertRaisesRegex(ValueError, "buffer_size"):
table_read.to_torch(
splits,
streaming=True,
shuffle=True,
buffer_size=0,
)
with self.assertRaisesRegex(ValueError, "max_buffer_input_splits"):
table_read.to_torch(
splits,
streaming=True,
shuffle=True,
max_buffer_input_splits=0,
)
def _create_shuffle_append_table(
self,
identifier,
total_rows=80,
rows_per_commit=10,
partition_keys=None,
):
schema = Schema.from_pyarrow_schema(
self.pa_schema,
partition_keys=partition_keys or [],
)
self.catalog.create_table(identifier, schema, False)
table = self.catalog.get_table(identifier)
write_builder = table.new_batch_write_builder()
for start in range(0, total_rows, rows_per_commit):
end = min(start + rows_per_commit, total_rows)
table_write = write_builder.new_write()
table_commit = write_builder.new_commit()
pa_table = pa.Table.from_pydict({
'user_id': list(range(start, end)),
'item_id': [1000 + i for i in range(start, end)],
'behavior': [chr(ord('a') + (i % 26)) for i in range(start, end)],
'dt': [f'p{i % 4}' for i in range(start, end)],
}, schema=self.pa_schema)
table_write.write_arrow(pa_table)
table_commit.commit(table_write.prepare_commit())
table_write.close()
table_commit.close()
return table
@staticmethod
def _collect_torch_user_ids(dataset, num_workers=0):
dataloader = DataLoader(
dataset,
batch_size=8,
num_workers=num_workers,
shuffle=False,
)
all_user_ids = []
for batch_data in dataloader:
all_user_ids.extend(batch_data['user_id'].tolist())
return all_user_ids
@staticmethod
def _collect_torch_user_ids_from_dataloader(dataloader):
all_user_ids = []
for batch_data in dataloader:
all_user_ids.extend(batch_data['user_id'].tolist())
return all_user_ids
def _write_test_table(self, table):
write_builder = table.new_batch_write_builder()
table_pa_schema = self.pk_pa_schema if table.primary_keys else self.pa_schema
# first write
table_write = write_builder.new_write()
table_commit = write_builder.new_commit()
data1 = {
'user_id': [1, 2, 3, 4],
'item_id': [1001, 1002, 1003, 1004],
'behavior': ['a', 'b', 'c', 'd'],
'dt': ['p1', 'p1', 'p2', 'p1'],
}
pa_table = pa.Table.from_pydict(data1, schema=table_pa_schema)
table_write.write_arrow(pa_table)
table_commit.commit(table_write.prepare_commit())
table_write.close()
table_commit.close()
# second write
table_write = write_builder.new_write()
table_commit = write_builder.new_commit()
data2 = {
'user_id': [5, 6, 7, 8],
'item_id': [1005, 1006, 1007, 1008],
'behavior': ['e', 'f', 'g', 'h'],
'dt': ['p2', 'p1', 'p2', 'p1'],
}
pa_table = pa.Table.from_pydict(data2, schema=table_pa_schema)
table_write.write_arrow(pa_table)
table_commit.commit(table_write.prepare_commit())
table_write.close()
table_commit.close()
def _read_test_table(self, read_builder):
table_read = read_builder.new_read()
splits = read_builder.new_scan().plan().splits()
return table_read.to_arrow(splits)