chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,96 @@
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
from ray.data._internal.arrow_block import ArrowBlockAccessor
|
||||
from ray.data._internal.block_batching.block_batching import batch_blocks
|
||||
|
||||
|
||||
def block_generator(num_rows: int, num_blocks: int):
|
||||
for i in range(num_blocks):
|
||||
yield pa.table({"foo": list(range(i * num_rows, (i + 1) * num_rows))})
|
||||
|
||||
|
||||
class TestBatchBlocks:
|
||||
"""Tests for batch_blocks function."""
|
||||
|
||||
@pytest.mark.parametrize("batch_format", ["pandas", "numpy", "pyarrow"])
|
||||
def test_basic(self, batch_format):
|
||||
"""Test that batch_blocks yields all data in the requested format."""
|
||||
blocks = block_generator(num_rows=3, num_blocks=2)
|
||||
batches = list(batch_blocks(blocks, batch_format=batch_format))
|
||||
|
||||
assert len(batches) == 2
|
||||
|
||||
if batch_format == "pandas":
|
||||
assert isinstance(batches[0], pd.DataFrame)
|
||||
assert isinstance(batches[1], pd.DataFrame)
|
||||
pd.testing.assert_frame_equal(
|
||||
batches[0],
|
||||
ArrowBlockAccessor(pa.table({"foo": [0, 1, 2]})).to_pandas(),
|
||||
)
|
||||
pd.testing.assert_frame_equal(
|
||||
batches[1],
|
||||
ArrowBlockAccessor(pa.table({"foo": [3, 4, 5]})).to_pandas(),
|
||||
)
|
||||
elif batch_format == "numpy":
|
||||
assert isinstance(batches[0], dict)
|
||||
assert isinstance(batches[1], dict)
|
||||
np.testing.assert_array_equal(batches[0]["foo"], np.array([0, 1, 2]))
|
||||
np.testing.assert_array_equal(batches[1]["foo"], np.array([3, 4, 5]))
|
||||
elif batch_format == "pyarrow":
|
||||
assert batches == [
|
||||
pa.table({"foo": [0, 1, 2]}),
|
||||
pa.table({"foo": [3, 4, 5]}),
|
||||
]
|
||||
else:
|
||||
pytest.fail(f"Unsupported batch format {batch_format}")
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_size,drop_last,expected_values",
|
||||
[
|
||||
# 6 rows, batch_size=2: yields 3 full batches
|
||||
(2, False, [[0, 1], [2, 3], [4, 5]]),
|
||||
# 6 rows, batch_size=4: yields 1 full + 1 partial batch
|
||||
(4, False, [[0, 1, 2, 3], [4, 5]]),
|
||||
# 6 rows, batch_size=4, drop_last: drops partial batch
|
||||
(4, True, [[0, 1, 2, 3]]),
|
||||
# 6 rows, batch_size=10, drop_last: no batches (all dropped)
|
||||
(10, True, []),
|
||||
],
|
||||
)
|
||||
def test_batch_size(self, batch_size, drop_last, expected_values):
|
||||
"""Test batch_size and drop_last parameters."""
|
||||
blocks = block_generator(num_rows=3, num_blocks=2)
|
||||
batches = list(
|
||||
batch_blocks(
|
||||
blocks,
|
||||
batch_size=batch_size,
|
||||
drop_last=drop_last,
|
||||
batch_format="numpy",
|
||||
)
|
||||
)
|
||||
|
||||
assert len(batches) == len(expected_values)
|
||||
for batch, expected in zip(batches, expected_values):
|
||||
np.testing.assert_array_equal(batch["foo"], np.array(expected))
|
||||
|
||||
def test_collate_fn(self):
|
||||
"""Test that collate_fn transforms batches."""
|
||||
|
||||
def double_values(batch):
|
||||
return {"foo": [x * 2 for x in batch["foo"].tolist()]}
|
||||
|
||||
blocks = block_generator(num_rows=3, num_blocks=2)
|
||||
batches = list(batch_blocks(blocks, collate_fn=double_values))
|
||||
|
||||
assert len(batches) == 2
|
||||
assert batches[0]["foo"] == [0, 2, 4]
|
||||
assert batches[1]["foo"] == [6, 8, 10]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,170 @@
|
||||
"""Tests for ray.data._internal.block_batching.interfaces."""
|
||||
|
||||
import pytest
|
||||
|
||||
from ray.data._internal.block_batching.interfaces import (
|
||||
Batch,
|
||||
BatchMetadata,
|
||||
BatchStageTimings,
|
||||
BlockStageTimings,
|
||||
)
|
||||
from ray.data._internal.stats import IterationStage, TimeSpan
|
||||
|
||||
|
||||
class TestAccumulateBlockTimings:
|
||||
"""Tests for BatchStageTimings.accumulate_block_timings().
|
||||
|
||||
accumulate_block_timings appends each block's spans to the batch's lists
|
||||
(no merging) so that overlap attribution can sum non-overlapping spans
|
||||
without double-counting.
|
||||
"""
|
||||
|
||||
def test_single_block(self):
|
||||
"""Accumulating a single block appends its span."""
|
||||
dst = BatchStageTimings()
|
||||
dst.accumulate_block_timings(
|
||||
BlockStageTimings(
|
||||
production_wait=TimeSpan(start_s=1.0, end_s=2.0),
|
||||
data_transfer=TimeSpan(start_s=2.0, end_s=3.0),
|
||||
)
|
||||
)
|
||||
assert len(dst.production_wait) == 1
|
||||
assert dst.production_wait[0].start_s == 1.0
|
||||
assert dst.production_wait[0].end_s == 2.0
|
||||
|
||||
def test_multiple_blocks_kept_separate(self):
|
||||
"""Multiple blocks produce a list of separate spans (no merge)."""
|
||||
dst = BatchStageTimings()
|
||||
|
||||
dst.accumulate_block_timings(
|
||||
BlockStageTimings(
|
||||
production_wait=TimeSpan(start_s=1.0, end_s=2.0),
|
||||
data_transfer=TimeSpan(start_s=2.0, end_s=3.0),
|
||||
)
|
||||
)
|
||||
dst.accumulate_block_timings(
|
||||
BlockStageTimings(
|
||||
production_wait=TimeSpan(start_s=3.0, end_s=4.0),
|
||||
data_transfer=TimeSpan(start_s=4.0, end_s=5.0),
|
||||
)
|
||||
)
|
||||
dst.accumulate_block_timings(
|
||||
BlockStageTimings(
|
||||
production_wait=TimeSpan(start_s=5.0, end_s=6.0),
|
||||
data_transfer=TimeSpan(start_s=6.0, end_s=7.0),
|
||||
)
|
||||
)
|
||||
|
||||
assert len(dst.production_wait) == 3
|
||||
assert [s.start_s for s in dst.production_wait] == [1.0, 3.0, 5.0]
|
||||
assert [s.end_s for s in dst.production_wait] == [2.0, 4.0, 6.0]
|
||||
|
||||
def test_overlapping_blocks_kept_separate(self):
|
||||
"""Overlapping windows are NOT merged — kept as separate spans."""
|
||||
dst = BatchStageTimings()
|
||||
|
||||
dst.accumulate_block_timings(
|
||||
BlockStageTimings(
|
||||
production_wait=TimeSpan(start_s=1.0, end_s=5.0),
|
||||
data_transfer=TimeSpan(start_s=5.0, end_s=6.0),
|
||||
)
|
||||
)
|
||||
dst.accumulate_block_timings(
|
||||
BlockStageTimings(
|
||||
production_wait=TimeSpan(start_s=3.0, end_s=7.0),
|
||||
data_transfer=TimeSpan(start_s=7.0, end_s=8.0),
|
||||
)
|
||||
)
|
||||
|
||||
assert len(dst.production_wait) == 2
|
||||
assert dst.production_wait[0].start_s == 1.0
|
||||
assert dst.production_wait[1].end_s == 7.0
|
||||
|
||||
def test_into_empty_destination(self):
|
||||
"""Accumulating into an empty BatchStageTimings appends the span."""
|
||||
dst = BatchStageTimings()
|
||||
dst.accumulate_block_timings(
|
||||
BlockStageTimings(
|
||||
production_wait=TimeSpan(start_s=10.0, end_s=20.0),
|
||||
data_transfer=TimeSpan(start_s=20.0, end_s=30.0),
|
||||
)
|
||||
)
|
||||
assert len(dst.production_wait) == 1
|
||||
assert dst.production_wait[0].start_s == 10.0
|
||||
|
||||
def test_data_transfer_multiple_blocks(self):
|
||||
"""data_transfer spans are kept separate across multiple blocks."""
|
||||
dst = BatchStageTimings()
|
||||
|
||||
dst.accumulate_block_timings(
|
||||
BlockStageTimings(
|
||||
production_wait=TimeSpan(start_s=0.0, end_s=1.0),
|
||||
data_transfer=TimeSpan(start_s=1.0, end_s=2.0),
|
||||
)
|
||||
)
|
||||
dst.accumulate_block_timings(
|
||||
BlockStageTimings(
|
||||
production_wait=TimeSpan(start_s=2.0, end_s=3.0),
|
||||
data_transfer=TimeSpan(start_s=3.0, end_s=4.0),
|
||||
)
|
||||
)
|
||||
|
||||
assert len(dst.data_transfer) == 2
|
||||
assert [s.start_s for s in dst.data_transfer] == [1.0, 3.0]
|
||||
|
||||
def test_both_stages_independent(self):
|
||||
"""production_wait and data_transfer lists are independent."""
|
||||
dst = BatchStageTimings()
|
||||
|
||||
# Block 1: prod [1,2], xfer [2,3]
|
||||
dst.accumulate_block_timings(
|
||||
BlockStageTimings(
|
||||
production_wait=TimeSpan(start_s=1.0, end_s=2.0),
|
||||
data_transfer=TimeSpan(start_s=2.0, end_s=3.0),
|
||||
)
|
||||
)
|
||||
# Block 2: prod [5,6], xfer [6,7]
|
||||
dst.accumulate_block_timings(
|
||||
BlockStageTimings(
|
||||
production_wait=TimeSpan(start_s=5.0, end_s=6.0),
|
||||
data_transfer=TimeSpan(start_s=6.0, end_s=7.0),
|
||||
)
|
||||
)
|
||||
|
||||
assert len(dst.production_wait) == 2
|
||||
assert len(dst.data_transfer) == 2
|
||||
assert [s.start_s for s in dst.production_wait] == [1.0, 5.0]
|
||||
assert [s.start_s for s in dst.data_transfer] == [2.0, 6.0]
|
||||
|
||||
|
||||
class TestStageTimingsFields:
|
||||
"""Tests that BatchStageTimings fields are accessible via stages()."""
|
||||
|
||||
def test_batch_carries_timings_through_pipeline(self):
|
||||
"""A Batch's metadata.stage_timings carries all stage windows."""
|
||||
timings = BatchStageTimings()
|
||||
timings.production_wait.append(TimeSpan(start_s=1.0, end_s=2.0))
|
||||
timings.batching = TimeSpan(start_s=2.0, end_s=3.0)
|
||||
timings.format = TimeSpan(start_s=3.0, end_s=4.0)
|
||||
timings.collate = TimeSpan(start_s=4.0, end_s=5.0)
|
||||
timings.finalize = TimeSpan(start_s=5.0, end_s=6.0)
|
||||
|
||||
batch = Batch(
|
||||
BatchMetadata(batch_idx=0, num_rows=50, stage_timings=timings), None
|
||||
)
|
||||
|
||||
# Verify all stages are accessible via stages() iterator
|
||||
stage_dict = dict(batch.metadata.stage_timings.stages())
|
||||
assert len(stage_dict) == 6
|
||||
assert stage_dict[IterationStage.PRODUCTION_WAIT][0].start_s == 1.0
|
||||
assert stage_dict[IterationStage.BATCHING][0].end_s == 3.0
|
||||
assert stage_dict[IterationStage.FORMAT][0].start_s == 3.0
|
||||
assert stage_dict[IterationStage.COLLATE][0].end_s == 5.0
|
||||
assert stage_dict[IterationStage.FINALIZE][0].start_s == 5.0
|
||||
assert batch.metadata.num_rows == 50
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,874 @@
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
from typing import Iterator, List, Optional
|
||||
from unittest.mock import patch
|
||||
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data._internal.block_batching.interfaces import (
|
||||
Batch,
|
||||
BatchMetadata,
|
||||
BatchStageTimings,
|
||||
BlockPrefetcher,
|
||||
)
|
||||
from ray.data._internal.block_batching.iter_batches import (
|
||||
BatchIterator,
|
||||
prefetch_batches_locally,
|
||||
restore_original_order,
|
||||
)
|
||||
from ray.data._internal.block_batching.util import (
|
||||
WaitBlockPrefetcher,
|
||||
)
|
||||
from ray.data._internal.execution.interfaces.ref_bundle import BlockEntry, RefBundle
|
||||
from ray.data._internal.stats import DatasetStats, TimeSpan
|
||||
from ray.data.block import Block, BlockAccessor, BlockMetadata
|
||||
from ray.types import ObjectRef
|
||||
|
||||
# Sleep duration injected into each scenario's bottleneck stage. Picked to be
|
||||
# large enough to dominate scheduling/measurement noise but small enough to
|
||||
# keep the test fast (5 batches × 0.3s ≈ 1.5s per scenario).
|
||||
SLEEP_S = 0.3
|
||||
|
||||
|
||||
def ref_bundle_generator(num_rows: int, num_blocks: int) -> Iterator[RefBundle]:
|
||||
for i in range(num_blocks):
|
||||
block = pa.table({"foo": [i] * num_rows})
|
||||
metadata = BlockMetadata(
|
||||
num_rows=num_rows,
|
||||
size_bytes=0,
|
||||
input_files=[],
|
||||
exec_stats=None,
|
||||
)
|
||||
schema = block.schema
|
||||
yield RefBundle(
|
||||
blocks=(BlockEntry(ray.put(block), metadata),),
|
||||
owns_blocks=True,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_batches_to_prefetch", [1, 2])
|
||||
@pytest.mark.parametrize("batch_size", [None, 1, 4])
|
||||
def test_prefetch_batches_locally(
|
||||
ray_start_regular_shared, num_batches_to_prefetch, batch_size
|
||||
):
|
||||
class DummyPrefetcher(BlockPrefetcher):
|
||||
def __init__(self):
|
||||
self.windows = []
|
||||
|
||||
def prefetch_blocks(self, block_refs: List[ObjectRef[Block]]):
|
||||
if batch_size is None:
|
||||
assert len(block_refs) == num_batches_to_prefetch
|
||||
else:
|
||||
assert (
|
||||
sum(len(ray.get(block_ref)) for block_ref in block_refs)
|
||||
>= batch_size * num_batches_to_prefetch
|
||||
)
|
||||
self.windows.append(block_refs)
|
||||
|
||||
num_blocks = 10
|
||||
num_rows = 2
|
||||
prefetcher = DummyPrefetcher()
|
||||
ref_bundles = list(ref_bundle_generator(num_blocks=num_blocks, num_rows=num_rows))
|
||||
prefetch_block_iter = prefetch_batches_locally(
|
||||
iter(ref_bundles),
|
||||
prefetcher=prefetcher,
|
||||
num_batches_to_prefetch=num_batches_to_prefetch,
|
||||
batch_size=batch_size,
|
||||
)
|
||||
|
||||
block_count = 0
|
||||
prefetched_blocks = []
|
||||
previous_num_windows = 1
|
||||
|
||||
for block in prefetch_block_iter:
|
||||
prefetched_blocks.append(block)
|
||||
block_count += 1
|
||||
remaining_rows = (num_blocks - block_count) * num_rows
|
||||
if batch_size is None and block_count < num_blocks - num_batches_to_prefetch:
|
||||
# Test that we are actually prefetching in advance if this is not the last
|
||||
# block.
|
||||
assert len(prefetcher.windows) == previous_num_windows + 1
|
||||
previous_num_windows = len(prefetcher.windows)
|
||||
elif (
|
||||
batch_size is not None
|
||||
and remaining_rows > batch_size * num_batches_to_prefetch
|
||||
):
|
||||
# Test that we are actually prefetching in advance if this is not the last
|
||||
# batch.
|
||||
assert len(prefetcher.windows) == previous_num_windows + 1
|
||||
previous_num_windows = len(prefetcher.windows)
|
||||
|
||||
# Test that original blocks are unchanged.
|
||||
expected_blocks = []
|
||||
for ref_bundle in ref_bundles:
|
||||
expected_blocks.extend(ref_bundle.block_refs)
|
||||
assert prefetched_blocks == expected_blocks
|
||||
|
||||
|
||||
def test_restore_from_original_order():
|
||||
base_iterator = [
|
||||
Batch(BatchMetadata(batch_idx=1), None),
|
||||
Batch(BatchMetadata(batch_idx=0), None),
|
||||
Batch(BatchMetadata(batch_idx=3), None),
|
||||
Batch(BatchMetadata(batch_idx=2), None),
|
||||
]
|
||||
|
||||
ordered = list(restore_original_order(iter(base_iterator)))
|
||||
idx = [batch.metadata.batch_idx for batch in ordered]
|
||||
assert idx == [0, 1, 2, 3]
|
||||
|
||||
|
||||
def test_attribute_blocked_time_overlap_attribution():
|
||||
stats = DatasetStats(metadata={}, parent=None)
|
||||
batch_iterator = BatchIterator(iter([]), stats=stats)
|
||||
timings = BatchStageTimings()
|
||||
timings.production_wait.append(TimeSpan(start_s=10.0, end_s=20.0))
|
||||
timings.batching = TimeSpan(start_s=20.0, end_s=30.0)
|
||||
timings.format = TimeSpan(start_s=30.0, end_s=40.0)
|
||||
timings.finalize = TimeSpan(start_s=50.0, end_s=60.0)
|
||||
batch = Batch(BatchMetadata(batch_idx=0, num_rows=8, stage_timings=timings), None)
|
||||
|
||||
batch_iterator._attribute_blocked_time(
|
||||
batch, blocked_start_s=15.0, blocked_end_s=35.0
|
||||
)
|
||||
|
||||
assert stats.iter_blocked_production_wait_s.get() == pytest.approx(5.0)
|
||||
assert stats.iter_blocked_batching_s.get() == pytest.approx(10.0)
|
||||
assert stats.iter_blocked_format_s.get() == pytest.approx(5.0)
|
||||
assert stats.iter_blocked_collate_s.get() == 0
|
||||
assert stats.iter_blocked_finalize_s.get() == 0
|
||||
assert stats.iter_batches_total == 1
|
||||
assert stats.iter_rows_total == 8
|
||||
|
||||
|
||||
def _make_span(start: Optional[float], end: Optional[float]) -> Optional[TimeSpan]:
|
||||
"""Create a TimeSpan, or None if the stage didn't run."""
|
||||
if start is None or end is None:
|
||||
return None
|
||||
return TimeSpan(start_s=start, end_s=end)
|
||||
|
||||
|
||||
def _make_batch_with_timings(
|
||||
production_wait_start: Optional[float] = None,
|
||||
production_wait_end: Optional[float] = None,
|
||||
data_transfer_start: Optional[float] = None,
|
||||
data_transfer_end: Optional[float] = None,
|
||||
batching_start: Optional[float] = None,
|
||||
batching_end: Optional[float] = None,
|
||||
format_start: Optional[float] = None,
|
||||
format_end: Optional[float] = None,
|
||||
collate_start: Optional[float] = None,
|
||||
collate_end: Optional[float] = None,
|
||||
finalize_start: Optional[float] = None,
|
||||
finalize_end: Optional[float] = None,
|
||||
num_rows: int = 0,
|
||||
):
|
||||
"""Helper to construct a Batch with specific stage timing windows."""
|
||||
timings = BatchStageTimings()
|
||||
pw = _make_span(production_wait_start, production_wait_end)
|
||||
if pw is not None:
|
||||
timings.production_wait.append(pw)
|
||||
dt = _make_span(data_transfer_start, data_transfer_end)
|
||||
if dt is not None:
|
||||
timings.data_transfer.append(dt)
|
||||
timings.batching = _make_span(batching_start, batching_end)
|
||||
timings.format = _make_span(format_start, format_end)
|
||||
timings.collate = _make_span(collate_start, collate_end)
|
||||
timings.finalize = _make_span(finalize_start, finalize_end)
|
||||
return Batch(
|
||||
BatchMetadata(batch_idx=0, num_rows=num_rows, stage_timings=timings), None
|
||||
)
|
||||
|
||||
|
||||
def _make_test_iterator(stats):
|
||||
"""Create a BatchIterator wired to the given stats without a real pipeline."""
|
||||
it = BatchIterator.__new__(BatchIterator)
|
||||
it._stats = stats
|
||||
return it
|
||||
|
||||
|
||||
class TestAttributeBlockedTimeEdgeCases:
|
||||
"""Edge case tests for overlap-based blocked attribution."""
|
||||
|
||||
def test_zero_overlap_stage_finished_before_blocked(self):
|
||||
"""Fetch [0, 1.5] finished before training blocked at t=2 → 0 attribution."""
|
||||
stats = DatasetStats(metadata={}, parent=None)
|
||||
it = _make_test_iterator(stats)
|
||||
batch = _make_batch_with_timings(
|
||||
production_wait_start=0.0, production_wait_end=1.5
|
||||
)
|
||||
it._attribute_blocked_time(batch, blocked_start_s=2.0, blocked_end_s=3.0)
|
||||
assert stats.iter_blocked_production_wait_s.get() == 0.0
|
||||
|
||||
def test_zero_overlap_blocked_before_stage(self):
|
||||
"""Training blocked [0, 1], stage ran [2, 3] → 0 attribution."""
|
||||
stats = DatasetStats(metadata={}, parent=None)
|
||||
it = _make_test_iterator(stats)
|
||||
batch = _make_batch_with_timings(format_start=2.0, format_end=3.0)
|
||||
it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=1.0)
|
||||
assert stats.iter_blocked_format_s.get() == 0.0
|
||||
|
||||
def test_partial_overlap(self):
|
||||
"""Fetch [0, 2], blocked [1, 3] → overlap = min(2,3)-max(0,1) = 1.0."""
|
||||
stats = DatasetStats(metadata={}, parent=None)
|
||||
it = _make_test_iterator(stats)
|
||||
batch = _make_batch_with_timings(
|
||||
production_wait_start=0.0, production_wait_end=2.0
|
||||
)
|
||||
it._attribute_blocked_time(batch, blocked_start_s=1.0, blocked_end_s=3.0)
|
||||
assert stats.iter_blocked_production_wait_s.get() == pytest.approx(1.0)
|
||||
|
||||
def test_full_overlap_stage_inside_blocked(self):
|
||||
"""Stage [1, 2] entirely inside blocked [0, 3] → full 1.0 credit."""
|
||||
stats = DatasetStats(metadata={}, parent=None)
|
||||
it = _make_test_iterator(stats)
|
||||
batch = _make_batch_with_timings(batching_start=1.0, batching_end=2.0)
|
||||
it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=3.0)
|
||||
assert stats.iter_blocked_batching_s.get() == pytest.approx(1.0)
|
||||
|
||||
def test_no_collate_fn_zero_attribution(self):
|
||||
"""collate stage has start_s=0 → skipped, 0 attribution."""
|
||||
stats = DatasetStats(metadata={}, parent=None)
|
||||
it = _make_test_iterator(stats)
|
||||
batch = _make_batch_with_timings(format_start=1.0, format_end=2.0)
|
||||
it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=3.0)
|
||||
assert stats.iter_blocked_format_s.get() == pytest.approx(1.0)
|
||||
assert stats.iter_blocked_collate_s.get() == 0.0
|
||||
|
||||
def test_no_finalize_fn_zero_attribution(self):
|
||||
"""finalize stage has start_s=0 → skipped, 0 attribution."""
|
||||
stats = DatasetStats(metadata={}, parent=None)
|
||||
it = _make_test_iterator(stats)
|
||||
batch = _make_batch_with_timings(collate_start=1.0, collate_end=2.0)
|
||||
it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=3.0)
|
||||
assert stats.iter_blocked_collate_s.get() == pytest.approx(1.0)
|
||||
assert stats.iter_blocked_finalize_s.get() == 0.0
|
||||
|
||||
def test_prefetch_hides_fetch_from_training(self):
|
||||
"""Effective prefetch: fetch done before training blocks → 0 fetch attribution."""
|
||||
stats = DatasetStats(metadata={}, parent=None)
|
||||
it = _make_test_iterator(stats)
|
||||
batch = _make_batch_with_timings(
|
||||
production_wait_start=0.0,
|
||||
production_wait_end=1.5,
|
||||
collate_start=2.3,
|
||||
collate_end=2.6,
|
||||
)
|
||||
# Training only starts blocking at t=2 (prefetch worked)
|
||||
it._attribute_blocked_time(batch, blocked_start_s=2.0, blocked_end_s=2.6)
|
||||
assert stats.iter_blocked_production_wait_s.get() == 0.0
|
||||
assert stats.iter_blocked_collate_s.get() == pytest.approx(0.3)
|
||||
|
||||
def test_accumulation_across_batches(self):
|
||||
"""Two batches each contribute to fetch — values accumulate."""
|
||||
stats = DatasetStats(metadata={}, parent=None)
|
||||
it = _make_test_iterator(stats)
|
||||
# Batch 1: fetch [0,1], blocked [0,2] → overlap 1.0
|
||||
b1 = _make_batch_with_timings(
|
||||
production_wait_start=0.0, production_wait_end=1.0, num_rows=10
|
||||
)
|
||||
it._attribute_blocked_time(b1, blocked_start_s=0.0, blocked_end_s=2.0)
|
||||
# Batch 2: fetch [5,6], blocked [5,7] → overlap 1.0
|
||||
b2 = _make_batch_with_timings(
|
||||
production_wait_start=5.0, production_wait_end=6.0, num_rows=20
|
||||
)
|
||||
it._attribute_blocked_time(b2, blocked_start_s=5.0, blocked_end_s=7.0)
|
||||
|
||||
assert stats.iter_blocked_production_wait_s.get() == pytest.approx(2.0)
|
||||
assert stats.iter_batches_total == 2
|
||||
assert stats.iter_rows_total == 30
|
||||
|
||||
def test_overlap_invariant_sum_leq_total(self):
|
||||
"""sum(iter_blocked_*) <= iter_total_blocked_s holds for non-overlapping stages."""
|
||||
stats = DatasetStats(metadata={}, parent=None)
|
||||
it = _make_test_iterator(stats)
|
||||
stats.iter_total_blocked_s.add(5.0)
|
||||
batch = _make_batch_with_timings(
|
||||
production_wait_start=0.0,
|
||||
production_wait_end=1.0,
|
||||
batching_start=1.0,
|
||||
batching_end=2.0,
|
||||
format_start=2.0,
|
||||
format_end=3.0,
|
||||
num_rows=5,
|
||||
)
|
||||
it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=5.0)
|
||||
|
||||
total = stats.iter_total_blocked_s.get()
|
||||
sum_stages = (
|
||||
stats.iter_blocked_production_wait_s.get()
|
||||
+ stats.iter_blocked_batching_s.get()
|
||||
+ stats.iter_blocked_format_s.get()
|
||||
+ stats.iter_blocked_collate_s.get()
|
||||
+ stats.iter_blocked_finalize_s.get()
|
||||
)
|
||||
assert sum_stages <= total + 1e-9
|
||||
|
||||
def test_blocked_inside_stage(self):
|
||||
"""Stage [0, 10] fully contains blocked [3, 5] → overlap = 2.0."""
|
||||
stats = DatasetStats(metadata={}, parent=None)
|
||||
it = _make_test_iterator(stats)
|
||||
batch = _make_batch_with_timings(
|
||||
production_wait_start=0.0, production_wait_end=10.0
|
||||
)
|
||||
it._attribute_blocked_time(batch, blocked_start_s=3.0, blocked_end_s=5.0)
|
||||
assert stats.iter_blocked_production_wait_s.get() == pytest.approx(2.0)
|
||||
|
||||
def test_all_stages_simultaneous_overlap(self):
|
||||
"""Multiple stages overlap with blocked window simultaneously."""
|
||||
stats = DatasetStats(metadata={}, parent=None)
|
||||
it = _make_test_iterator(stats)
|
||||
batch = _make_batch_with_timings(
|
||||
production_wait_start=0.0,
|
||||
production_wait_end=1.0,
|
||||
batching_start=1.0,
|
||||
batching_end=2.0,
|
||||
format_start=2.0,
|
||||
format_end=3.0,
|
||||
collate_start=3.0,
|
||||
collate_end=4.0,
|
||||
finalize_start=4.0,
|
||||
finalize_end=5.0,
|
||||
num_rows=100,
|
||||
)
|
||||
# Blocked window covers all stages
|
||||
it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=5.0)
|
||||
assert stats.iter_blocked_production_wait_s.get() == pytest.approx(1.0)
|
||||
assert stats.iter_blocked_batching_s.get() == pytest.approx(1.0)
|
||||
assert stats.iter_blocked_format_s.get() == pytest.approx(1.0)
|
||||
assert stats.iter_blocked_collate_s.get() == pytest.approx(1.0)
|
||||
assert stats.iter_blocked_finalize_s.get() == pytest.approx(1.0)
|
||||
assert stats.iter_batches_total == 1
|
||||
assert stats.iter_rows_total == 100
|
||||
|
||||
def test_overlapping_spans_not_double_counted(self):
|
||||
"""Two overlapping production_wait spans: union, not sum."""
|
||||
stats = DatasetStats(metadata={}, parent=None)
|
||||
it = _make_test_iterator(stats)
|
||||
# Block 1: prod [0, 100], Block 2: prod [50, 150] — overlap [50, 100]
|
||||
# Blocked [0, 200] covers both
|
||||
batch = _make_batch_with_timings(
|
||||
production_wait_start=0.0,
|
||||
production_wait_end=100.0,
|
||||
num_rows=10,
|
||||
)
|
||||
# Add a second production_wait span (multi-block batch)
|
||||
batch.metadata.stage_timings.production_wait.append(
|
||||
TimeSpan(start_s=50.0, end_s=150.0)
|
||||
)
|
||||
it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=200.0)
|
||||
# Union of [0,100] and [50,150] = [0,150] = 150, NOT 100+100=200
|
||||
assert stats.iter_blocked_production_wait_s.get() == pytest.approx(150.0)
|
||||
|
||||
|
||||
def test_attribute_blocked_time_all_stages_full_overlap():
|
||||
"""All stages with realistic timing, full overlap with blocked window."""
|
||||
stats = DatasetStats(metadata={}, parent=None)
|
||||
it = _make_test_iterator(stats)
|
||||
stats.iter_total_blocked_s.add(5.0)
|
||||
|
||||
batch = _make_batch_with_timings(
|
||||
production_wait_start=0.0,
|
||||
production_wait_end=0.5,
|
||||
batching_start=0.5,
|
||||
batching_end=1.0,
|
||||
format_start=1.0,
|
||||
format_end=2.0,
|
||||
collate_start=2.0,
|
||||
collate_end=2.5,
|
||||
finalize_start=2.5,
|
||||
finalize_end=3.0,
|
||||
num_rows=256,
|
||||
)
|
||||
|
||||
# Blocked window covers all stages
|
||||
it._attribute_blocked_time(batch, blocked_start_s=0.0, blocked_end_s=5.0)
|
||||
|
||||
# Each stage gets its full duration
|
||||
assert stats.iter_blocked_production_wait_s.get() == pytest.approx(0.5)
|
||||
assert stats.iter_blocked_batching_s.get() == pytest.approx(0.5)
|
||||
assert stats.iter_blocked_format_s.get() == pytest.approx(1.0)
|
||||
assert stats.iter_blocked_collate_s.get() == pytest.approx(0.5)
|
||||
assert stats.iter_blocked_finalize_s.get() == pytest.approx(0.5)
|
||||
assert stats.iter_batches_total == 1
|
||||
assert stats.iter_rows_total == 256
|
||||
|
||||
# Invariant: sum = 3.0 <= total_blocked = 5.0
|
||||
sum_stages = (
|
||||
stats.iter_blocked_production_wait_s.get()
|
||||
+ stats.iter_blocked_batching_s.get()
|
||||
+ stats.iter_blocked_format_s.get()
|
||||
+ stats.iter_blocked_collate_s.get()
|
||||
+ stats.iter_blocked_finalize_s.get()
|
||||
)
|
||||
assert sum_stages == pytest.approx(3.0)
|
||||
assert sum_stages <= stats.iter_total_blocked_s.get() + 1e-9
|
||||
|
||||
|
||||
def test_finalize_fn_uses_single_thread(ray_start_regular_shared):
|
||||
"""Tests that finalize_fn is not run with multiple threads."""
|
||||
ref_bundles_iter = ref_bundle_generator(num_blocks=20, num_rows=2)
|
||||
|
||||
q = queue.Queue()
|
||||
semaphore = threading.Semaphore(value=1)
|
||||
|
||||
def finalize_enforce_single_thread(batch):
|
||||
already_acquired = not semaphore.acquire(blocking=False)
|
||||
if already_acquired:
|
||||
e = AssertionError("finalize_fn is being run concurrently.")
|
||||
q.put(e, block=True)
|
||||
semaphore.release()
|
||||
return batch
|
||||
|
||||
# Test that finalize_fn is called in a single thread,
|
||||
# even if prefetch_batches is set.
|
||||
output_batches = BatchIterator(
|
||||
ref_bundles_iter,
|
||||
collate_fn=lambda batch: batch,
|
||||
finalize_fn=finalize_enforce_single_thread,
|
||||
prefetch_batches=4,
|
||||
)
|
||||
|
||||
# Force execution of the iterator.
|
||||
# This step should not raise an exception.
|
||||
list(output_batches)
|
||||
|
||||
try:
|
||||
e = q.get(block=False, timeout=0.1)
|
||||
raise e
|
||||
except queue.Empty:
|
||||
pass
|
||||
|
||||
|
||||
# Test for 3 cases
|
||||
# 1. Batch size is less than block size
|
||||
# 2. Batch size is more than block size
|
||||
# 3. Block size is not divisble by batch size
|
||||
@pytest.mark.parametrize("batch_size", [1, 4, 3])
|
||||
@pytest.mark.parametrize("drop_last", [True, False])
|
||||
@pytest.mark.parametrize("prefetch_batches", [0, 1])
|
||||
def test_iter_batches_e2e(
|
||||
ray_start_regular_shared, batch_size, drop_last, prefetch_batches
|
||||
):
|
||||
def collate_fn(batch: pd.DataFrame):
|
||||
return batch + 1
|
||||
|
||||
ref_bundles_iter = ref_bundle_generator(num_blocks=4, num_rows=2)
|
||||
|
||||
output_batches = BatchIterator(
|
||||
ref_bundles_iter,
|
||||
batch_size=batch_size,
|
||||
prefetch_batches=prefetch_batches,
|
||||
batch_format="pandas",
|
||||
collate_fn=collate_fn,
|
||||
drop_last=drop_last,
|
||||
preserve_order=True,
|
||||
)
|
||||
|
||||
output_batches = list(output_batches)
|
||||
|
||||
assert len(output_batches) > 0
|
||||
for df in output_batches:
|
||||
# Check batch formatting.
|
||||
assert isinstance(df, pd.DataFrame)
|
||||
# Check batch size.
|
||||
if batch_size == 3 and not drop_last:
|
||||
assert len(df) in {2, 3}
|
||||
else:
|
||||
assert len(df) == batch_size
|
||||
|
||||
concat_df = pd.concat(output_batches)
|
||||
# Test that collate_fn is applied.
|
||||
assert concat_df["foo"].iloc[0] == 1
|
||||
# Make sure order is preserved.
|
||||
for i in range(len(concat_df) - 1):
|
||||
assert concat_df["foo"].iloc[i + 1] >= concat_df["foo"].iloc[i]
|
||||
|
||||
|
||||
def test_iter_batches_counts_rows_at_pipeline_exit(ray_start_regular_shared):
|
||||
stats = DatasetStats(metadata={}, parent=None)
|
||||
ref_bundles_iter = ref_bundle_generator(num_blocks=4, num_rows=2)
|
||||
|
||||
output_batches = list(
|
||||
BatchIterator(
|
||||
ref_bundles_iter,
|
||||
stats=stats,
|
||||
batch_size=3,
|
||||
prefetch_batches=0,
|
||||
batch_format="pandas",
|
||||
drop_last=True,
|
||||
)
|
||||
)
|
||||
|
||||
assert len(output_batches) == 2
|
||||
assert [len(batch) for batch in output_batches] == [3, 3]
|
||||
assert stats.iter_batches_total == 2
|
||||
assert stats.iter_rows_total == 6
|
||||
|
||||
|
||||
def test_iter_batches_e2e_async(ray_start_regular_shared):
|
||||
"""We add time.sleep in 3 places:
|
||||
1. In the base generator to simulate streaming executor blocking on next results.
|
||||
2. In the collate_fn to simulate expensive slicing/formatting/collation
|
||||
3. In the user thread to simulate training.
|
||||
"""
|
||||
|
||||
def collate_fn(batch):
|
||||
time.sleep(2)
|
||||
return batch
|
||||
|
||||
ref_bundles = ref_bundle_generator(num_blocks=20, num_rows=2)
|
||||
start_time = time.time()
|
||||
output_batches = BatchIterator(
|
||||
ref_bundles,
|
||||
batch_size=None,
|
||||
collate_fn=collate_fn,
|
||||
prefetch_batches=4,
|
||||
)
|
||||
batches = []
|
||||
for batch in output_batches:
|
||||
time.sleep(1.5)
|
||||
batches.append(batch)
|
||||
end_time = time.time()
|
||||
|
||||
# 20 batches, 1.5 second sleep. Should be less than 45 seconds, even with some
|
||||
# overhead.
|
||||
# If there was no overlap, then we would expect this to take at least 20*2.5 = 50
|
||||
assert end_time - start_time < 45, end_time - start_time
|
||||
|
||||
assert len(batches) == 20
|
||||
assert all(len(batch) == 2 for batch in batches)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("preserve_order", [True, False])
|
||||
def test_iter_batches_preserve_order_flag(
|
||||
ray_start_regular_shared, preserve_order, restore_data_context
|
||||
):
|
||||
"""When `execution_options.preserve_order` is True, batches must come
|
||||
out in input order even with a multi-worker format threadpool. When
|
||||
False, ordering is not guaranteed (but the full set of batches must
|
||||
still be produced)."""
|
||||
# Variable per-batch collate cost makes worker-completion order
|
||||
# arbitrary so the reorder path actually does work when enabled.
|
||||
def collate_fn(batch):
|
||||
idx = int(batch["foo"][0])
|
||||
time.sleep(0.05 * (idx % 4))
|
||||
return batch
|
||||
|
||||
num_blocks = 16
|
||||
ref_bundles = ref_bundle_generator(num_blocks=num_blocks, num_rows=1)
|
||||
output_batches = list(
|
||||
BatchIterator(
|
||||
ref_bundles,
|
||||
batch_size=1,
|
||||
collate_fn=collate_fn,
|
||||
batch_format="pandas",
|
||||
prefetch_batches=4,
|
||||
preserve_order=preserve_order,
|
||||
)
|
||||
)
|
||||
|
||||
indices = [int(df["foo"].iloc[0]) for df in output_batches]
|
||||
assert sorted(indices) == list(range(num_blocks))
|
||||
if preserve_order:
|
||||
assert indices == list(range(num_blocks)), indices
|
||||
|
||||
|
||||
def test_finalize_fn_runs_after_restore_original_order(ray_start_regular_shared):
|
||||
"""When preserve_order=True, finalize_fn must run after the reorder
|
||||
buffer so that the buffer holds CPU batches rather than finalize_fn
|
||||
outputs (e.g., GPU tensors). Asserts finalize_fn sees batches in
|
||||
monotonically increasing order even when the format/collate threadpool
|
||||
completes them out of order."""
|
||||
|
||||
def collate_fn(batch):
|
||||
# Variable per-batch cost so worker-completion order is arbitrary.
|
||||
idx = int(batch["foo"].iloc[0])
|
||||
time.sleep(0.05 * (idx % 4))
|
||||
return batch
|
||||
|
||||
seen_by_finalize = []
|
||||
seen_lock = threading.Lock()
|
||||
|
||||
def finalize_fn(batch):
|
||||
idx = int(batch["foo"].iloc[0])
|
||||
with seen_lock:
|
||||
seen_by_finalize.append(idx)
|
||||
return batch
|
||||
|
||||
num_blocks = 16
|
||||
ref_bundles = ref_bundle_generator(num_blocks=num_blocks, num_rows=1)
|
||||
list(
|
||||
BatchIterator(
|
||||
ref_bundles,
|
||||
batch_size=1,
|
||||
collate_fn=collate_fn,
|
||||
finalize_fn=finalize_fn,
|
||||
batch_format="pandas",
|
||||
prefetch_batches=4,
|
||||
preserve_order=True,
|
||||
)
|
||||
)
|
||||
|
||||
assert seen_by_finalize == list(range(num_blocks)), seen_by_finalize
|
||||
|
||||
|
||||
def _ref_bundles_with_size(
|
||||
num_blocks: int, num_rows: int, size_bytes_per_block: int
|
||||
) -> Iterator[RefBundle]:
|
||||
"""Create ref bundles with explicit size_bytes for testing."""
|
||||
for i in range(num_blocks):
|
||||
block = pa.table({"foo": [i] * num_rows})
|
||||
metadata = BlockMetadata(
|
||||
num_rows=num_rows,
|
||||
size_bytes=size_bytes_per_block,
|
||||
input_files=[],
|
||||
exec_stats=None,
|
||||
)
|
||||
schema = block.schema
|
||||
yield RefBundle(
|
||||
blocks=(BlockEntry(ray.put(block), metadata),),
|
||||
owns_blocks=True,
|
||||
schema=schema,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"num_batches_to_prefetch,expected_bytes_sequence",
|
||||
[
|
||||
# No prefetching: all 5 blocks report 0 prefetched bytes
|
||||
(0, [0, 0, 0, 0, 0]),
|
||||
# prefetch 2 blocks: with 5 blocks of 100 bytes each
|
||||
# After yield block 0: window has 1,2 -> 200 (added block 2)
|
||||
# After yield block 1: window has 2,3 -> 200 (added block 3)
|
||||
# After yield block 2: window has 3,4 -> 200 (added block 4)
|
||||
# After yield block 3: window has 4 -> 100 (no more to add)
|
||||
# After yield block 4: window empty -> 0
|
||||
(2, [200, 200, 200, 100, 0]),
|
||||
],
|
||||
)
|
||||
def test_prefetch_bytes_tracking(
|
||||
ray_start_regular_shared, num_batches_to_prefetch, expected_bytes_sequence
|
||||
):
|
||||
"""Test iter_prefetched_bytes is set correctly during prefetching.
|
||||
|
||||
Tests prefetch_batches_locally directly to verify exact values,
|
||||
bypassing async BatchIterator which has non-deterministic timing.
|
||||
"""
|
||||
stats = DatasetStats(metadata={}, parent=None)
|
||||
|
||||
# Create 5 ref bundles, each with size_bytes=100
|
||||
num_blocks = 5
|
||||
ref_bundles = list(
|
||||
_ref_bundles_with_size(num_blocks, num_rows=2, size_bytes_per_block=100)
|
||||
)
|
||||
|
||||
prefetcher = WaitBlockPrefetcher()
|
||||
block_iter = prefetch_batches_locally(
|
||||
iter(ref_bundles),
|
||||
prefetcher=prefetcher,
|
||||
num_batches_to_prefetch=num_batches_to_prefetch,
|
||||
batch_size=None,
|
||||
stats=stats,
|
||||
)
|
||||
|
||||
# Track iter_prefetched_bytes after each block is yielded
|
||||
recorded_bytes = []
|
||||
for _ in block_iter:
|
||||
recorded_bytes.append(stats.iter_prefetched_bytes)
|
||||
|
||||
assert recorded_bytes == expected_bytes_sequence, f"Got {recorded_bytes}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefetch_batches", [0, 2])
|
||||
def test_prefetch_bytes_callback(ray_start_regular_shared, prefetch_batches):
|
||||
"""Test prefetch_bytes_callback is invoked correctly by BatchIterator."""
|
||||
reported_bytes = []
|
||||
|
||||
def prefetch_callback(num_bytes: int):
|
||||
reported_bytes.append(num_bytes)
|
||||
|
||||
stats = DatasetStats(metadata={}, parent=None)
|
||||
|
||||
# Create 5 ref bundles
|
||||
num_blocks = 5
|
||||
ref_bundles = list(
|
||||
_ref_bundles_with_size(num_blocks, num_rows=2, size_bytes_per_block=100)
|
||||
)
|
||||
|
||||
output_batches = BatchIterator(
|
||||
iter(ref_bundles),
|
||||
stats=stats,
|
||||
batch_size=None,
|
||||
prefetch_batches=prefetch_batches,
|
||||
prefetch_bytes_callback=prefetch_callback,
|
||||
)
|
||||
|
||||
# Consume all batches
|
||||
batches = list(output_batches)
|
||||
|
||||
assert len(batches) == 5
|
||||
|
||||
# Callback is called 5 times (per batch) + 1 time at epoch end
|
||||
assert len(reported_bytes) == 6, f"Expected 6, got {len(reported_bytes)}"
|
||||
|
||||
# All values should be non-negative
|
||||
assert all(b >= 0 for b in reported_bytes), f"Negative: {reported_bytes}"
|
||||
|
||||
# Last value should be 0 (after_epoch_end)
|
||||
assert reported_bytes[-1] == 0, f"Last should be 0: {reported_bytes}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"scenario,bound_stage",
|
||||
[
|
||||
("production", "iter_blocked_production_wait_s"),
|
||||
("data_transfer", "iter_blocked_data_transfer_s"),
|
||||
("batching", "iter_blocked_batching_s"),
|
||||
("collate", "iter_blocked_collate_s"),
|
||||
("format", "iter_blocked_format_s"),
|
||||
("finalize", "iter_blocked_finalize_s"),
|
||||
],
|
||||
)
|
||||
def test_e2e_blocked_attribution_by_scenario(
|
||||
ray_start_regular_shared, scenario, bound_stage
|
||||
):
|
||||
"""E2e: when a specific stage is the bottleneck, its blocked metric
|
||||
should be the largest among all stages, and at least SLEEP_S."""
|
||||
from ray.data._internal.stats import _StatsManager
|
||||
|
||||
iter_kwargs = {"batch_size": 10, "prefetch_batches": 0}
|
||||
patches = []
|
||||
|
||||
if scenario == "production":
|
||||
# Slow upstream map → production_wait dominates.
|
||||
def slow_map(batch):
|
||||
time.sleep(SLEEP_S)
|
||||
return batch
|
||||
|
||||
ds = ray.data.range(50, override_num_blocks=5).map(slow_map)
|
||||
|
||||
elif scenario == "data_transfer":
|
||||
# Patch ray.get ONLY in util.resolve_block_refs (not globally) so the
|
||||
# streaming executor's own ray.get calls aren't slowed (which would
|
||||
# inflate production_wait). We replace util_mod.ray with a proxy that
|
||||
# has a slow `get` but delegates everything else to the real ray.
|
||||
from ray.data._internal.block_batching import util as util_mod
|
||||
|
||||
orig_get = ray.get
|
||||
|
||||
class _SlowGetRayProxy:
|
||||
"""Proxy that sleeps on `.get` but delegates everything else."""
|
||||
|
||||
def __getattr__(self, name):
|
||||
return getattr(ray, name)
|
||||
|
||||
@staticmethod
|
||||
def get(ref):
|
||||
time.sleep(SLEEP_S)
|
||||
return orig_get(ref)
|
||||
|
||||
patches.append(patch.object(util_mod, "ray", _SlowGetRayProxy()))
|
||||
ds = ray.data.range(50, override_num_blocks=5)
|
||||
|
||||
elif scenario == "batching":
|
||||
# Patch Batcher.next_batch to inject slow batching.
|
||||
from ray.data._internal.batcher import Batcher
|
||||
|
||||
orig_next_batch = Batcher.next_batch
|
||||
|
||||
def slow_next_batch(self):
|
||||
time.sleep(SLEEP_S)
|
||||
return orig_next_batch(self)
|
||||
|
||||
patches.append(patch.object(Batcher, "next_batch", slow_next_batch))
|
||||
ds = ray.data.range(50, override_num_blocks=5)
|
||||
|
||||
elif scenario == "collate":
|
||||
# Pass _collate_fn via _iter_batches (private signature accepts it;
|
||||
# public iter_batches does not).
|
||||
def slow_collate(batch):
|
||||
time.sleep(SLEEP_S)
|
||||
return batch
|
||||
|
||||
iter_kwargs["_collate_fn"] = slow_collate
|
||||
ds = ray.data.range(50, override_num_blocks=5)
|
||||
|
||||
elif scenario == "format":
|
||||
# Patch BlockAccessor.to_batch_format — it's called INSIDE
|
||||
# _format_batch's _maybe_time context, so the sleep is captured by
|
||||
# the format timing span.
|
||||
orig_to_batch_format = BlockAccessor.to_batch_format
|
||||
|
||||
def slow_to_batch_format(self, batch_format):
|
||||
time.sleep(SLEEP_S)
|
||||
return orig_to_batch_format(self, batch_format)
|
||||
|
||||
patches.append(
|
||||
patch.object(BlockAccessor, "to_batch_format", slow_to_batch_format)
|
||||
)
|
||||
ds = ray.data.range(50, override_num_blocks=5)
|
||||
|
||||
elif scenario == "finalize":
|
||||
# Pass _finalize_fn via _iter_batches (private signature accepts it).
|
||||
def slow_finalize(data):
|
||||
time.sleep(SLEEP_S)
|
||||
return data
|
||||
|
||||
iter_kwargs["_finalize_fn"] = slow_finalize
|
||||
ds = ray.data.range(50, override_num_blocks=5)
|
||||
|
||||
it = ds.iterator()
|
||||
captured = []
|
||||
orig = _StatsManager.update_iteration_metrics
|
||||
|
||||
def spy(stats, dataset_tag):
|
||||
captured.append(stats)
|
||||
return orig(stats, dataset_tag)
|
||||
|
||||
patches.append(patch.object(_StatsManager, "update_iteration_metrics", spy))
|
||||
|
||||
import contextlib
|
||||
|
||||
with contextlib.ExitStack() as stack:
|
||||
for p in patches:
|
||||
stack.enter_context(p)
|
||||
# Use _iter_batches (private) so we can pass _collate_fn / _finalize_fn
|
||||
# which the public iter_batches signature does not expose.
|
||||
for _ in it._iter_batches(**iter_kwargs):
|
||||
pass
|
||||
|
||||
stats = captured[-1]
|
||||
all_stages = [
|
||||
stats.iter_blocked_production_wait_s.get(),
|
||||
stats.iter_blocked_data_transfer_s.get(),
|
||||
stats.iter_blocked_batching_s.get(),
|
||||
stats.iter_blocked_format_s.get(),
|
||||
stats.iter_blocked_collate_s.get(),
|
||||
stats.iter_blocked_finalize_s.get(),
|
||||
]
|
||||
bound_value = getattr(stats, bound_stage).get()
|
||||
# The bottleneck stage should be at least the sleep time we injected,
|
||||
# proving the timing capture is actually recording the stall.
|
||||
assert bound_value >= SLEEP_S, (
|
||||
f"{scenario}-bound: {bound_stage}={bound_value} < SLEEP_S={SLEEP_S}; "
|
||||
"timing capture missed the injected stall"
|
||||
)
|
||||
# The bottleneck stage should be strictly greater than all others.
|
||||
for v in all_stages:
|
||||
if v == bound_value:
|
||||
continue
|
||||
assert (
|
||||
bound_value > v
|
||||
), f"{scenario}-bound: {bound_stage}={bound_value} not > {v}"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,660 @@
|
||||
import logging
|
||||
import random
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from collections import Counter
|
||||
from os import urandom
|
||||
from typing import Callable, Iterator
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data._internal.block_batching.interfaces import (
|
||||
Batch,
|
||||
BatchMetadata,
|
||||
ResolvedBlock,
|
||||
)
|
||||
from ray.data._internal.block_batching.util import (
|
||||
_calculate_ref_hits,
|
||||
blocks_to_batches,
|
||||
collate,
|
||||
finalize_batches,
|
||||
format_batches,
|
||||
iter_threaded,
|
||||
resolve_block_refs,
|
||||
)
|
||||
from ray.data._internal.stats import DatasetStats
|
||||
from ray.data._internal.util import make_async_gen
|
||||
|
||||
logger = logging.getLogger(__file__)
|
||||
|
||||
|
||||
def block_generator(num_rows: int, num_blocks: int):
|
||||
for _ in range(num_blocks):
|
||||
yield pa.table({"foo": [1] * num_rows})
|
||||
|
||||
|
||||
def test_resolve_block_refs(ray_start_regular_shared):
|
||||
block_refs = [ray.put(0), ray.put(1), ray.put(2)]
|
||||
|
||||
resolved_iter = resolve_block_refs(iter(block_refs))
|
||||
resolved = list(resolved_iter)
|
||||
assert all(isinstance(b, ResolvedBlock) for b in resolved)
|
||||
assert [b.block for b in resolved] == [0, 1, 2]
|
||||
|
||||
|
||||
def test_resolve_block_refs_accumulates_data_transfer_timer(
|
||||
ray_start_regular_shared,
|
||||
):
|
||||
"""resolve_block_refs accumulates ray.get() time into iter_get_s and
|
||||
captures a per-block data_transfer TimeSpan."""
|
||||
block_refs = [ray.put(i) for i in range(3)]
|
||||
|
||||
stats = DatasetStats(metadata={}, parent=None)
|
||||
resolved = list(resolve_block_refs(iter(block_refs), stats=stats))
|
||||
|
||||
assert len(resolved) == 3
|
||||
|
||||
# data_transfer TimeSpan captured per block.
|
||||
for r in resolved:
|
||||
assert r.stage_timings is not None
|
||||
assert r.stage_timings.data_transfer is not None
|
||||
assert r.stage_timings.data_transfer.duration >= 0.0
|
||||
|
||||
|
||||
def test_resolve_block_refs_captures_production_wait_span(
|
||||
ray_start_regular_shared,
|
||||
):
|
||||
"""resolve_block_refs captures a per-block production_wait TimeSpan
|
||||
around ``next(block_ref_iter)`` (manual capture, no Timer accumulation)."""
|
||||
block_refs = [ray.put(i) for i in range(3)]
|
||||
|
||||
stats = DatasetStats(metadata={}, parent=None)
|
||||
resolved = list(resolve_block_refs(iter(block_refs), stats=stats))
|
||||
|
||||
assert len(resolved) == 3
|
||||
for r in resolved:
|
||||
assert r.stage_timings is not None
|
||||
assert r.stage_timings.production_wait is not None
|
||||
assert r.stage_timings.production_wait.duration >= 0.0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("block_size", [1, 10])
|
||||
@pytest.mark.parametrize("drop_last", [True, False])
|
||||
def test_blocks_to_batches(block_size, drop_last):
|
||||
num_blocks = 5
|
||||
block_iter = block_generator(num_rows=block_size, num_blocks=num_blocks)
|
||||
# Wrap raw blocks in ResolvedBlock (stage_timings=None) as blocks_to_batches now expects
|
||||
wrapped_blocks = (ResolvedBlock(block=b) for b in block_iter)
|
||||
|
||||
batch_size = 3
|
||||
batch_iter = list(
|
||||
blocks_to_batches(wrapped_blocks, batch_size=batch_size, drop_last=drop_last)
|
||||
)
|
||||
|
||||
if drop_last:
|
||||
for batch in batch_iter:
|
||||
assert len(batch.data) == batch_size
|
||||
else:
|
||||
full_batches = 0
|
||||
leftover_batches = 0
|
||||
|
||||
dataset_size = block_size * num_blocks
|
||||
for batch in batch_iter:
|
||||
if len(batch.data) == batch_size:
|
||||
full_batches += 1
|
||||
if len(batch.data) == (dataset_size % batch_size):
|
||||
leftover_batches += 1
|
||||
|
||||
assert leftover_batches == 1
|
||||
assert full_batches == (dataset_size // batch_size)
|
||||
|
||||
assert [batch.metadata.batch_idx for batch in batch_iter] == list(
|
||||
range(len(batch_iter))
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_format", ["pandas", "numpy", "pyarrow"])
|
||||
def test_format_batches(batch_format):
|
||||
block_iter = block_generator(num_rows=2, num_blocks=2)
|
||||
batch_iter = (
|
||||
Batch(BatchMetadata(batch_idx=i), block) for i, block in enumerate(block_iter)
|
||||
)
|
||||
batch_iter = list(format_batches(batch_iter, batch_format=batch_format))
|
||||
|
||||
for batch in batch_iter:
|
||||
if batch_format == "pandas":
|
||||
assert isinstance(batch.data, pd.DataFrame)
|
||||
elif batch_format == "arrow":
|
||||
assert isinstance(batch.data, pa.Table)
|
||||
elif batch_format == "numpy":
|
||||
assert isinstance(batch.data, dict)
|
||||
assert isinstance(batch.data["foo"], np.ndarray)
|
||||
|
||||
assert [batch.metadata.batch_idx for batch in batch_iter] == list(
|
||||
range(len(batch_iter))
|
||||
)
|
||||
|
||||
|
||||
def test_collate():
|
||||
def collate_fn(batch):
|
||||
return pa.table({"bar": [1] * 2})
|
||||
|
||||
batches = [
|
||||
Batch(BatchMetadata(batch_idx=i), data)
|
||||
for i, data in enumerate(block_generator(num_rows=2, num_blocks=2))
|
||||
]
|
||||
batch_iter = collate(batches, collate_fn=collate_fn)
|
||||
|
||||
for i, batch in enumerate(batch_iter):
|
||||
assert batch.metadata.batch_idx == i
|
||||
assert batch.data == pa.table({"bar": [1] * 2})
|
||||
|
||||
|
||||
def test_finalize():
|
||||
def finalize_fn(batch):
|
||||
return pa.table({"bar": [1] * 2})
|
||||
|
||||
batches = [
|
||||
Batch(BatchMetadata(batch_idx=i), data)
|
||||
for i, data in enumerate(block_generator(num_rows=2, num_blocks=2))
|
||||
]
|
||||
batch_iter = finalize_batches(batches, finalize_fn=finalize_fn)
|
||||
|
||||
for i, batch in enumerate(batch_iter):
|
||||
assert batch.metadata.batch_idx == i
|
||||
assert batch.data == pa.table({"bar": [1] * 2})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("preserve_ordering", [True, False])
|
||||
@pytest.mark.parametrize("buffer_size", [0, 1, 2])
|
||||
def test_make_async_gen_fail(buffer_size: int, preserve_ordering):
|
||||
"""Tests that any errors raised in async threads are propagated to the main
|
||||
thread."""
|
||||
|
||||
def gen(base_iterator):
|
||||
raise ValueError("Fail")
|
||||
|
||||
iterator = make_async_gen(
|
||||
base_iterator=iter([1]),
|
||||
fn=gen,
|
||||
preserve_ordering=preserve_ordering,
|
||||
buffer_size=buffer_size,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError) as e:
|
||||
for _ in iterator:
|
||||
pass
|
||||
|
||||
assert e.match("Fail")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("preserve_ordering", [True, False])
|
||||
def test_make_async_gen_varying_seq_length_stress_test(preserve_ordering):
|
||||
"""This test executes make_async_gen against a function generating variable
|
||||
length sequences to stress test its concurrency control.
|
||||
"""
|
||||
|
||||
num_workers = 4
|
||||
|
||||
c = 0
|
||||
|
||||
# Roll the dice 100 times
|
||||
for i in range(100):
|
||||
# Fetch 8b seed from urandom
|
||||
seed = int.from_bytes(urandom(8), byteorder=sys.byteorder)
|
||||
r = random.Random(seed)
|
||||
|
||||
print(f">>> Seed: {seed}")
|
||||
|
||||
# NOTE: Number of seqs >> number of workers
|
||||
# to saturate the input queue
|
||||
num_seqs = num_workers * 10
|
||||
|
||||
lens = list(range(num_seqs))
|
||||
|
||||
r.shuffle(lens)
|
||||
|
||||
source = [range(len_) for len_ in lens]
|
||||
|
||||
print("===" * 8)
|
||||
print(source)
|
||||
print("===" * 8)
|
||||
|
||||
def flatten(list_iter):
|
||||
for l in list_iter:
|
||||
print(f">>> Flattening: {l}")
|
||||
yield from l
|
||||
|
||||
it = make_async_gen(
|
||||
iter(source),
|
||||
flatten,
|
||||
preserve_ordering=preserve_ordering,
|
||||
num_workers=4,
|
||||
buffer_size=1,
|
||||
)
|
||||
|
||||
total = 0
|
||||
|
||||
for i in it:
|
||||
total += i
|
||||
|
||||
assert total == 9880
|
||||
c += 1
|
||||
|
||||
assert c == 100
|
||||
|
||||
|
||||
@pytest.mark.parametrize("preserve_ordering", [True, False])
|
||||
def test_make_async_gen_non_reentrant(preserve_ordering):
|
||||
"""This test is asserting that make_async_gen iterating over the
|
||||
sequence as a whole and not re-entering provided transformation,
|
||||
as this might have substantial performance impact in extreme case
|
||||
of re-entering for every element of the sequence
|
||||
"""
|
||||
|
||||
logs = []
|
||||
finished = False
|
||||
|
||||
def _transform_inner(it):
|
||||
nonlocal finished
|
||||
|
||||
assert not finished
|
||||
|
||||
logs.append(">>> Entering Inner")
|
||||
|
||||
for i in it:
|
||||
logs.append(f">>> Inner: {i}")
|
||||
yield i
|
||||
|
||||
logs.append(">>> Leaving Inner")
|
||||
|
||||
# Once this transform finishes
|
||||
finished = True
|
||||
|
||||
def _transform_b(it):
|
||||
logs.append(">>> Entering Outer")
|
||||
|
||||
for i in _transform_inner(it):
|
||||
logs.append(f">>> Outer: {i}")
|
||||
yield i
|
||||
|
||||
logs.append(">>> Leaving Outer")
|
||||
|
||||
for _ in make_async_gen(
|
||||
iter(range(3)),
|
||||
_transform_b,
|
||||
preserve_ordering=preserve_ordering,
|
||||
):
|
||||
pass
|
||||
|
||||
assert [
|
||||
">>> Entering Outer",
|
||||
">>> Entering Inner",
|
||||
">>> Inner: 0",
|
||||
">>> Outer: 0",
|
||||
">>> Inner: 1",
|
||||
">>> Outer: 1",
|
||||
">>> Inner: 2",
|
||||
">>> Outer: 2",
|
||||
">>> Leaving Inner",
|
||||
">>> Leaving Outer",
|
||||
] == logs
|
||||
|
||||
|
||||
@pytest.mark.parametrize("preserve_ordering", [True, False])
|
||||
@pytest.mark.parametrize(
|
||||
"buffer_size, expected_gen_time",
|
||||
[
|
||||
(0, 5.5), # 5 x 1s + 0.5s buffer
|
||||
(1, 7.5), # 3 x 1s + 2 x 2s (limited buffer delay) + 0.5s buffer
|
||||
(2, 5.5), # 5 x 1s + 0.5s buffer
|
||||
],
|
||||
)
|
||||
def test_make_async_gen_x(buffer_size: int, expected_gen_time, preserve_ordering):
|
||||
"""Tests that make_async_gen overlaps compute."""
|
||||
|
||||
num_items = 5
|
||||
|
||||
def gen(base_iterator):
|
||||
gen_start = time.perf_counter()
|
||||
|
||||
for i in base_iterator:
|
||||
time.sleep(1)
|
||||
yield i
|
||||
print(f">>> ({time.time()}) Generating {i}")
|
||||
|
||||
gen_finish = time.perf_counter()
|
||||
|
||||
# 0.5s buffer
|
||||
assert gen_finish - gen_start < expected_gen_time
|
||||
|
||||
def sleepy_udf(item):
|
||||
time.sleep(2)
|
||||
return item
|
||||
|
||||
iterator = make_async_gen(
|
||||
base_iterator=iter(range(num_items)),
|
||||
fn=gen,
|
||||
preserve_ordering=preserve_ordering,
|
||||
num_workers=1,
|
||||
buffer_size=buffer_size,
|
||||
)
|
||||
|
||||
outputs = []
|
||||
|
||||
iter_start = time.perf_counter()
|
||||
for item in iterator:
|
||||
print(f">>> ({time.time()}) Iterating over {item}")
|
||||
print(item)
|
||||
outputs.append(sleepy_udf(item))
|
||||
iter_finish = time.perf_counter()
|
||||
|
||||
dur_s = iter_finish - iter_start
|
||||
|
||||
print(f">>> Took {dur_s}")
|
||||
|
||||
# 1s to yield first element
|
||||
# 10s to iterate t/h all 5
|
||||
# 0.5s extra buffer
|
||||
assert dur_s < num_items * 2 + 1.5
|
||||
|
||||
# Assert ordering is preserved
|
||||
assert outputs == list(range(num_items))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("preserve_ordering", [True, False])
|
||||
@pytest.mark.parametrize("buffer_size", [0, 1, 2])
|
||||
def test_make_async_gen_multiple_threads(buffer_size: int, preserve_ordering):
|
||||
"""Tests that using multiple threads can overlap compute even more."""
|
||||
|
||||
num_items = 5
|
||||
|
||||
gen_sleep = 2
|
||||
iter_sleep = 3
|
||||
|
||||
def gen(base_iterator):
|
||||
for i in base_iterator:
|
||||
time.sleep(gen_sleep)
|
||||
yield i
|
||||
|
||||
def sleep_udf(item):
|
||||
time.sleep(iter_sleep)
|
||||
return item
|
||||
|
||||
# All 5 items should be fetched concurrently.
|
||||
iterator = make_async_gen(
|
||||
base_iterator=iter(range(num_items)),
|
||||
fn=gen,
|
||||
preserve_ordering=preserve_ordering,
|
||||
num_workers=5,
|
||||
buffer_size=buffer_size,
|
||||
)
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Only sleep for first item.
|
||||
elements = [sleep_udf(next(iterator))] + list(iterator)
|
||||
|
||||
# All subsequent items should already be prefetched and should be ready.
|
||||
end_time = time.time()
|
||||
|
||||
# Assert ordering is preserved
|
||||
if preserve_ordering:
|
||||
assert elements == list(range(num_items))
|
||||
|
||||
# - 2 second for every worker to handle their single element
|
||||
# - 3 seconds for overlapping one
|
||||
# - 0.5 seconds buffer
|
||||
assert end_time - start_time < gen_sleep + iter_sleep + 0.5
|
||||
|
||||
|
||||
@pytest.mark.parametrize("preserve_ordering", [True, False])
|
||||
@pytest.mark.parametrize("buffer_size", [0, 1, 2])
|
||||
def test_make_async_gen_multiple_threads_unfinished(
|
||||
buffer_size: int, preserve_ordering
|
||||
):
|
||||
"""Tests that using multiple threads can overlap compute even more.
|
||||
Do not finish iteration with break in the middle.
|
||||
"""
|
||||
|
||||
num_items = 5
|
||||
|
||||
def gen(base_iterator):
|
||||
for i in base_iterator:
|
||||
time.sleep(4)
|
||||
yield i
|
||||
|
||||
def sleep_udf(item):
|
||||
time.sleep(5)
|
||||
return item
|
||||
|
||||
# All 5 items should be fetched concurrently.
|
||||
iterator = make_async_gen(
|
||||
base_iterator=iter(range(num_items)),
|
||||
fn=gen,
|
||||
preserve_ordering=preserve_ordering,
|
||||
num_workers=5,
|
||||
buffer_size=buffer_size,
|
||||
)
|
||||
|
||||
start_time = time.time()
|
||||
|
||||
# Only sleep for first item.
|
||||
sleep_udf(next(iterator))
|
||||
|
||||
# All subsequent items should already be prefetched and should be ready.
|
||||
for i, _ in enumerate(iterator):
|
||||
if i > 2:
|
||||
break
|
||||
end_time = time.time()
|
||||
|
||||
# 4 second for first item, 5 seconds for udf, 0.5 seconds buffer
|
||||
assert end_time - start_time < 9.5
|
||||
|
||||
|
||||
def test_calculate_ref_hits(ray_start_regular_shared):
|
||||
refs = [ray.put(0), ray.put(1)]
|
||||
hits, misses, unknowns = _calculate_ref_hits(refs)
|
||||
# With ctx.enable_get_object_locations_for_metrics set to False
|
||||
# by default, `_calculate_ref_hits` returns -1 for all, since
|
||||
# getting object locations is disabled.
|
||||
assert hits == 0
|
||||
assert misses == 0
|
||||
assert unknowns == 0
|
||||
|
||||
ctx = ray.data.DataContext.get_current()
|
||||
prev_enable_get_object_locations_for_metrics = (
|
||||
ctx.enable_get_object_locations_for_metrics
|
||||
)
|
||||
try:
|
||||
ctx.enable_get_object_locations_for_metrics = True
|
||||
hits, misses, unknowns = _calculate_ref_hits(refs)
|
||||
assert hits == 2
|
||||
assert misses == 0
|
||||
assert unknowns == 0
|
||||
finally:
|
||||
ctx.enable_get_object_locations_for_metrics = (
|
||||
prev_enable_get_object_locations_for_metrics
|
||||
)
|
||||
|
||||
|
||||
def _identity(it: Iterator[int]) -> Iterator[int]:
|
||||
return it
|
||||
|
||||
|
||||
def _duplicate_each(it: Iterator[int]) -> Iterator[int]:
|
||||
for item in it:
|
||||
yield item
|
||||
yield item
|
||||
|
||||
|
||||
class TestIterThreaded:
|
||||
"""Unit tests for ``iter_threaded``."""
|
||||
|
||||
@pytest.mark.parametrize("num_workers", [1, 2, 4])
|
||||
@pytest.mark.parametrize("output_buffer_size", [1, 2, 4])
|
||||
@pytest.mark.parametrize(
|
||||
"fn,multiplier",
|
||||
[(_identity, 1), (_duplicate_each, 2)],
|
||||
ids=["identity", "duplicate"],
|
||||
)
|
||||
def test_processes_all_exactly_once(
|
||||
self,
|
||||
num_workers: int,
|
||||
output_buffer_size: int,
|
||||
fn: Callable[[Iterator[int]], Iterator[int]],
|
||||
multiplier: int,
|
||||
):
|
||||
"""Every input item is consumed and produced exactly the expected
|
||||
number of times across the worker pool (no losses, no duplicates).
|
||||
Output ordering is not required."""
|
||||
items = list(range(50))
|
||||
output = list(
|
||||
iter_threaded(
|
||||
iter(items),
|
||||
fn,
|
||||
num_workers=num_workers,
|
||||
output_buffer_size=output_buffer_size,
|
||||
)
|
||||
)
|
||||
assert len(output) == len(items) * multiplier
|
||||
assert Counter(output) == Counter(items * multiplier)
|
||||
|
||||
def test_stateful_base_iterator_thread_safe(self):
|
||||
"""Python generators are not thread-safe; concurrent ``next()``
|
||||
calls raise ``ValueError: generator already executing`` without
|
||||
a lock. This test passes only if ``iter_threaded`` serializes
|
||||
the underlying ``next()`` properly."""
|
||||
|
||||
def stateful_gen():
|
||||
for i in range(200):
|
||||
# Encourage interleaving across workers.
|
||||
time.sleep(0.001)
|
||||
yield i
|
||||
|
||||
output = list(iter_threaded(stateful_gen(), _identity, num_workers=4))
|
||||
assert sorted(output) == list(range(200))
|
||||
|
||||
@pytest.mark.parametrize("num_workers", [1, 4])
|
||||
def test_fn_exception_propagates(self, num_workers: int):
|
||||
"""An exception raised inside ``fn`` is surfaced to the consumer
|
||||
rather than silently swallowed or hanging the iterator."""
|
||||
|
||||
def fn(it: Iterator[int]) -> Iterator[int]:
|
||||
for i, item in enumerate(it):
|
||||
if i >= 3:
|
||||
raise ValueError("boom")
|
||||
yield item
|
||||
|
||||
it = iter_threaded(iter(range(100)), fn, num_workers=num_workers)
|
||||
with pytest.raises(ValueError, match="boom"):
|
||||
list(it)
|
||||
|
||||
@pytest.mark.parametrize("num_workers", [1, 4])
|
||||
def test_non_generator_fn_construction_raises(self, num_workers: int):
|
||||
"""When ``fn`` is a non-generator function that raises during
|
||||
construction (e.g., setup code before returning the iterator), the
|
||||
exception must surface to the consumer rather than hang. Regression
|
||||
for the case where ``fn(_locked_iter())`` was called outside the
|
||||
worker's try/finally."""
|
||||
|
||||
def fn(it: Iterator[int]) -> Iterator[int]:
|
||||
# Body runs eagerly at call time (not a generator function).
|
||||
raise ValueError("boom in fn construction")
|
||||
|
||||
it = iter_threaded(iter(range(100)), fn, num_workers=num_workers)
|
||||
with pytest.raises(ValueError, match="boom in fn construction"):
|
||||
list(it)
|
||||
|
||||
@pytest.mark.parametrize("num_workers", [1, 4])
|
||||
def test_consumer_break_stops_workers(self, num_workers: int):
|
||||
"""When the consumer breaks early and the iterator is no longer
|
||||
referenced, CPython GCs the generator immediately, which runs the
|
||||
``finally: stopped.set()`` cleanup path. Worker threads should
|
||||
terminate within the ``_put`` poll interval (~100ms) rather than
|
||||
leak."""
|
||||
|
||||
def slow_fn(it: Iterator[int]) -> Iterator[int]:
|
||||
for item in it:
|
||||
time.sleep(0.05)
|
||||
yield item
|
||||
|
||||
# Inline so `break` drops the last reference → GC → finally.
|
||||
for i, _ in enumerate(
|
||||
iter_threaded(iter(range(10_000)), slow_fn, num_workers=num_workers)
|
||||
):
|
||||
if i >= 5:
|
||||
break
|
||||
|
||||
# Workers poll `stopped` every 100ms inside `_put`; give a generous
|
||||
# margin for CI under load.
|
||||
deadline = time.time() + 5.0
|
||||
while time.time() < deadline:
|
||||
alive = [t for t in threading.enumerate() if t.name == "iter_threaded"]
|
||||
if not alive:
|
||||
break
|
||||
time.sleep(0.05)
|
||||
else:
|
||||
pytest.fail(
|
||||
f"iter_threaded workers did not exit within 5s: "
|
||||
f"{[t.name for t in threading.enumerate() if t.name == 'iter_threaded']}"
|
||||
)
|
||||
|
||||
def test_num_workers_validation(self):
|
||||
with pytest.raises(ValueError, match="num_workers must be at least 1"):
|
||||
list(iter_threaded(iter([1]), _identity, num_workers=0))
|
||||
|
||||
def test_output_buffer_size_validation(self):
|
||||
with pytest.raises(ValueError, match="output_buffer_size must be at least 1"):
|
||||
list(iter_threaded(iter([1]), _identity, output_buffer_size=0))
|
||||
|
||||
def test_empty_base_iterator(self):
|
||||
output = list(iter_threaded(iter([]), _identity, num_workers=4))
|
||||
assert output == []
|
||||
|
||||
@pytest.mark.parametrize("num_workers,output_buffer_size", [(1, 1), (2, 2), (4, 2)])
|
||||
def test_in_flight_items_bounded_by_output_buffer_size(
|
||||
self, num_workers: int, output_buffer_size: int
|
||||
):
|
||||
"""Without consumption, workers must not pull more than
|
||||
``output_buffer_size`` items from the base iterator. Pulled-but-not-
|
||||
consumed items are 'in flight', and the bound caps them."""
|
||||
pulled = 0
|
||||
pulled_lock = threading.Lock()
|
||||
|
||||
def counting_iter() -> Iterator[int]:
|
||||
nonlocal pulled
|
||||
for i in range(1_000_000):
|
||||
with pulled_lock:
|
||||
pulled += 1
|
||||
yield i
|
||||
|
||||
it = iter_threaded(
|
||||
counting_iter(),
|
||||
_identity,
|
||||
num_workers=num_workers,
|
||||
output_buffer_size=output_buffer_size,
|
||||
)
|
||||
|
||||
# Trigger the generator body (which starts the workers), then stop
|
||||
# consuming. Workers fill in-flight up to the bound and then block
|
||||
# on _acquire_slot.
|
||||
next(it)
|
||||
time.sleep(0.3)
|
||||
|
||||
with pulled_lock:
|
||||
# Consumer took 1 → in-flight ≤ K. Plus the 1 already consumed.
|
||||
assert (
|
||||
pulled <= 1 + output_buffer_size
|
||||
), f"Pulled {pulled}, expected <= {1 + output_buffer_size}"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,815 @@
|
||||
import copy
|
||||
import os
|
||||
import posixpath
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray._common.test_utils import wait_for_condition
|
||||
from ray._private.internal_api import get_memory_info_reply, get_state_from_address
|
||||
from ray.data._internal.execution.block_ref_counter import BlockRefCounter
|
||||
from ray.data._internal.execution.operators.base_physical_operator import (
|
||||
AllToAllOperator,
|
||||
)
|
||||
from ray.data._internal.tensor_extensions.arrow import ArrowTensorArray
|
||||
from ray.data._internal.utils.arrow_utils import get_pyarrow_version
|
||||
from ray.data.block import BlockExecStats, BlockMetadata
|
||||
from ray.data.constants import TENSOR_COLUMN_NAME
|
||||
from ray.data.context import DEFAULT_TARGET_MAX_BLOCK_SIZE, DataContext, ShuffleStrategy
|
||||
from ray.data.tests.mock_server import * # noqa
|
||||
|
||||
# Trigger pytest hook to automatically zip test cluster logs to archive dir on failure
|
||||
from ray.tests.conftest import * # noqa
|
||||
from ray.tests.conftest import _ray_start
|
||||
from ray.util.debug import reset_log_once
|
||||
from ray.util.state import list_actors
|
||||
|
||||
|
||||
def mock_all_to_all_op(input_op, name="MockAllToAll"):
|
||||
"""Create a mock AllToAllOperator for testing.
|
||||
|
||||
Creates an AllToAllOperator which is NOT eligible for resource allocation
|
||||
(throttling_disabled=True) but is a blocking materializing operator.
|
||||
|
||||
Note: Creating this operator automatically adds it to input_op._output_dependencies.
|
||||
"""
|
||||
op = AllToAllOperator(
|
||||
bulk_fn=MagicMock(),
|
||||
input_op=input_op,
|
||||
data_context=ray.data.DataContext.get_current(),
|
||||
name=name,
|
||||
)
|
||||
op.start = MagicMock(side_effect=lambda *_: None)
|
||||
return op
|
||||
|
||||
|
||||
def noop_counter():
|
||||
"""BlockRefCounter that works without a Ray cluster."""
|
||||
return BlockRefCounter(add_object_out_of_scope_callback=lambda *_: True)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def data_context_override(request):
|
||||
overrides = getattr(request, "param", {})
|
||||
|
||||
ctx = DataContext.get_current()
|
||||
copy = ctx.copy()
|
||||
|
||||
for k, v in overrides.items():
|
||||
assert hasattr(ctx, k), f"Key '{k}' not found in DataContext"
|
||||
|
||||
setattr(ctx, k, v)
|
||||
|
||||
yield ctx
|
||||
|
||||
DataContext._set_current(copy)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def ray_start_2_cpus_shared(request):
|
||||
param = getattr(request, "param", {})
|
||||
with _ray_start(num_cpus=2, **param) as res:
|
||||
yield res
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def ray_start_10_cpus_shared(request):
|
||||
param = getattr(request, "param", {})
|
||||
with _ray_start(num_cpus=10, **param) as res:
|
||||
yield res
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def aws_credentials():
|
||||
import os
|
||||
|
||||
# Credentials dict that can be passed as kwargs to pa.fs.S3FileSystem
|
||||
credentials = dict(
|
||||
access_key="testing", secret_key="testing", session_token="testing"
|
||||
)
|
||||
|
||||
old_env = os.environ
|
||||
os.environ["AWS_ACCESS_KEY_ID"] = credentials["access_key"]
|
||||
os.environ["AWS_SECRET_ACCESS_KEY"] = credentials["secret_key"]
|
||||
os.environ["AWS_SECURITY_TOKEN"] = "testing"
|
||||
os.environ["AWS_SESSION_TOKEN"] = credentials["session_token"]
|
||||
|
||||
yield credentials
|
||||
os.environ = old_env
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def data_dir():
|
||||
yield "test_data"
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def data_dir_with_space():
|
||||
yield "test data"
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def data_dir_with_special_chars():
|
||||
yield "test data#fragment?query=test/"
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def s3_path(tmp_path, data_dir):
|
||||
yield "s3://" + posixpath.join(tmp_path, data_dir).strip("/")
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def s3_path_with_space(tmp_path, data_dir_with_space):
|
||||
yield "s3://" + posixpath.join(tmp_path, data_dir_with_space).strip("/")
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def s3_path_with_special_chars(tmp_path, data_dir_with_special_chars):
|
||||
yield "s3://" + posixpath.join(tmp_path, data_dir_with_special_chars).lstrip("/")
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def s3_path_with_anonymous_crendential(tmp_path, data_dir):
|
||||
yield "s3://" + "anonymous@" + posixpath.join(tmp_path, data_dir).lstrip("/")
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def s3_fs(aws_credentials, s3_server, s3_path):
|
||||
yield from _s3_fs(aws_credentials, s3_server, s3_path)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def s3_fs_with_space(aws_credentials, s3_server, s3_path_with_space):
|
||||
yield from _s3_fs(aws_credentials, s3_server, s3_path_with_space)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def s3_fs_with_special_chars(aws_credentials, s3_server, s3_path_with_special_chars):
|
||||
yield from _s3_fs(aws_credentials, s3_server, s3_path_with_special_chars)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def s3_fs_with_anonymous_crendential(
|
||||
aws_credentials, s3_server, s3_path_with_anonymous_crendential
|
||||
):
|
||||
yield from _s3_fs(aws_credentials, s3_server, s3_path_with_anonymous_crendential)
|
||||
|
||||
|
||||
def _s3_fs(aws_credentials, s3_server, s3_path):
|
||||
import urllib.parse
|
||||
|
||||
from packaging.version import parse as parse_version
|
||||
|
||||
kwargs = aws_credentials.copy()
|
||||
|
||||
if get_pyarrow_version() >= parse_version("9.0.0"):
|
||||
kwargs["allow_bucket_creation"] = True
|
||||
kwargs["allow_bucket_deletion"] = True
|
||||
|
||||
fs = None
|
||||
try:
|
||||
fs = pa.fs.S3FileSystem(
|
||||
region="us-west-2",
|
||||
endpoint_override=s3_server,
|
||||
**kwargs,
|
||||
)
|
||||
if s3_path.startswith("s3://"):
|
||||
if "@" in s3_path:
|
||||
s3_path = s3_path.split("@")[-1]
|
||||
else:
|
||||
s3_path = s3_path[len("s3://") :]
|
||||
s3_path = urllib.parse.quote(s3_path)
|
||||
fs.create_dir(s3_path)
|
||||
yield fs
|
||||
|
||||
finally:
|
||||
# Explicit cleanup for S3FileSystem resources
|
||||
if fs is not None:
|
||||
try:
|
||||
# Clean up test directory if it exists
|
||||
try:
|
||||
file_info = fs.get_file_info(s3_path)
|
||||
if file_info.type != pa.fs.FileType.NotFound:
|
||||
fs.delete_dir(s3_path)
|
||||
except (OSError, pa.lib.ArrowIOError):
|
||||
# Directory doesn't exist or can't be deleted, that's fine
|
||||
pass
|
||||
except Exception as e:
|
||||
print(f"Warning: S3 filesystem cleanup error: {e}")
|
||||
finally:
|
||||
fs = None
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def local_path(tmp_path, data_dir):
|
||||
path = os.path.join(tmp_path, data_dir)
|
||||
os.mkdir(path)
|
||||
yield path
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def local_fs():
|
||||
yield pa.fs.LocalFileSystem()
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def base_partitioned_df():
|
||||
yield pd.DataFrame(
|
||||
{"one": [1, 1, 1, 3, 3, 3], "two": ["a", "b", "c", "e", "f", "g"]}
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def write_partitioned_df():
|
||||
def _write_partitioned_df(
|
||||
df,
|
||||
partition_keys,
|
||||
partition_path_encoder,
|
||||
file_writer_fn,
|
||||
file_name_suffix="_1",
|
||||
):
|
||||
import urllib.parse
|
||||
|
||||
df_partitions = [df for _, df in df.groupby(partition_keys, as_index=False)]
|
||||
paths = []
|
||||
for df_partition in df_partitions:
|
||||
partition_values = []
|
||||
for key in partition_keys:
|
||||
partition_values.append(str(df_partition[key].iloc[0]))
|
||||
path = partition_path_encoder(partition_values)
|
||||
partition_path_encoder.scheme.resolved_filesystem.create_dir(path)
|
||||
base_dir = partition_path_encoder.scheme.base_dir
|
||||
parsed_base_dir = urllib.parse.urlparse(base_dir)
|
||||
file_name = f"test_{file_name_suffix}.tmp"
|
||||
if parsed_base_dir.scheme:
|
||||
# replace the protocol removed by the partition path generator
|
||||
path = posixpath.join(f"{parsed_base_dir.scheme}://{path}", file_name)
|
||||
else:
|
||||
path = os.path.join(path, file_name)
|
||||
file_writer_fn(df_partition, path)
|
||||
paths.append(path)
|
||||
return paths
|
||||
|
||||
yield _write_partitioned_df
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def restore_data_context(request):
|
||||
"""Restore any DataContext changes after the test runs"""
|
||||
ctx = ray.data.context.DataContext.get_current()
|
||||
original = copy.deepcopy(ctx)
|
||||
yield ctx
|
||||
ray.data.context.DataContext._set_current(original)
|
||||
|
||||
|
||||
def _get_supported_tensor_formats():
|
||||
"""Get list of supported tensor formats based on PyArrow version.
|
||||
|
||||
Returns V1, V2, and ARROW_NATIVE only if PyArrow >= 16 (which supports
|
||||
native FixedShapeTensorScalar, FixedShapeTensorType, FixedShapeTensorArray).
|
||||
"""
|
||||
from ray.data._internal.tensor_extensions.arrow import (
|
||||
MIN_PYARROW_VERSION_FIXED_SHAPE_TENSOR_SCALAR,
|
||||
FixedShapeTensorFormat,
|
||||
)
|
||||
|
||||
formats = [FixedShapeTensorFormat.V1, FixedShapeTensorFormat.V2]
|
||||
if get_pyarrow_version() >= MIN_PYARROW_VERSION_FIXED_SHAPE_TENSOR_SCALAR:
|
||||
formats.append(FixedShapeTensorFormat.ARROW_NATIVE)
|
||||
return formats
|
||||
|
||||
|
||||
@pytest.fixture(params=_get_supported_tensor_formats())
|
||||
def tensor_format(request):
|
||||
"""Fixture that yields supported tensor formats.
|
||||
|
||||
Yields V1, V2 for all PyArrow versions.
|
||||
Yields ARROW_NATIVE only when PyArrow >= 16.
|
||||
|
||||
This allows tests to use `tensor_format.to_type()` safely without
|
||||
needing fallback logic for unsupported PyArrow versions.
|
||||
"""
|
||||
return request.param
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tensor_format_context(request, restore_data_context, tensor_format):
|
||||
"""Fixture that sets the DataContext to use the given tensor format.
|
||||
|
||||
Combines restore_data_context with tensor_format to automatically
|
||||
configure the context for tensor format testing.
|
||||
"""
|
||||
ctx = ray.data.context.DataContext.get_current()
|
||||
ctx.arrow_fixed_shape_tensor_format = tensor_format
|
||||
return tensor_format
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def disable_fallback_to_object_extension(request, restore_data_context):
|
||||
"""Disables fallback to ArrowPythonObjectType"""
|
||||
ray.data.context.DataContext.get_current().enable_fallback_to_arrow_object_ext_type = (
|
||||
False
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(
|
||||
params=[
|
||||
s
|
||||
for s in ShuffleStrategy
|
||||
if s != ShuffleStrategy.GPU_SHUFFLE
|
||||
or os.environ.get("RAY_PYTEST_USE_GPU") == "1"
|
||||
]
|
||||
)
|
||||
def configure_shuffle_method(request):
|
||||
shuffle_strategy = request.param
|
||||
|
||||
ctx = ray.data.context.DataContext.get_current()
|
||||
|
||||
original_shuffle_strategy = ctx.shuffle_strategy
|
||||
original_default_hash_shuffle_parallelism = ctx.default_hash_shuffle_parallelism
|
||||
original_gpu_shuffle_num_actors = ctx.gpu_shuffle_num_actors
|
||||
|
||||
ctx.shuffle_strategy = shuffle_strategy
|
||||
|
||||
# NOTE: We override default parallelism for hash-based shuffling to
|
||||
# avoid excessive partitioning of the data (to achieve desired
|
||||
# parallelism
|
||||
if shuffle_strategy in [ShuffleStrategy.HASH_SHUFFLE, ShuffleStrategy.GPU_SHUFFLE]:
|
||||
ctx.default_hash_shuffle_parallelism = 8
|
||||
|
||||
if shuffle_strategy == ShuffleStrategy.GPU_SHUFFLE:
|
||||
ctx.gpu_shuffle_num_actors = 1
|
||||
|
||||
yield request.param
|
||||
|
||||
ctx.shuffle_strategy = original_shuffle_strategy
|
||||
ctx.default_hash_shuffle_parallelism = original_default_hash_shuffle_parallelism
|
||||
ctx.gpu_shuffle_num_actors = original_gpu_shuffle_num_actors
|
||||
|
||||
|
||||
@pytest.fixture(params=[True, False])
|
||||
def use_polars_sort(request):
|
||||
use_polars_sort = request.param
|
||||
|
||||
ctx = ray.data.context.DataContext.get_current()
|
||||
|
||||
original_use_polars = ctx.use_polars_sort
|
||||
|
||||
ctx.use_polars_sort = use_polars_sort
|
||||
|
||||
yield request.param
|
||||
|
||||
ctx.use_polars_sort = original_use_polars
|
||||
|
||||
|
||||
@pytest.fixture(params=[True, False])
|
||||
def enable_automatic_tensor_extension_cast(request):
|
||||
ctx = ray.data.context.DataContext.get_current()
|
||||
original = ctx.enable_tensor_extension_casting
|
||||
ctx.enable_tensor_extension_casting = request.param
|
||||
yield request.param
|
||||
ctx.enable_tensor_extension_casting = original
|
||||
|
||||
|
||||
@pytest.fixture(params=[True, False])
|
||||
def enable_auto_log_stats(request):
|
||||
ctx = ray.data.context.DataContext.get_current()
|
||||
original = ctx.enable_auto_log_stats
|
||||
ctx.enable_auto_log_stats = request.param
|
||||
yield request.param
|
||||
ctx.enable_auto_log_stats = original
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def reset_log_once_fixture():
|
||||
reset_log_once()
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture(params=[1024])
|
||||
def target_max_block_size(request):
|
||||
ctx = ray.data.context.DataContext.get_current()
|
||||
original = ctx.target_max_block_size
|
||||
ctx.target_max_block_size = request.param
|
||||
yield request.param
|
||||
ctx.target_max_block_size = original
|
||||
|
||||
|
||||
@pytest.fixture(params=[None, DEFAULT_TARGET_MAX_BLOCK_SIZE])
|
||||
def target_max_block_size_infinite_or_default(request):
|
||||
"""Fixture that sets target_max_block_size to None/DEFAULT_TARGET_MAX_BLOCK_SIZE and resets after test finishes."""
|
||||
ctx = ray.data.context.DataContext.get_current()
|
||||
original = ctx.target_max_block_size
|
||||
ctx.target_max_block_size = request.param
|
||||
yield
|
||||
ctx.target_max_block_size = original
|
||||
|
||||
|
||||
@pytest.fixture(params=[None])
|
||||
def target_max_block_size_infinite(request):
|
||||
"""Fixture that sets target_max_block_size to None and resets after test finishes."""
|
||||
ctx = ray.data.context.DataContext.get_current()
|
||||
original = ctx.target_max_block_size
|
||||
ctx.target_max_block_size = request.param
|
||||
yield
|
||||
ctx.target_max_block_size = original
|
||||
|
||||
|
||||
# ===== Pandas dataset formats =====
|
||||
@pytest.fixture(scope="function")
|
||||
def ds_pandas_single_column_format(ray_start_regular_shared):
|
||||
in_df = pd.DataFrame({"column_1": [1, 2, 3, 4]})
|
||||
yield ray.data.from_pandas(in_df)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def ds_pandas_multi_column_format(ray_start_regular_shared):
|
||||
in_df = pd.DataFrame({"column_1": [1, 2, 3, 4], "column_2": [1, -1, 1, -1]})
|
||||
yield ray.data.from_pandas(in_df)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def ds_pandas_list_multi_column_format(ray_start_regular_shared):
|
||||
in_df = pd.DataFrame({"column_1": [1], "column_2": [1]})
|
||||
yield ray.data.from_pandas([in_df] * 4)
|
||||
|
||||
|
||||
# ===== Arrow dataset formats =====
|
||||
@pytest.fixture(scope="function")
|
||||
def ds_arrow_single_column_format(ray_start_regular_shared):
|
||||
yield ray.data.from_arrow(pa.table({"column_1": [1, 2, 3, 4]}))
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def ds_arrow_single_column_tensor_format(ray_start_regular_shared):
|
||||
yield ray.data.from_arrow(
|
||||
pa.table(
|
||||
{
|
||||
TENSOR_COLUMN_NAME: ArrowTensorArray.from_numpy(
|
||||
np.arange(16).reshape((4, 2, 2))
|
||||
)
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def ds_arrow_multi_column_format(ray_start_regular_shared):
|
||||
yield ray.data.from_arrow(
|
||||
pa.table(
|
||||
{
|
||||
"column_1": [1, 2, 3, 4],
|
||||
"column_2": [1, -1, 1, -1],
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def ds_list_arrow_multi_column_format(ray_start_regular_shared):
|
||||
yield ray.data.from_arrow([pa.table({"column_1": [1], "column_2": [1]})] * 4)
|
||||
|
||||
|
||||
# ===== Numpy dataset formats =====
|
||||
@pytest.fixture(scope="function")
|
||||
def ds_numpy_single_column_tensor_format(ray_start_regular_shared):
|
||||
yield ray.data.from_numpy(np.arange(16).reshape((4, 2, 2)))
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def ds_numpy_list_of_ndarray_tensor_format(ray_start_regular_shared):
|
||||
yield ray.data.from_numpy([np.arange(4).reshape((1, 2, 2))] * 4)
|
||||
|
||||
|
||||
# ===== Observability & Logging Fixtures =====
|
||||
@pytest.fixture
|
||||
def op_two_block():
|
||||
block_params = {
|
||||
"num_rows": [10000, 5000],
|
||||
"size_bytes": [100, 50],
|
||||
"wall_time": [5, 10],
|
||||
"cpu_time": [1.2, 3.4],
|
||||
"udf_time": [1.1, 1.7],
|
||||
"node_id": ["a1", "b2"],
|
||||
"task_idx": [0, 1],
|
||||
}
|
||||
|
||||
block_delay = 20
|
||||
block_meta_list = []
|
||||
for i in range(len(block_params["num_rows"])):
|
||||
start_time_s = time.perf_counter() + i * block_delay
|
||||
# The blocks are executing from [0, 5] and [20, 30].
|
||||
block_exec_stats = BlockExecStats(
|
||||
start_time_s=start_time_s,
|
||||
end_time_s=start_time_s + block_params["wall_time"][i],
|
||||
wall_time_s=block_params["wall_time"][i],
|
||||
cpu_time_s=block_params["cpu_time"][i],
|
||||
udf_time_s=block_params["udf_time"][i],
|
||||
node_id=block_params["node_id"][i],
|
||||
task_idx=block_params["task_idx"][i],
|
||||
)
|
||||
block_meta_list.append(
|
||||
BlockMetadata(
|
||||
num_rows=block_params["num_rows"][i],
|
||||
size_bytes=block_params["size_bytes"][i],
|
||||
input_files=None,
|
||||
exec_stats=block_exec_stats,
|
||||
)
|
||||
)
|
||||
return block_params, block_meta_list
|
||||
|
||||
|
||||
def equals_or_true(count, expected_count):
|
||||
if isinstance(expected_count, int):
|
||||
if count != expected_count:
|
||||
return False
|
||||
else:
|
||||
if not expected_count(count):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
class CoreExecutionMetrics:
|
||||
def __init__(self, task_count=None, object_store_stats=None, actor_count=None):
|
||||
self.task_count = task_count
|
||||
self.object_store_stats = object_store_stats
|
||||
self.actor_count = actor_count
|
||||
|
||||
def get_task_count(self):
|
||||
return self.task_count
|
||||
|
||||
def get_object_store_stats(self):
|
||||
return self.object_store_stats
|
||||
|
||||
def get_actor_count(self):
|
||||
return self.actor_count
|
||||
|
||||
def _assert_count_equals(self, actual_count, expected_count):
|
||||
diff = {}
|
||||
# Check that all tasks in expected tasks match those in actual task
|
||||
# count.
|
||||
for name, count in expected_count.items():
|
||||
if not equals_or_true(actual_count[name], count):
|
||||
diff[name] = (actual_count[name], count)
|
||||
|
||||
assert len(diff) == 0, "\nTask diff:\n" + "\n".join(
|
||||
f" - {key}: expected {val[1]}, got {val[0]}" for key, val in diff.items()
|
||||
)
|
||||
|
||||
def assert_task_metrics(self, expected_metrics):
|
||||
"""
|
||||
Assert equality to the given { <task name>: <task count> }.
|
||||
A lambda that takes in the count and returns a bool to assert can also
|
||||
be given instead of an integer task count.
|
||||
|
||||
An empty dict means that we expected no tasks to run. Pass None to skip
|
||||
the check.
|
||||
"""
|
||||
if expected_metrics.get_task_count() is None:
|
||||
return
|
||||
|
||||
expected_task_count = expected_metrics.get_task_count()
|
||||
actual_task_count = self.get_task_count()
|
||||
self._assert_count_equals(actual_task_count, expected_task_count)
|
||||
|
||||
def assert_object_store_metrics(self, expected_metrics):
|
||||
"""
|
||||
By default this checks that no objects were spilled or restored.
|
||||
Collected stats only apply to plasma store objects and exclude inlined
|
||||
or in-memory objects.
|
||||
|
||||
Caller can also override the following fields with a value or lambda to assert.
|
||||
- spilled_bytes_total
|
||||
- restored_bytes_total
|
||||
- cumulative_created_plasma_bytes
|
||||
- cumulative_created_plasma_objects
|
||||
"""
|
||||
expected_object_store_stats = (
|
||||
CoreExecutionMetrics.get_default_object_store_stats()
|
||||
)
|
||||
if expected_metrics.get_object_store_stats() is not None:
|
||||
for key, val in expected_metrics.get_object_store_stats().items():
|
||||
expected_object_store_stats[key] = val
|
||||
|
||||
actual_object_store_stats = self.get_object_store_stats()
|
||||
for key, val in expected_object_store_stats.items():
|
||||
print(f"{key}: Expect {val}, got {actual_object_store_stats[key]}")
|
||||
assert equals_or_true(
|
||||
actual_object_store_stats[key], val
|
||||
), f"{key}: expected {val} got {actual_object_store_stats[key]}"
|
||||
|
||||
def assert_actor_metrics(self, expected_metrics):
|
||||
if expected_metrics.get_actor_count() is None:
|
||||
return
|
||||
|
||||
expected_actor_count = expected_metrics.get_actor_count()
|
||||
actual_actor_count = self.get_actor_count()
|
||||
self._assert_count_equals(actual_actor_count, expected_actor_count)
|
||||
|
||||
@staticmethod
|
||||
def get_default_object_store_stats():
|
||||
return {
|
||||
"spilled_bytes_total": 0,
|
||||
"restored_bytes_total": 0,
|
||||
}
|
||||
|
||||
|
||||
class PhysicalCoreExecutionMetrics(CoreExecutionMetrics):
|
||||
"""Generated from a snapshot of the metrics collected by Ray Core during
|
||||
the physical execution.
|
||||
|
||||
NOTE(swang): Currently object store stats only include objects stored in
|
||||
plasma shared memory.
|
||||
"""
|
||||
|
||||
def __init__(self, last_snapshot=None):
|
||||
self.task_metrics = ray.util.state.list_tasks(detail=True, limit=10_000)
|
||||
self.last_snapshot = last_snapshot
|
||||
|
||||
memory_info = get_memory_info_reply(
|
||||
get_state_from_address(ray.get_runtime_context().gcs_address)
|
||||
)
|
||||
self.object_store_stats = {
|
||||
"spilled_bytes_total": memory_info.store_stats.spilled_bytes_total,
|
||||
"restored_bytes_total": memory_info.store_stats.restored_bytes_total,
|
||||
"cumulative_created_plasma_bytes": (
|
||||
memory_info.store_stats.cumulative_created_bytes
|
||||
),
|
||||
"cumulative_created_plasma_objects": (
|
||||
memory_info.store_stats.cumulative_created_objects
|
||||
),
|
||||
}
|
||||
|
||||
self.actor_metrics = list_actors(limit=10_000)
|
||||
|
||||
def clear_task_count(self):
|
||||
self.task_metrics = []
|
||||
|
||||
def clear_object_store_stats(self):
|
||||
self.object_store_stats = {}
|
||||
|
||||
def clear_actor_count(self):
|
||||
self.actor_metrics = []
|
||||
|
||||
def get_task_count(self):
|
||||
task_count = defaultdict(int)
|
||||
tasks = self.task_metrics
|
||||
tasks = [t for t in tasks if t.name != "barrier"]
|
||||
|
||||
for task in tasks:
|
||||
task_count[task.name] += 1
|
||||
|
||||
# Filter out previous and dummy tasks.
|
||||
if self.last_snapshot is not None:
|
||||
prev_task_count = self.last_snapshot.get_task_count()
|
||||
if prev_task_count is not None:
|
||||
for name, count in prev_task_count.items():
|
||||
task_count[name] -= count
|
||||
if task_count[name] < 0:
|
||||
task_count[name] = 0
|
||||
return task_count
|
||||
|
||||
def get_actor_count(self):
|
||||
actor_count = defaultdict(int)
|
||||
for actor in self.actor_metrics:
|
||||
actor_count[actor.class_name] += 1
|
||||
if self.last_snapshot is not None:
|
||||
prev_actor_count = self.last_snapshot.get_actor_count()
|
||||
if prev_actor_count is not None:
|
||||
for name, count in prev_actor_count.items():
|
||||
actor_count[name] -= count
|
||||
if actor_count[name] < 0:
|
||||
actor_count[name] = 0
|
||||
return actor_count
|
||||
|
||||
def get_object_store_stats(self):
|
||||
object_store_stats = self.object_store_stats.copy()
|
||||
if self.last_snapshot is not None:
|
||||
prev_object_store_stats = self.last_snapshot.get_object_store_stats()
|
||||
if prev_object_store_stats is not None:
|
||||
for key, val in prev_object_store_stats.items():
|
||||
object_store_stats[key] -= val
|
||||
return object_store_stats
|
||||
|
||||
|
||||
# Dummy task used to make sure that we wait until (most) stats are available.
|
||||
@ray.remote
|
||||
def barrier():
|
||||
time.sleep(1)
|
||||
return
|
||||
|
||||
|
||||
@ray.remote
|
||||
def warmup():
|
||||
time.sleep(1)
|
||||
return np.zeros(1024 * 1024, dtype=np.uint8)
|
||||
|
||||
|
||||
def task_metrics_flushed(refs):
|
||||
task_ids = [t.task_id for t in ray.util.state.list_tasks(limit=10_000)]
|
||||
# All tasks appear in the metrics.
|
||||
return all(ref.task_id().hex() in task_ids for ref in refs)
|
||||
|
||||
|
||||
def get_initial_core_execution_metrics_snapshot():
|
||||
# Warmup plasma store and workers.
|
||||
refs = [warmup.remote() for _ in range(int(ray.cluster_resources()["CPU"]))]
|
||||
ray.get(refs)
|
||||
wait_for_condition(lambda: task_metrics_flushed(refs))
|
||||
|
||||
last_snapshot = assert_core_execution_metrics_equals(
|
||||
CoreExecutionMetrics(
|
||||
task_count={"warmup": lambda count: True}, object_store_stats={}
|
||||
),
|
||||
last_snapshot=None,
|
||||
)
|
||||
return last_snapshot
|
||||
|
||||
|
||||
def assert_core_execution_metrics_equals(
|
||||
expected_metrics: CoreExecutionMetrics,
|
||||
last_snapshot=None,
|
||||
):
|
||||
# Wait for one task per CPU to finish to prevent a race condition where not
|
||||
# all of the task metrics have been collected yet.
|
||||
if expected_metrics.get_task_count() is not None:
|
||||
refs = [barrier.remote() for _ in range(int(ray.cluster_resources()["CPU"]))]
|
||||
ray.get(refs)
|
||||
wait_for_condition(lambda: task_metrics_flushed(refs))
|
||||
|
||||
metrics = PhysicalCoreExecutionMetrics(last_snapshot)
|
||||
metrics.assert_task_metrics(expected_metrics)
|
||||
metrics.assert_object_store_metrics(expected_metrics)
|
||||
metrics.assert_actor_metrics(expected_metrics)
|
||||
|
||||
# Return a last_snapshot to the current snapshot of metrics to make subsequent
|
||||
# queries easier. Don't return a last_snapshot for metrics that weren't asserted.
|
||||
last_snapshot = PhysicalCoreExecutionMetrics()
|
||||
if expected_metrics.get_task_count() is None:
|
||||
last_snapshot.clear_task_count()
|
||||
elif expected_metrics.get_object_store_stats() is None:
|
||||
last_snapshot.clear_object_store_stats()
|
||||
elif expected_metrics.get_actor_count() is None:
|
||||
last_snapshot.clear_actor_count()
|
||||
|
||||
return last_snapshot
|
||||
|
||||
|
||||
def assert_blocks_expected_in_plasma(
|
||||
last_snapshot,
|
||||
num_blocks_expected,
|
||||
block_size_expected=None,
|
||||
):
|
||||
total_bytes_expected = None
|
||||
|
||||
if block_size_expected is not None:
|
||||
total_bytes_expected = num_blocks_expected * block_size_expected
|
||||
|
||||
print(f"Expecting {total_bytes_expected} bytes, {num_blocks_expected} blocks")
|
||||
|
||||
def _assert(last_snapshot):
|
||||
assert_core_execution_metrics_equals(
|
||||
CoreExecutionMetrics(
|
||||
object_store_stats={
|
||||
"cumulative_created_plasma_objects": (
|
||||
lambda count: num_blocks_expected * 0.5
|
||||
<= count
|
||||
<= 1.5 * num_blocks_expected
|
||||
),
|
||||
"cumulative_created_plasma_bytes": (
|
||||
lambda count: total_bytes_expected is None
|
||||
or total_bytes_expected * 0.5
|
||||
<= count
|
||||
<= 1.5 * total_bytes_expected
|
||||
),
|
||||
},
|
||||
),
|
||||
last_snapshot,
|
||||
)
|
||||
return True
|
||||
|
||||
wait_for_condition(lambda: _assert(last_snapshot))
|
||||
|
||||
# Get the latest last_snapshot.
|
||||
last_snapshot = assert_core_execution_metrics_equals(
|
||||
CoreExecutionMetrics(
|
||||
object_store_stats={
|
||||
"cumulative_created_plasma_objects": lambda count: True,
|
||||
"cumulative_created_plasma_bytes": lambda count: True,
|
||||
}
|
||||
),
|
||||
last_snapshot,
|
||||
)
|
||||
|
||||
return last_snapshot
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True, scope="function")
|
||||
def log_internal_stack_trace_to_stdout(restore_data_context):
|
||||
ray.data.context.DataContext.get_current().log_internal_stack_trace_to_stdout = True
|
||||
@@ -0,0 +1,69 @@
|
||||
"""Test utilities for Databricks datasource tests."""
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
from ray.data._internal.datasource.databricks_credentials import (
|
||||
DatabricksCredentialProvider,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockResponse:
|
||||
"""Mock HTTP response for testing.
|
||||
|
||||
Args:
|
||||
status_code: HTTP status code. Defaults to 200.
|
||||
content: Response content as bytes. Defaults to None.
|
||||
_json_data: JSON response data. Defaults to None.
|
||||
raise_on_error: If True, raise_for_status() raises for status >= 400.
|
||||
Defaults to True.
|
||||
"""
|
||||
|
||||
status_code: int = 200
|
||||
content: Optional[bytes] = None
|
||||
_json_data: Optional[dict] = None
|
||||
raise_on_error: bool = field(default=True, repr=False)
|
||||
|
||||
def raise_for_status(self):
|
||||
"""Raise an exception if status code indicates an error."""
|
||||
if self.raise_on_error and self.status_code >= 400:
|
||||
raise Exception(f"HTTP Error {self.status_code}")
|
||||
|
||||
def json(self):
|
||||
"""Return the JSON data."""
|
||||
return self._json_data
|
||||
|
||||
|
||||
class RefreshableCredentialProvider(DatabricksCredentialProvider):
|
||||
"""A credential provider that simulates token refresh on invalidate.
|
||||
|
||||
Useful for testing 401 retry logic. When invalidate() is called,
|
||||
the token changes from initial_token to "refreshed_token".
|
||||
|
||||
Args:
|
||||
initial_token: The initial token value. Defaults to "expired_token".
|
||||
host: The host URL to return. Defaults to "https://test-host.databricks.com".
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
initial_token: str = "expired_token",
|
||||
host: str = "https://test-host.databricks.com",
|
||||
):
|
||||
self.current_token = initial_token
|
||||
self.invalidate_count = 0
|
||||
self._host = host
|
||||
|
||||
def get_token(self) -> str:
|
||||
"""Get the current token."""
|
||||
return self.current_token
|
||||
|
||||
def get_host(self) -> str:
|
||||
"""Get the host URL."""
|
||||
return self._host
|
||||
|
||||
def invalidate(self) -> None:
|
||||
"""Simulate token refresh by changing to 'refreshed_token'."""
|
||||
self.invalidate_count += 1
|
||||
self.current_token = "refreshed_token"
|
||||
@@ -0,0 +1,134 @@
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_dataframes():
|
||||
"""Fixture providing sample pandas DataFrames for testing.
|
||||
|
||||
Returns:
|
||||
tuple: (df1, df2) where df1 has 3 rows and df2 has 3 rows
|
||||
"""
|
||||
df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
||||
df2 = pd.DataFrame({"one": [4, 5, 6], "two": ["e", "f", "g"]})
|
||||
return df1, df2
|
||||
|
||||
|
||||
def test_from_arrow(ray_start_regular_shared, sample_dataframes):
|
||||
"""Test basic from_arrow functionality with single and multiple tables."""
|
||||
df1, df2 = sample_dataframes
|
||||
|
||||
ds = ray.data.from_arrow([pa.Table.from_pandas(df1), pa.Table.from_pandas(df2)])
|
||||
values = [(r["one"], r["two"]) for r in ds.take(6)]
|
||||
rows = [(r.one, r.two) for _, r in pd.concat([df1, df2]).iterrows()]
|
||||
assert values == rows
|
||||
# Check that metadata fetch is included in stats.
|
||||
assert "FromArrow" in ds.stats()
|
||||
|
||||
# test from single pyarrow table
|
||||
ds = ray.data.from_arrow(pa.Table.from_pandas(df1))
|
||||
values = [(r["one"], r["two"]) for r in ds.take(3)]
|
||||
rows = [(r.one, r.two) for _, r in df1.iterrows()]
|
||||
assert values == rows
|
||||
# Check that metadata fetch is included in stats.
|
||||
assert "FromArrow" in ds.stats()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"tables,override_num_blocks,expected_blocks,expected_rows",
|
||||
[
|
||||
# Single table scenarios
|
||||
("single", 1, 1, 3), # Single table, 1 block
|
||||
("single", 2, 2, 3), # Single table split into 2 blocks
|
||||
("single", 5, 5, 3), # Single table, more blocks than rows
|
||||
(
|
||||
"single",
|
||||
10,
|
||||
10,
|
||||
3,
|
||||
), # Edge case: 3 rows split into 10 blocks (creates empty blocks)
|
||||
# Multiple tables scenarios
|
||||
("multiple", 3, 3, 6), # Multiple tables split into 3 blocks
|
||||
("multiple", 10, 10, 6), # Multiple tables, more blocks than rows
|
||||
# Empty table scenarios
|
||||
("empty", 1, 1, 0), # Empty table, 1 block
|
||||
("empty", 5, 5, 0), # Empty table, more blocks than rows
|
||||
],
|
||||
)
|
||||
def test_from_arrow_override_num_blocks(
|
||||
ray_start_regular_shared,
|
||||
sample_dataframes,
|
||||
tables,
|
||||
override_num_blocks,
|
||||
expected_blocks,
|
||||
expected_rows,
|
||||
):
|
||||
"""Test from_arrow with override_num_blocks parameter."""
|
||||
df1, df2 = sample_dataframes
|
||||
empty_df = pd.DataFrame({"one": [], "two": []})
|
||||
|
||||
# Prepare tables based on test case
|
||||
if tables == "single":
|
||||
arrow_tables = pa.Table.from_pandas(df1)
|
||||
expected_data = [(r.one, r.two) for _, r in df1.iterrows()]
|
||||
elif tables == "multiple":
|
||||
arrow_tables = [pa.Table.from_pandas(df1), pa.Table.from_pandas(df2)]
|
||||
expected_data = [(r.one, r.two) for _, r in pd.concat([df1, df2]).iterrows()]
|
||||
elif tables == "empty":
|
||||
arrow_tables = pa.Table.from_pandas(empty_df)
|
||||
expected_data = []
|
||||
|
||||
# Create dataset with override_num_blocks
|
||||
ds = ray.data.from_arrow(arrow_tables, override_num_blocks=override_num_blocks)
|
||||
|
||||
# Verify number of blocks
|
||||
assert ds.num_blocks() == expected_blocks
|
||||
|
||||
# Verify row count
|
||||
assert ds.count() == expected_rows
|
||||
|
||||
# Verify data integrity (only for non-empty datasets)
|
||||
if expected_rows > 0:
|
||||
values = [(r["one"], r["two"]) for r in ds.take_all()]
|
||||
assert values == expected_data
|
||||
|
||||
|
||||
def test_from_arrow_refs(ray_start_regular_shared, sample_dataframes):
|
||||
df1, df2 = sample_dataframes
|
||||
ds = ray.data.from_arrow_refs(
|
||||
[ray.put(pa.Table.from_pandas(df1)), ray.put(pa.Table.from_pandas(df2))]
|
||||
)
|
||||
values = [(r["one"], r["two"]) for r in ds.take(6)]
|
||||
rows = [(r.one, r.two) for _, r in pd.concat([df1, df2]).iterrows()]
|
||||
assert values == rows
|
||||
# Check that metadata fetch is included in stats.
|
||||
assert "FromArrow" in ds.stats()
|
||||
|
||||
# test from single pyarrow table ref
|
||||
ds = ray.data.from_arrow_refs(ray.put(pa.Table.from_pandas(df1)))
|
||||
values = [(r["one"], r["two"]) for r in ds.take(3)]
|
||||
rows = [(r.one, r.two) for _, r in df1.iterrows()]
|
||||
assert values == rows
|
||||
# Check that metadata fetch is included in stats.
|
||||
assert "FromArrow" in ds.stats()
|
||||
|
||||
|
||||
def test_to_arrow_refs(ray_start_regular_shared):
|
||||
n = 5
|
||||
df = pd.DataFrame({"id": list(range(n))})
|
||||
ds = ray.data.range(n)
|
||||
dfds = pd.concat(
|
||||
[t.to_pandas() for t in ray.get(ds.to_arrow_refs())], ignore_index=True
|
||||
)
|
||||
assert df.equals(dfds)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,38 @@
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
NUM_AUDIO_FILES = 10
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def audio_uri():
|
||||
root = "s3://anonymous@air-example-data-2/6G-audio-data-LibriSpeech-train-clean-100-flac" # noqa: E501
|
||||
return [
|
||||
f"{root}/train-clean-100/5022/29411/5022-29411-{n:04}.flac"
|
||||
for n in range(NUM_AUDIO_FILES)
|
||||
]
|
||||
|
||||
|
||||
def test_read_audio(ray_start_regular_shared, audio_uri):
|
||||
ds = ray.data.read_audio(audio_uri)
|
||||
|
||||
# Verify basic audio properties
|
||||
assert ds.count() == NUM_AUDIO_FILES, ds.count()
|
||||
assert ds.schema().names == ["amplitude", "sample_rate"], ds.schema()
|
||||
|
||||
# Check the sample rate
|
||||
assert all(row["sample_rate"] == 16000 for row in ds.take_all())
|
||||
|
||||
for row in ds.take_all():
|
||||
assert row["amplitude"].ndim == 2
|
||||
assert row["amplitude"].shape[0] == 1
|
||||
assert row["amplitude"].dtype == np.float32
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,40 @@
|
||||
import os
|
||||
|
||||
import fastavro
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
|
||||
schema = {
|
||||
"type": "record",
|
||||
"name": "TestRecord",
|
||||
"fields": [{"name": "test_field", "type": "string"}],
|
||||
}
|
||||
|
||||
|
||||
def test_read_basic_avro_file(ray_start_regular_shared, tmp_path):
|
||||
path = os.path.join(tmp_path, "sample.avro")
|
||||
records = [{"test_field": "test_value1"}, {"test_field": "test_value2"}]
|
||||
with open(path, "wb") as out:
|
||||
fastavro.writer(out, schema, records)
|
||||
|
||||
ds = ray.data.read_avro(path)
|
||||
|
||||
expected = [{"test_field": "test_value1"}, {"test_field": "test_value2"}]
|
||||
assert ds.take_all() == expected
|
||||
|
||||
|
||||
def test_read_empty_avro_files(ray_start_regular_shared, tmp_path):
|
||||
path = os.path.join(tmp_path, "empty.avro")
|
||||
# Write an empty Avro file with the schema
|
||||
with open(path, "wb") as out:
|
||||
# Write the schema with no records
|
||||
fastavro.writer(out, schema, [])
|
||||
|
||||
ds = ray.data.read_avro(path)
|
||||
|
||||
assert ds.count() == 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main(["-v", __file__])
|
||||
@@ -0,0 +1,284 @@
|
||||
from typing import Iterator
|
||||
from unittest import mock
|
||||
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from google.api_core import exceptions, operation
|
||||
from google.cloud import bigquery, bigquery_storage
|
||||
from google.cloud.bigquery import job
|
||||
from google.cloud.bigquery_storage_v1.types import stream as gcbqs_stream
|
||||
|
||||
import ray
|
||||
from ray.data._internal.datasource.bigquery_datasink import BigQueryDatasink
|
||||
from ray.data._internal.datasource.bigquery_datasource import BigQueryDatasource
|
||||
from ray.data._internal.execution.interfaces.task_context import TaskContext
|
||||
from ray.data._internal.planner.plan_write_op import generate_collect_write_stats_fn
|
||||
from ray.data.block import Block
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.data.tests.mock_http_server import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
_TEST_GCP_PROJECT_ID = "mock-test-project-id"
|
||||
_TEST_BQ_DATASET_ID = "mockdataset"
|
||||
_TEST_BQ_TABLE_ID = "mocktable"
|
||||
_TEST_BQ_DATASET = _TEST_BQ_DATASET_ID + "." + _TEST_BQ_TABLE_ID
|
||||
_TEST_BQ_TEMP_DESTINATION = _TEST_GCP_PROJECT_ID + ".tempdataset.temptable"
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def bq_client_full_mock(monkeypatch):
|
||||
client_mock = mock.create_autospec(bigquery.Client)
|
||||
client_mock.return_value = client_mock
|
||||
|
||||
def bq_get_dataset_mock(dataset_id):
|
||||
if dataset_id != _TEST_BQ_DATASET_ID:
|
||||
raise exceptions.NotFound(
|
||||
"Dataset {} is not found. Please ensure that it exists.".format(
|
||||
_TEST_BQ_DATASET
|
||||
)
|
||||
)
|
||||
|
||||
def bq_get_table_mock(table_id):
|
||||
if table_id != _TEST_BQ_DATASET:
|
||||
raise exceptions.NotFound(
|
||||
"Table {} is not found. Please ensure that it exists.".format(
|
||||
_TEST_BQ_DATASET
|
||||
)
|
||||
)
|
||||
|
||||
def bq_create_dataset_mock(dataset_id, **kwargs):
|
||||
if dataset_id == "existingdataset":
|
||||
raise exceptions.Conflict("Dataset already exists")
|
||||
return mock.Mock(operation.Operation)
|
||||
|
||||
def bq_delete_table_mock(table, **kwargs):
|
||||
return None
|
||||
|
||||
def bq_query_mock(query):
|
||||
fake_job_ref = job._JobReference(
|
||||
"fake_job_id", _TEST_GCP_PROJECT_ID, "us-central1"
|
||||
)
|
||||
fake_query_job = job.QueryJob(fake_job_ref, query, None)
|
||||
fake_query_job.configuration.destination = _TEST_BQ_TEMP_DESTINATION
|
||||
return fake_query_job
|
||||
|
||||
client_mock.get_dataset = bq_get_dataset_mock
|
||||
client_mock.get_table = bq_get_table_mock
|
||||
client_mock.create_dataset = bq_create_dataset_mock
|
||||
client_mock.delete_table = bq_delete_table_mock
|
||||
client_mock.query = bq_query_mock
|
||||
|
||||
monkeypatch.setattr(bigquery, "Client", client_mock)
|
||||
return client_mock
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def bqs_client_full_mock(monkeypatch):
|
||||
client_mock = mock.create_autospec(bigquery_storage.BigQueryReadClient)
|
||||
client_mock.return_value = client_mock
|
||||
|
||||
def bqs_create_read_session(max_stream_count=0, **kwargs):
|
||||
read_session_proto = gcbqs_stream.ReadSession()
|
||||
read_session_proto.streams = [
|
||||
gcbqs_stream.ReadStream() for _ in range(max_stream_count)
|
||||
]
|
||||
return read_session_proto
|
||||
|
||||
client_mock.create_read_session = bqs_create_read_session
|
||||
|
||||
monkeypatch.setattr(bigquery_storage, "BigQueryReadClient", client_mock)
|
||||
client_mock.reset_mock()
|
||||
return client_mock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def bq_query_result_mock():
|
||||
with mock.patch.object(bigquery.job.QueryJob, "result") as query_result_mock:
|
||||
yield query_result_mock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def bq_query_result_mock_fail():
|
||||
with mock.patch.object(bigquery.job.QueryJob, "result") as query_result_mock_fail:
|
||||
query_result_mock_fail.side_effect = exceptions.BadRequest("400 Syntax error")
|
||||
yield query_result_mock_fail
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ray_get_mock():
|
||||
with mock.patch.object(ray, "get") as ray_get:
|
||||
ray_get.return_value = None
|
||||
yield ray_get
|
||||
|
||||
|
||||
class TestReadBigQuery:
|
||||
"""Tests for BigQuery Read."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"parallelism",
|
||||
[1, 2, 3, 4, 10, 100],
|
||||
)
|
||||
def test_create_read_tasks(self, parallelism):
|
||||
bq_ds = BigQueryDatasource(
|
||||
project_id=_TEST_GCP_PROJECT_ID,
|
||||
dataset=_TEST_BQ_DATASET,
|
||||
)
|
||||
read_tasks_list = bq_ds.get_read_tasks(parallelism)
|
||||
assert len(read_tasks_list) == parallelism
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"parallelism",
|
||||
[1, 2, 3, 4, 10, 100],
|
||||
)
|
||||
def test_create_reader_query(self, parallelism, bq_query_result_mock):
|
||||
bq_ds = BigQueryDatasource(
|
||||
project_id=_TEST_GCP_PROJECT_ID,
|
||||
query="SELECT * FROM mockdataset.mocktable",
|
||||
)
|
||||
read_tasks_list = bq_ds.get_read_tasks(parallelism)
|
||||
bq_query_result_mock.assert_called_once()
|
||||
assert len(read_tasks_list) == parallelism
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"parallelism",
|
||||
[1, 2, 3, 4, 10, 100],
|
||||
)
|
||||
def test_create_reader_query_bad_request(
|
||||
self,
|
||||
parallelism,
|
||||
bq_query_result_mock_fail,
|
||||
):
|
||||
bq_ds = BigQueryDatasource(
|
||||
project_id=_TEST_GCP_PROJECT_ID,
|
||||
query="SELECT * FROM mockdataset.mocktable",
|
||||
)
|
||||
with pytest.raises(exceptions.BadRequest):
|
||||
bq_ds.get_read_tasks(parallelism)
|
||||
bq_query_result_mock_fail.assert_called()
|
||||
|
||||
def test_dataset_query_kwargs_provided(self):
|
||||
with pytest.raises(ValueError) as exception:
|
||||
BigQueryDatasource(
|
||||
project_id=_TEST_GCP_PROJECT_ID,
|
||||
dataset=_TEST_BQ_DATASET,
|
||||
query="SELECT * FROM mockdataset.mocktable",
|
||||
)
|
||||
expected_message = (
|
||||
"Query and dataset kwargs cannot both be provided"
|
||||
+ " (must be mutually exclusive)."
|
||||
)
|
||||
assert str(exception.value) == expected_message
|
||||
|
||||
def test_create_reader_dataset_not_found(self):
|
||||
parallelism = 4
|
||||
bq_ds = BigQueryDatasource(
|
||||
project_id=_TEST_GCP_PROJECT_ID,
|
||||
dataset="nonexistentdataset.mocktable",
|
||||
)
|
||||
with pytest.raises(ValueError) as exception:
|
||||
bq_ds.get_read_tasks(parallelism)
|
||||
expected_message = (
|
||||
"Dataset nonexistentdataset is not found. Please ensure that it exists."
|
||||
)
|
||||
assert str(exception.value) == expected_message
|
||||
|
||||
def test_create_reader_table_not_found(self):
|
||||
parallelism = 4
|
||||
bq_ds = BigQueryDatasource(
|
||||
project_id=_TEST_GCP_PROJECT_ID,
|
||||
dataset="mockdataset.nonexistenttable",
|
||||
)
|
||||
with pytest.raises(ValueError) as exception:
|
||||
bq_ds.get_read_tasks(parallelism)
|
||||
expected_message = (
|
||||
"Table mockdataset.nonexistenttable is not found."
|
||||
+ " Please ensure that it exists."
|
||||
)
|
||||
assert str(exception.value) == expected_message
|
||||
|
||||
|
||||
class TestWriteBigQuery:
|
||||
"""Tests for BigQuery Write."""
|
||||
|
||||
def _extract_write_result(self, stats: Iterator[Block]):
|
||||
return dict(next(stats).iloc[0])
|
||||
|
||||
def test_write(self, ray_get_mock):
|
||||
bq_datasink = BigQueryDatasink(
|
||||
project_id=_TEST_GCP_PROJECT_ID,
|
||||
dataset=_TEST_BQ_DATASET,
|
||||
)
|
||||
arr = pa.array([2, 4, 5, 100])
|
||||
block = pa.Table.from_arrays([arr], names=["data"])
|
||||
ctx = TaskContext(1, "")
|
||||
bq_datasink.write(
|
||||
blocks=[block],
|
||||
ctx=ctx,
|
||||
)
|
||||
|
||||
collect_stats_fn = generate_collect_write_stats_fn()
|
||||
stats = collect_stats_fn([block], ctx)
|
||||
pd.testing.assert_frame_equal(
|
||||
next(stats),
|
||||
pd.DataFrame(
|
||||
{
|
||||
"num_rows": [4],
|
||||
"size_bytes": [32],
|
||||
"write_return": [None],
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
def test_write_dataset_exists(self, ray_get_mock):
|
||||
bq_datasink = BigQueryDatasink(
|
||||
project_id=_TEST_GCP_PROJECT_ID,
|
||||
dataset="existingdataset" + "." + _TEST_BQ_TABLE_ID,
|
||||
)
|
||||
arr = pa.array([2, 4, 5, 100])
|
||||
block = pa.Table.from_arrays([arr], names=["data"])
|
||||
ctx = TaskContext(1, "")
|
||||
bq_datasink.write(
|
||||
blocks=[block],
|
||||
ctx=ctx,
|
||||
)
|
||||
collect_stats_fn = generate_collect_write_stats_fn()
|
||||
stats = collect_stats_fn([block], ctx)
|
||||
pd.testing.assert_frame_equal(
|
||||
next(stats),
|
||||
pd.DataFrame(
|
||||
{
|
||||
"num_rows": [4],
|
||||
"size_bytes": [32],
|
||||
"write_return": [None],
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
def test_write_empty_block(self, ray_get_mock):
|
||||
"""Test that writing a zero-sized block doesn't crash.
|
||||
|
||||
See https://github.com/ray-project/ray/issues/51892
|
||||
"""
|
||||
bq_datasink = BigQueryDatasink(
|
||||
project_id=_TEST_GCP_PROJECT_ID,
|
||||
dataset=_TEST_BQ_DATASET,
|
||||
)
|
||||
# Create an empty block with schema but no rows
|
||||
block = pa.Table.from_arrays([pa.array([], type=pa.int64())], names=["data"])
|
||||
ctx = TaskContext(1, "")
|
||||
# This should not raise an error - empty blocks should be skipped
|
||||
bq_datasink.write(
|
||||
blocks=[block],
|
||||
ctx=ctx,
|
||||
)
|
||||
|
||||
# write() always calls ray.get(), but with an empty list since the
|
||||
# zero-row block is filtered out (no remote write tasks launched).
|
||||
ray_get_mock.assert_called_once_with([])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,53 @@
|
||||
import os
|
||||
from io import BytesIO
|
||||
|
||||
import pytest
|
||||
import snappy
|
||||
|
||||
import ray
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.data.tests.util import extract_values, gen_bin_files
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
|
||||
def test_read_binary_files(ray_start_regular_shared):
|
||||
with gen_bin_files(10) as (_, paths):
|
||||
ds = ray.data.read_binary_files(paths)
|
||||
for i, item in enumerate(ds.iter_rows()):
|
||||
expected = open(paths[i], "rb").read()
|
||||
assert expected == item["bytes"]
|
||||
# Test metadata ops.
|
||||
assert ds.count() == 10
|
||||
assert "bytes" in str(ds.schema()), ds
|
||||
assert "bytes" in str(ds), ds
|
||||
|
||||
|
||||
def test_read_binary_snappy(ray_start_regular_shared, tmp_path):
|
||||
path = os.path.join(tmp_path, "test_binary_snappy")
|
||||
os.mkdir(path)
|
||||
with open(os.path.join(path, "file"), "wb") as f:
|
||||
byte_str = "hello, world".encode()
|
||||
bytes = BytesIO(byte_str)
|
||||
snappy.stream_compress(bytes, f)
|
||||
ds = ray.data.read_binary_files(
|
||||
path,
|
||||
arrow_open_stream_args=dict(compression="snappy"),
|
||||
)
|
||||
assert sorted(extract_values("bytes", ds.take())) == [byte_str]
|
||||
|
||||
|
||||
def test_read_binary_snappy_inferred(ray_start_regular_shared, tmp_path):
|
||||
path = os.path.join(tmp_path, "test_binary_snappy_inferred")
|
||||
os.mkdir(path)
|
||||
with open(os.path.join(path, "file.snappy"), "wb") as f:
|
||||
byte_str = "hello, world".encode()
|
||||
bytes = BytesIO(byte_str)
|
||||
snappy.stream_compress(bytes, f)
|
||||
ds = ray.data.read_binary_files(path)
|
||||
assert sorted(extract_values("bytes", ds.take())) == [byte_str]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,833 @@
|
||||
import re
|
||||
from unittest import mock
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from clickhouse_connect.driver.summary import QuerySummary
|
||||
|
||||
from ray.data._internal.datasource.clickhouse_datasink import (
|
||||
ClickHouseDatasink,
|
||||
ClickHouseTableSettings,
|
||||
SinkMode,
|
||||
)
|
||||
from ray.data._internal.datasource.clickhouse_datasource import ClickHouseDatasource
|
||||
from ray.data._internal.execution.interfaces.task_context import TaskContext
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def patch_clickhouse_get_client():
|
||||
with patch("clickhouse_connect.get_client") as mock_factory:
|
||||
mock_instance = MagicMock()
|
||||
mock_instance.insert_arrow.return_value = QuerySummary({"written_rows": 3})
|
||||
mock_factory.return_value = mock_instance
|
||||
yield mock_instance
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_clickhouse_client():
|
||||
client_mock = mock.MagicMock()
|
||||
client_mock.return_value = client_mock
|
||||
return client_mock
|
||||
|
||||
|
||||
class TestClickHouseDatasource:
|
||||
"""Tests for ClickHouseDatasource."""
|
||||
|
||||
@pytest.fixture
|
||||
def datasource(self, mock_clickhouse_client):
|
||||
datasource = ClickHouseDatasource(
|
||||
table="default.table_name",
|
||||
dsn="clickhouse://user:password@localhost:8123/default",
|
||||
columns=["column1", "column2"],
|
||||
order_by=(["column1"], False),
|
||||
client_settings={"setting1": "value1"},
|
||||
client_kwargs={"client_name": "test-client"},
|
||||
)
|
||||
datasource._client = mock_clickhouse_client
|
||||
return datasource
|
||||
|
||||
def test_init(self, datasource):
|
||||
expected_query = (
|
||||
"SELECT column1, column2 FROM default.table_name ORDER BY column1"
|
||||
)
|
||||
assert datasource._query == expected_query
|
||||
|
||||
@mock.patch.object(ClickHouseDatasource, "_init_client")
|
||||
def test_init_with_filter(self, mock_init_client):
|
||||
mock_client = MagicMock()
|
||||
mock_init_client.return_value = mock_client
|
||||
mock_client.query.return_value = MagicMock()
|
||||
ds_with_filter = ClickHouseDatasource(
|
||||
table="default.table_name",
|
||||
dsn="clickhouse://user:password@localhost:8123/default",
|
||||
columns=["column1", "column2"],
|
||||
filter="label = 2 AND text IS NOT NULL",
|
||||
order_by=(["column1"], False),
|
||||
)
|
||||
assert (
|
||||
ds_with_filter._query == "SELECT column1, column2 FROM default.table_name "
|
||||
"WHERE label = 2 AND text IS NOT NULL "
|
||||
"ORDER BY column1"
|
||||
)
|
||||
|
||||
def test_estimate_inmemory_data_size(self, datasource):
|
||||
mock_client = mock.MagicMock()
|
||||
datasource._init_client = MagicMock(return_value=mock_client)
|
||||
mock_client.query.return_value.result_rows = [[12345]]
|
||||
size = datasource.estimate_inmemory_data_size()
|
||||
assert size == 12345
|
||||
mock_client.query.assert_called_once_with(
|
||||
f"SELECT SUM(byteSize(*)) AS estimate FROM ({datasource._query})"
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"limit_row_count, offset_row_count, expected_query",
|
||||
[
|
||||
(
|
||||
10,
|
||||
0,
|
||||
"""
|
||||
SELECT column1, column2 FROM default.table_name ORDER BY column1
|
||||
FETCH FIRST 10 ROWS ONLY
|
||||
""".strip(),
|
||||
),
|
||||
(
|
||||
1,
|
||||
0,
|
||||
"""
|
||||
SELECT column1, column2 FROM default.table_name ORDER BY column1
|
||||
FETCH FIRST 1 ROW ONLY
|
||||
""".strip(),
|
||||
),
|
||||
(
|
||||
10,
|
||||
5,
|
||||
"""
|
||||
SELECT column1, column2 FROM default.table_name ORDER BY column1
|
||||
OFFSET 5 ROWS
|
||||
FETCH NEXT 10 ROWS ONLY
|
||||
""".strip(),
|
||||
),
|
||||
(
|
||||
1,
|
||||
1,
|
||||
"""
|
||||
SELECT column1, column2 FROM default.table_name ORDER BY column1
|
||||
OFFSET 1 ROW
|
||||
FETCH NEXT 1 ROW ONLY
|
||||
""".strip(),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_build_block_query(
|
||||
self, datasource, limit_row_count, offset_row_count, expected_query
|
||||
):
|
||||
generated_query = datasource._build_block_query(
|
||||
limit_row_count, offset_row_count
|
||||
)
|
||||
clean_generated_query = re.sub(r"\s+", " ", generated_query.strip())
|
||||
clean_expected_query = re.sub(r"\s+", " ", expected_query.strip())
|
||||
assert clean_generated_query == clean_expected_query
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"columns, expected_query_part",
|
||||
[
|
||||
(
|
||||
["field1"],
|
||||
"SELECT field1 FROM default.table_name",
|
||||
),
|
||||
(["field1", "field2"], "SELECT field1, field2 FROM default.table_name"),
|
||||
(None, "SELECT * FROM default.table_name"),
|
||||
],
|
||||
)
|
||||
def test_generate_query_columns(self, datasource, columns, expected_query_part):
|
||||
datasource._columns = columns
|
||||
generated_query = datasource._generate_query()
|
||||
assert expected_query_part in generated_query
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"order_by, expected_query_part",
|
||||
[
|
||||
((["field1"], False), "ORDER BY field1"),
|
||||
((["field2"], True), "ORDER BY field2 DESC"),
|
||||
((["field1", "field2"], False), "ORDER BY (field1, field2)"),
|
||||
],
|
||||
)
|
||||
def test_generate_query_with_order_by(
|
||||
self, datasource, order_by, expected_query_part
|
||||
):
|
||||
datasource._order_by = order_by
|
||||
generated_query = datasource._generate_query()
|
||||
assert expected_query_part in generated_query
|
||||
|
||||
@mock.patch.object(ClickHouseDatasource, "_init_client")
|
||||
@pytest.mark.parametrize(
|
||||
"query_params, expected_query",
|
||||
[
|
||||
(
|
||||
{},
|
||||
"SELECT * FROM default.table_name",
|
||||
),
|
||||
(
|
||||
{
|
||||
"columns": ["field1"],
|
||||
},
|
||||
"SELECT field1 FROM default.table_name",
|
||||
),
|
||||
(
|
||||
{
|
||||
"columns": ["field1"],
|
||||
"order_by": (["field1"], False),
|
||||
},
|
||||
"SELECT field1 FROM default.table_name ORDER BY field1",
|
||||
),
|
||||
(
|
||||
{
|
||||
"columns": ["field1", "field2"],
|
||||
"order_by": (["field1"], True),
|
||||
},
|
||||
"SELECT field1, field2 FROM default.table_name ORDER BY field1 DESC",
|
||||
),
|
||||
(
|
||||
{
|
||||
"columns": ["field1", "field2", "field3"],
|
||||
"order_by": (["field1", "field2"], False),
|
||||
},
|
||||
"SELECT field1, field2, field3 FROM default.table_name "
|
||||
"ORDER BY (field1, field2)",
|
||||
),
|
||||
(
|
||||
{
|
||||
"columns": ["field1", "field2", "field3"],
|
||||
"order_by": (["field1", "field2"], True),
|
||||
},
|
||||
"SELECT field1, field2, field3 FROM default.table_name "
|
||||
"ORDER BY (field1, field2) DESC",
|
||||
),
|
||||
(
|
||||
{
|
||||
"columns": ["field1", "field2", "field3"],
|
||||
"order_by": (["field1", "field2", "field3"], True),
|
||||
},
|
||||
"SELECT field1, field2, field3 FROM default.table_name "
|
||||
"ORDER BY (field1, field2, field3) DESC",
|
||||
),
|
||||
(
|
||||
{
|
||||
"columns": None,
|
||||
"filter": "label = 2",
|
||||
},
|
||||
"SELECT * FROM default.table_name WHERE label = 2",
|
||||
),
|
||||
(
|
||||
{
|
||||
"columns": ["field1", "field2"],
|
||||
"filter": "label = 2 AND text IS NOT NULL",
|
||||
"order_by": (["field1"], False),
|
||||
},
|
||||
"SELECT field1, field2 FROM default.table_name WHERE label = 2 AND "
|
||||
"text IS NOT NULL ORDER BY field1",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_generate_query_full(
|
||||
self, mock_init_client, datasource, query_params, expected_query
|
||||
):
|
||||
mock_client = MagicMock()
|
||||
mock_init_client.return_value = mock_client
|
||||
mock_client.query.return_value = MagicMock()
|
||||
datasource._columns = query_params.get("columns")
|
||||
datasource._filter = query_params.get("filter")
|
||||
datasource._order_by = query_params.get("order_by")
|
||||
generated_query = datasource._generate_query()
|
||||
assert expected_query == generated_query
|
||||
|
||||
@pytest.mark.parametrize("parallelism", [1, 2, 3, 4])
|
||||
def test_get_read_tasks_ordered_table(self, datasource, parallelism):
|
||||
batch1 = pa.record_batch([pa.array([1, 2, 3, 4, 5, 6, 7, 8])], names=["field1"])
|
||||
batch2 = pa.record_batch(
|
||||
[pa.array([9, 10, 11, 12, 13, 14, 15, 16])], names=["field1"]
|
||||
)
|
||||
mock_stream = MagicMock()
|
||||
mock_client = mock.MagicMock()
|
||||
mock_client.query_arrow_stream.return_value.__enter__.return_value = mock_stream
|
||||
mock_stream.__iter__.return_value = [batch1, batch2]
|
||||
datasource.MIN_ROWS_PER_READ_TASK = 4
|
||||
datasource._init_client = MagicMock(return_value=mock_client)
|
||||
datasource._get_estimate_count = MagicMock(return_value=16)
|
||||
datasource._get_sampled_estimates = MagicMock(return_value=(100, batch1.schema))
|
||||
read_tasks = datasource.get_read_tasks(parallelism)
|
||||
expected_num_tasks = parallelism
|
||||
assert len(read_tasks) == expected_num_tasks
|
||||
total_rows = sum(batch.num_rows for batch in [batch1, batch2])
|
||||
rows_per_task = total_rows // parallelism
|
||||
extra_rows = total_rows % parallelism
|
||||
for i, read_task in enumerate(read_tasks):
|
||||
expected_rows = rows_per_task + (1 if i < extra_rows else 0)
|
||||
assert read_task.metadata.num_rows == expected_rows
|
||||
|
||||
@pytest.mark.parametrize("parallelism", [1, 4])
|
||||
def test_get_read_tasks_no_ordering(self, datasource, parallelism):
|
||||
datasource._order_by = None
|
||||
batch1 = pa.record_batch([pa.array([1, 2, 3, 4, 5, 6, 7, 8])], names=["field2"])
|
||||
batch2 = pa.record_batch(
|
||||
[pa.array([9, 10, 11, 12, 13, 14, 15, 16])], names=["field2"]
|
||||
)
|
||||
mock_stream = MagicMock()
|
||||
mock_client = mock.MagicMock()
|
||||
mock_client.query_arrow_stream.return_value.__enter__.return_value = mock_stream
|
||||
mock_stream.__iter__.return_value = [batch1, batch2]
|
||||
datasource.MIN_ROWS_PER_READ_TASK = 4
|
||||
datasource._init_client = MagicMock(return_value=mock_client)
|
||||
datasource._get_estimate_count = MagicMock(return_value=16)
|
||||
datasource._get_sampled_estimates = MagicMock(return_value=(100, batch1.schema))
|
||||
read_tasks = datasource.get_read_tasks(parallelism)
|
||||
assert len(read_tasks) == 1
|
||||
for i, read_task in enumerate(read_tasks):
|
||||
assert read_task.metadata.num_rows == 16
|
||||
|
||||
def test_get_read_tasks_no_batches(self, datasource, mock_clickhouse_client):
|
||||
mock_reader = mock.MagicMock()
|
||||
mock_reader.__iter__.return_value = iter([])
|
||||
datasource._init_client = MagicMock(return_value=mock_clickhouse_client)
|
||||
datasource._get_estimate_count = MagicMock(return_value=0)
|
||||
mock_block_accessor = mock.MagicMock()
|
||||
datasource._get_sampled_estimates = MagicMock(return_value=(0, None))
|
||||
datasource._get_sample_block = MagicMock(return_value=mock_block_accessor)
|
||||
read_tasks = datasource.get_read_tasks(parallelism=2)
|
||||
assert len(read_tasks) == 0
|
||||
|
||||
@mock.patch.object(ClickHouseDatasource, "_init_client")
|
||||
@pytest.mark.parametrize(
|
||||
"filter_str, expect_error, expected_error_substring",
|
||||
[
|
||||
("label = 2 AND text IS NOT NULL", False, None),
|
||||
("some_col = 'my;string' AND another_col > 10", False, None),
|
||||
("AND label = 2", True, "Error: Simulated parse error"),
|
||||
("some_col =", True, "Error: Simulated parse error"),
|
||||
("col = 'someval", True, "Error: Simulated parse error"),
|
||||
("col = NULL", True, "Error: Simulated parse error"),
|
||||
(
|
||||
"col = 123; DROP TABLE foobar",
|
||||
True,
|
||||
"Invalid characters outside of string literals",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_filter_validation(
|
||||
self, mock_init_client, filter_str, expect_error, expected_error_substring
|
||||
):
|
||||
mock_client = MagicMock()
|
||||
mock_init_client.return_value = mock_client
|
||||
if expect_error:
|
||||
if "Invalid characters" not in expected_error_substring:
|
||||
mock_client.query.side_effect = Exception("Simulated parse error")
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
ClickHouseDatasource(
|
||||
table="default.table_name",
|
||||
dsn="clickhouse://user:password@localhost:8123/default",
|
||||
filter=filter_str,
|
||||
)
|
||||
assert expected_error_substring in str(exc_info.value), (
|
||||
f"Expected substring '{expected_error_substring}' "
|
||||
f"not found in: {exc_info.value}"
|
||||
)
|
||||
else:
|
||||
mock_client.query.return_value = MagicMock()
|
||||
ds = ClickHouseDatasource(
|
||||
table="default.table_name",
|
||||
dsn="clickhouse://user:password@localhost:8123/default",
|
||||
filter=filter_str,
|
||||
)
|
||||
assert f"WHERE {filter_str}" in ds._query
|
||||
|
||||
@pytest.mark.parametrize("parallelism", [1, 4])
|
||||
def test_get_read_tasks_with_filter(self, datasource, parallelism):
|
||||
datasource._filter = "label = 2 AND text IS NOT NULL"
|
||||
batch1 = pa.record_batch([pa.array([1, 2, 3, 4, 5, 6, 7, 8])], names=["field2"])
|
||||
batch2 = pa.record_batch(
|
||||
[pa.array([9, 10, 11, 12, 13, 14, 15, 16])], names=["field2"]
|
||||
)
|
||||
mock_stream = MagicMock()
|
||||
mock_client = mock.MagicMock()
|
||||
mock_client.query_arrow_stream.return_value.__enter__.return_value = mock_stream
|
||||
mock_stream.__iter__.return_value = [batch1, batch2]
|
||||
datasource.MIN_ROWS_PER_READ_TASK = 4
|
||||
datasource._init_client = MagicMock(return_value=mock_client)
|
||||
datasource._get_estimate_count = MagicMock(return_value=16)
|
||||
datasource._get_sampled_estimates = MagicMock(return_value=(100, batch1.schema))
|
||||
read_tasks = datasource.get_read_tasks(parallelism)
|
||||
assert len(read_tasks) == 1
|
||||
assert read_tasks[0].metadata.num_rows == 16
|
||||
|
||||
def test_filter_none(self):
|
||||
table_name = "default.table_name"
|
||||
dsn = "clickhouse://user:password@localhost:8123/default"
|
||||
with mock.patch.object(ClickHouseDatasource, "_init_client") as mocked_init:
|
||||
mock_client = MagicMock()
|
||||
mocked_init.return_value = mock_client
|
||||
ds = ClickHouseDatasource(table=table_name, dsn=dsn, filter=None)
|
||||
assert "WHERE" not in ds._query
|
||||
assert ds._filter is None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_clickhouse_sink_client():
|
||||
client = MagicMock()
|
||||
client.insert_arrow.return_value = QuerySummary({"written_rows": 3})
|
||||
return client
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def patch_global_get_client(mock_clickhouse_sink_client):
|
||||
with patch(
|
||||
"clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.mark.usefixtures("ray_start_2_cpus_shared")
|
||||
class TestClickHouseDatasink:
|
||||
@pytest.fixture
|
||||
def datasink(self, mock_clickhouse_sink_client):
|
||||
sink = ClickHouseDatasink(
|
||||
table="default.test_table",
|
||||
dsn="clickhouse+http://user:pass@localhost:8123/default",
|
||||
mode=SinkMode.APPEND,
|
||||
table_settings=ClickHouseTableSettings(engine="MergeTree()"),
|
||||
)
|
||||
return sink
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mode",
|
||||
[
|
||||
SinkMode.OVERWRITE,
|
||||
SinkMode.APPEND,
|
||||
SinkMode.CREATE,
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("table_exists", [True, False])
|
||||
def test_on_write_start_modes(
|
||||
self, datasink, mock_clickhouse_sink_client, mode, table_exists
|
||||
):
|
||||
datasink._mode = mode
|
||||
if (mode in [SinkMode.OVERWRITE, SinkMode.CREATE]) or (
|
||||
mode == SinkMode.APPEND and not table_exists
|
||||
):
|
||||
datasink._schema = pa.schema([("col1", pa.int32())])
|
||||
with patch.object(
|
||||
datasink, "_table_exists", return_value=table_exists
|
||||
) as mock_tbl_exists, patch.object(
|
||||
datasink, "_get_existing_order_by", return_value="(prev_col)"
|
||||
) as mock_get_order:
|
||||
if mode == SinkMode.CREATE and table_exists:
|
||||
with pytest.raises(ValueError, match="already exists.*CREATE"):
|
||||
datasink.on_write_start()
|
||||
mock_tbl_exists.assert_called_once()
|
||||
mock_get_order.assert_not_called()
|
||||
mock_clickhouse_sink_client.command.assert_not_called()
|
||||
else:
|
||||
datasink.on_write_start()
|
||||
mock_tbl_exists.assert_called_once()
|
||||
if mode == SinkMode.OVERWRITE:
|
||||
drop_cmd = "DROP TABLE IF EXISTS default.test_table"
|
||||
mock_clickhouse_sink_client.command.assert_any_call(drop_cmd)
|
||||
if table_exists:
|
||||
mock_get_order.assert_called_once()
|
||||
else:
|
||||
mock_get_order.assert_not_called()
|
||||
|
||||
elif mode == SinkMode.APPEND:
|
||||
if table_exists:
|
||||
mock_get_order.assert_called_once()
|
||||
else:
|
||||
mock_get_order.assert_not_called()
|
||||
create_cmds = [
|
||||
call_args[0][0]
|
||||
for call_args in mock_clickhouse_sink_client.command.call_args_list
|
||||
if "CREATE TABLE" in call_args[0][0]
|
||||
]
|
||||
assert (
|
||||
len(create_cmds) == 1
|
||||
), "Expected one CREATE TABLE for append + !exists."
|
||||
elif mode == SinkMode.CREATE:
|
||||
if not table_exists:
|
||||
mock_get_order.assert_not_called()
|
||||
create_cmds = [
|
||||
call_args[0][0]
|
||||
for call_args in mock_clickhouse_sink_client.command.call_args_list
|
||||
if "CREATE TABLE" in call_args[0][0]
|
||||
]
|
||||
assert (
|
||||
len(create_cmds) == 1
|
||||
), "Expected one CREATE TABLE for create + !exists."
|
||||
|
||||
@pytest.mark.parametrize("mode", [SinkMode.OVERWRITE, SinkMode.APPEND])
|
||||
@pytest.mark.parametrize("table_exists", [True, False])
|
||||
@pytest.mark.parametrize("user_order_by", [None, "user_defined_col", "tuple()"])
|
||||
def test_write_behavior(
|
||||
self,
|
||||
datasink,
|
||||
mock_clickhouse_sink_client,
|
||||
mode,
|
||||
table_exists,
|
||||
user_order_by,
|
||||
):
|
||||
datasink._mode = mode
|
||||
if user_order_by is not None:
|
||||
datasink._table_settings.order_by = user_order_by
|
||||
else:
|
||||
datasink._table_settings.order_by = None
|
||||
with patch.object(datasink, "_table_exists", return_value=table_exists), patch(
|
||||
"clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client
|
||||
):
|
||||
if not table_exists or mode == SinkMode.OVERWRITE:
|
||||
datasink._schema = pa.schema([("col1", pa.int32())])
|
||||
datasink.on_write_start()
|
||||
rb = pa.record_batch([pa.array([1, 2, 3])], names=["col1"])
|
||||
block_data = pa.Table.from_batches([rb])
|
||||
ctx = TaskContext(1, "")
|
||||
results = datasink.write([block_data], ctx=ctx)
|
||||
assert results == [3]
|
||||
mock_clickhouse_sink_client.insert_arrow.assert_called()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"schema, expected_order_by",
|
||||
[
|
||||
(pa.schema([]), "tuple()"),
|
||||
(pa.schema([("ts", pa.timestamp("ns")), ("col2", pa.string())]), "ts"),
|
||||
(pa.schema([("col1", pa.string()), ("val", pa.int64())]), "val"),
|
||||
(pa.schema([("s1", pa.string()), ("s2", pa.large_string())]), "s1"),
|
||||
],
|
||||
)
|
||||
def test_pick_best_arrow_field_for_order_by(
|
||||
self, datasink, mock_clickhouse_sink_client, schema, expected_order_by
|
||||
):
|
||||
datasink._mode = SinkMode.OVERWRITE
|
||||
datasink._table_settings.order_by = None
|
||||
datasink._schema = schema
|
||||
with patch.object(datasink, "_table_exists", return_value=False), patch(
|
||||
"clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client
|
||||
):
|
||||
datasink.on_write_start()
|
||||
# Build an empty table: 0 rows
|
||||
empty_table = pa.Table.from_batches([], schema=schema)
|
||||
datasink.write([empty_table], ctx=None)
|
||||
# Since we're skipping empty inserts now, we expect 0 calls:
|
||||
mock_clickhouse_sink_client.insert_arrow.assert_not_called()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ddl_str, expected_order_by",
|
||||
[
|
||||
(
|
||||
"CREATE TABLE default.test_table (col1 Int32) ENGINE = MergeTree() ORDER BY col1",
|
||||
"col1",
|
||||
),
|
||||
("CREATE TABLE default.test_table (col1 Int32) ENGINE = MergeTree()", None),
|
||||
(
|
||||
"CREATE TABLE default.test_table (col1 Int32) ORDER BY city ENGINE = MergeTree()",
|
||||
"city",
|
||||
),
|
||||
(
|
||||
"CREATE TABLE default.test_table (col1 Int32) ENGINE = MergeTree() PARTITION BY toYYYYMMDD(date_col)",
|
||||
None,
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_get_existing_order_by(
|
||||
self, datasink, mock_clickhouse_sink_client, ddl_str, expected_order_by
|
||||
):
|
||||
mock_clickhouse_sink_client.command.return_value = ddl_str
|
||||
result = datasink._get_existing_order_by(mock_clickhouse_sink_client)
|
||||
assert result == expected_order_by
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"table_settings, schema, expected_engine, expected_order_by_part, expected_clauses",
|
||||
[
|
||||
(
|
||||
ClickHouseTableSettings(),
|
||||
pa.schema([("col1", pa.int32())]),
|
||||
"MergeTree()",
|
||||
"ORDER BY col1",
|
||||
[],
|
||||
),
|
||||
(
|
||||
ClickHouseTableSettings(engine="ReplacingMergeTree()"),
|
||||
pa.schema([("col1", pa.int32())]),
|
||||
"ReplacingMergeTree()",
|
||||
"ORDER BY col1",
|
||||
[],
|
||||
),
|
||||
(
|
||||
ClickHouseTableSettings(order_by="user_col"),
|
||||
pa.schema([("col1", pa.int32())]),
|
||||
"MergeTree()",
|
||||
"ORDER BY user_col",
|
||||
[],
|
||||
),
|
||||
(
|
||||
ClickHouseTableSettings(partition_by="toYYYYMMDD(ts)"),
|
||||
pa.schema([("ts", pa.timestamp("ns"))]),
|
||||
"MergeTree()",
|
||||
"ORDER BY ts",
|
||||
["PARTITION BY toYYYYMMDD(ts)"],
|
||||
),
|
||||
(
|
||||
ClickHouseTableSettings(primary_key="id"),
|
||||
pa.schema([("id", pa.int64()), ("val", pa.string())]),
|
||||
"MergeTree()",
|
||||
"ORDER BY id",
|
||||
["PRIMARY KEY (id)"],
|
||||
),
|
||||
(
|
||||
ClickHouseTableSettings(settings="index_granularity=8192"),
|
||||
pa.schema([("id", pa.int64())]),
|
||||
"MergeTree()",
|
||||
"ORDER BY id",
|
||||
["SETTINGS index_granularity=8192"],
|
||||
),
|
||||
(
|
||||
ClickHouseTableSettings(
|
||||
engine="SummingMergeTree()",
|
||||
order_by="col2",
|
||||
partition_by="toYYYYMMDD(ts)",
|
||||
primary_key="id",
|
||||
settings="index_granularity=8192",
|
||||
),
|
||||
pa.schema(
|
||||
[
|
||||
("id", pa.int64()),
|
||||
("col2", pa.float64()),
|
||||
("ts", pa.timestamp("ns")),
|
||||
]
|
||||
),
|
||||
"SummingMergeTree()",
|
||||
"ORDER BY col2",
|
||||
[
|
||||
"PARTITION BY toYYYYMMDD(ts)",
|
||||
"PRIMARY KEY (id)",
|
||||
"SETTINGS index_granularity=8192",
|
||||
],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_generate_create_table_sql(
|
||||
self,
|
||||
datasink,
|
||||
mock_clickhouse_sink_client,
|
||||
table_settings,
|
||||
schema,
|
||||
expected_engine,
|
||||
expected_order_by_part,
|
||||
expected_clauses,
|
||||
):
|
||||
datasink._mode = SinkMode.OVERWRITE
|
||||
datasink._table_settings = table_settings
|
||||
datasink._schema = schema
|
||||
with patch.object(datasink, "_table_exists", return_value=False), patch(
|
||||
"clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client
|
||||
):
|
||||
datasink.on_write_start()
|
||||
arrays = []
|
||||
for field in schema:
|
||||
if pa.types.is_integer(field.type):
|
||||
arrays.append(pa.array([1, 2, 3], type=field.type))
|
||||
elif pa.types.is_floating(field.type):
|
||||
arrays.append(pa.array([1.1, 2.2, 3.3], type=field.type))
|
||||
elif pa.types.is_timestamp(field.type):
|
||||
arrays.append(pa.array([1, 2, 3], type=field.type))
|
||||
else:
|
||||
arrays.append(pa.array(["a", "b", "c"], type=field.type))
|
||||
block_data = pa.Table.from_arrays(arrays, names=[f.name for f in schema])
|
||||
datasink.write([block_data], ctx=TaskContext(1, ""))
|
||||
create_sql = None
|
||||
for call_arg in mock_clickhouse_sink_client.command.call_args_list:
|
||||
sql_arg = call_arg[0][0]
|
||||
if "CREATE TABLE" in sql_arg:
|
||||
create_sql = sql_arg
|
||||
break
|
||||
assert create_sql is not None, "No CREATE TABLE statement was generated!"
|
||||
assert f"ENGINE = {expected_engine}" in create_sql
|
||||
assert expected_order_by_part in create_sql
|
||||
for clause in expected_clauses:
|
||||
assert clause in create_sql
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"provided_schema,block_fields,expected_create_columns",
|
||||
[
|
||||
(
|
||||
pa.schema([("my_col", pa.float64()), ("ts", pa.timestamp("ns"))]),
|
||||
[("my_col", pa.int32()), ("ts", pa.int64())],
|
||||
["`my_col` Float64", "`ts` DateTime64(3)"],
|
||||
),
|
||||
(
|
||||
pa.schema([("my_col", pa.float64()), ("col2", pa.string())]),
|
||||
[("my_col", pa.int64()), ("col2", pa.large_string())],
|
||||
[
|
||||
"`my_col` Float64",
|
||||
"`col2` String",
|
||||
],
|
||||
),
|
||||
(
|
||||
pa.schema([("id", pa.int32()), ("val", pa.string())]),
|
||||
[("id", pa.int64()), ("val", pa.large_string())],
|
||||
[
|
||||
"`id` Int32",
|
||||
"`val` String",
|
||||
],
|
||||
),
|
||||
(
|
||||
pa.schema([("f1", pa.int32()), ("f2", pa.float64())]),
|
||||
[("f1", pa.int32()), ("f2", pa.int32())],
|
||||
[
|
||||
"`f1` Int32",
|
||||
"`f2` Float64",
|
||||
],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_write_schema_override(
|
||||
self,
|
||||
datasink,
|
||||
mock_clickhouse_sink_client,
|
||||
provided_schema,
|
||||
block_fields,
|
||||
expected_create_columns,
|
||||
):
|
||||
datasink._mode = SinkMode.CREATE
|
||||
datasink._table_settings.order_by = None
|
||||
with patch.object(datasink, "_table_exists", return_value=False), patch(
|
||||
"clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client
|
||||
):
|
||||
datasink._schema = provided_schema
|
||||
datasink.on_write_start()
|
||||
arrays = []
|
||||
for name, typ in block_fields:
|
||||
if pa.types.is_integer(typ):
|
||||
arrays.append(pa.array([1, 2, 3], type=typ))
|
||||
elif pa.types.is_string(typ) or pa.types.is_large_string(typ):
|
||||
arrays.append(pa.array(["a", "b", "c"], type=typ))
|
||||
elif pa.types.is_timestamp(typ):
|
||||
arrays.append(pa.array([1, 2, 3], type=typ))
|
||||
else:
|
||||
arrays.append(pa.array([1.0, 2.0, 3.0], type=typ))
|
||||
block_data = pa.Table.from_arrays(
|
||||
arrays, names=[n for (n, _) in block_fields]
|
||||
)
|
||||
datasink.write([block_data], ctx=TaskContext(1, ""))
|
||||
create_sql = None
|
||||
for call_arg in mock_clickhouse_sink_client.command.call_args_list:
|
||||
sql_arg = call_arg[0][0]
|
||||
if "CREATE TABLE" in sql_arg:
|
||||
create_sql = sql_arg
|
||||
break
|
||||
assert create_sql is not None, "Expected CREATE TABLE to be issued."
|
||||
for expected_col_def in expected_create_columns:
|
||||
assert expected_col_def in create_sql
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"max_insert_block_rows,block_sizes,expected_insert_calls",
|
||||
[
|
||||
(2, [6], [3]),
|
||||
(2, [6, 3], [3, 2]),
|
||||
(None, [6, 3], [1, 1]),
|
||||
(3, [3, 5, 2], [1, 2, 1]),
|
||||
],
|
||||
)
|
||||
def test_chunked_inserts(
|
||||
self,
|
||||
datasink,
|
||||
mock_clickhouse_sink_client,
|
||||
max_insert_block_rows,
|
||||
block_sizes,
|
||||
expected_insert_calls,
|
||||
):
|
||||
datasink._mode = SinkMode.CREATE
|
||||
datasink._schema = pa.schema([("col1", pa.int32())])
|
||||
datasink._max_insert_block_rows = max_insert_block_rows
|
||||
with patch.object(datasink, "_table_exists", return_value=False), patch(
|
||||
"clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client
|
||||
):
|
||||
datasink.on_write_start()
|
||||
blocks = []
|
||||
for size in block_sizes:
|
||||
arr = pa.array(range(size), type=pa.int32())
|
||||
block_table = pa.Table.from_arrays([arr], names=["col1"])
|
||||
blocks.append(block_table)
|
||||
datasink.write(blocks, ctx=TaskContext(1, ""))
|
||||
insert_calls = [
|
||||
call_args[0][1]
|
||||
for call_args in mock_clickhouse_sink_client.insert_arrow.call_args_list
|
||||
]
|
||||
actual_inserts = len(insert_calls)
|
||||
assert actual_inserts == sum(expected_insert_calls), (
|
||||
f"Expected total insert calls {sum(expected_insert_calls)}, "
|
||||
f"got {actual_inserts}."
|
||||
)
|
||||
offset = 0
|
||||
for block_idx, size in enumerate(block_sizes):
|
||||
calls_for_block = expected_insert_calls[block_idx]
|
||||
chunk_tables = insert_calls[offset : offset + calls_for_block]
|
||||
offset += calls_for_block
|
||||
total_rows = sum(tbl.num_rows for tbl in chunk_tables)
|
||||
assert total_rows == size, (
|
||||
f"Block of size {size} was split incorrectly. "
|
||||
f"Sum of chunk sizes is {total_rows}."
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"table_exists,mode,user_schema,block_fields,expected_error_regex",
|
||||
[
|
||||
(
|
||||
False,
|
||||
SinkMode.CREATE,
|
||||
pa.schema([("id", pa.int32())]),
|
||||
[("id", pa.int32()), ("extra_col", pa.int32())],
|
||||
r"(ArrowInvalid|Could not convert|field names are not matching|columns not in target schema.*)",
|
||||
),
|
||||
(
|
||||
True,
|
||||
SinkMode.OVERWRITE,
|
||||
pa.schema([("id", pa.timestamp("ns"))]),
|
||||
[("id", pa.int32())],
|
||||
r"(ArrowInvalid|Could not convert|field names are not matching|columns not in target schema|Unsupported cast.*)",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_user_schema_block_mismatch(
|
||||
self,
|
||||
datasink,
|
||||
mock_clickhouse_sink_client,
|
||||
table_exists,
|
||||
mode,
|
||||
user_schema,
|
||||
block_fields,
|
||||
expected_error_regex,
|
||||
):
|
||||
datasink._mode = mode
|
||||
datasink._schema = user_schema
|
||||
with patch.object(datasink, "_table_exists", return_value=table_exists), patch(
|
||||
"clickhouse_connect.get_client", return_value=mock_clickhouse_sink_client
|
||||
):
|
||||
try:
|
||||
datasink.on_write_start()
|
||||
except ValueError:
|
||||
pass
|
||||
arrays = []
|
||||
for name, typ in block_fields:
|
||||
arrays.append(pa.array([1, 2, 3], type=typ))
|
||||
block_data = pa.Table.from_arrays(
|
||||
arrays, names=[n for (n, _) in block_fields]
|
||||
)
|
||||
with pytest.raises(
|
||||
(ValueError, pa.lib.ArrowInvalid, pa.lib.ArrowNotImplementedError),
|
||||
match=expected_error_regex,
|
||||
):
|
||||
datasink.write([block_data], ctx=TaskContext(1, ""))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,213 @@
|
||||
import os
|
||||
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pyarrow.parquet as pq
|
||||
import pytest
|
||||
from packaging.version import Version
|
||||
|
||||
import ray
|
||||
from ray.data import Schema
|
||||
from ray.data._internal.util import rows_same
|
||||
from ray.data.block import BlockAccessor
|
||||
from ray.data.datasource.path_util import _unwrap_protocol
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.data.tests.mock_http_server import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
|
||||
def df_to_csv(dataframe, path, **kwargs):
|
||||
dataframe.to_csv(path, **kwargs)
|
||||
|
||||
|
||||
def test_csv_read(
|
||||
ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
|
||||
):
|
||||
# Single file.
|
||||
df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
||||
path1 = os.path.join(tmp_path, "test1.csv")
|
||||
df1.to_csv(path1, index=False)
|
||||
ds = ray.data.read_csv(path1, partitioning=None)
|
||||
dsdf = ds.to_pandas().sort_values(by=["one", "two"]).reset_index(drop=True)
|
||||
pd.testing.assert_frame_equal(df1.astype(dsdf.dtypes.to_dict()), dsdf)
|
||||
# Test metadata ops.
|
||||
assert ds.count() == 3
|
||||
assert ds.input_files() == [_unwrap_protocol(path1)]
|
||||
assert ds.schema() == Schema(pa.schema([("one", pa.int64()), ("two", pa.string())]))
|
||||
|
||||
# Two files, override_num_blocks=2.
|
||||
df2 = pd.DataFrame({"one": [4, 5, 6], "two": ["e", "f", "g"]})
|
||||
path2 = os.path.join(tmp_path, "test2.csv")
|
||||
df2.to_csv(path2, index=False)
|
||||
ds = ray.data.read_csv([path1, path2], override_num_blocks=2, partitioning=None)
|
||||
dsdf = ds.to_pandas().sort_values(by=["one", "two"]).reset_index(drop=True)
|
||||
df = pd.concat([df1, df2], ignore_index=True)
|
||||
pd.testing.assert_frame_equal(df.astype(dsdf.dtypes.to_dict()), dsdf)
|
||||
# Test metadata ops.
|
||||
for entry in ds._execute().blocks:
|
||||
assert (
|
||||
# pyrefly: ignore[no-matching-overload]
|
||||
BlockAccessor.for_block(ray.get(entry.ref)).size_bytes()
|
||||
== entry.metadata.size_bytes
|
||||
)
|
||||
|
||||
# Three files, override_num_blocks=2.
|
||||
df3 = pd.DataFrame({"one": [7, 8, 9], "two": ["h", "i", "j"]})
|
||||
path3 = os.path.join(tmp_path, "test3.csv")
|
||||
df3.to_csv(path3, index=False)
|
||||
ds = ray.data.read_csv(
|
||||
[path1, path2, path3],
|
||||
override_num_blocks=2,
|
||||
partitioning=None,
|
||||
)
|
||||
df = pd.concat([df1, df2, df3], ignore_index=True)
|
||||
dsdf = ds.to_pandas().sort_values(by=["one", "two"]).reset_index(drop=True)
|
||||
pd.testing.assert_frame_equal(df.astype(dsdf.dtypes.to_dict()), dsdf)
|
||||
|
||||
|
||||
def test_csv_write(
|
||||
ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
|
||||
):
|
||||
input_df = pd.DataFrame({"id": [0]})
|
||||
ds = ray.data.from_blocks([input_df])
|
||||
|
||||
ds.write_csv(tmp_path)
|
||||
|
||||
output_df = pd.concat(
|
||||
[
|
||||
pd.read_csv(os.path.join(tmp_path, filename))
|
||||
for filename in os.listdir(tmp_path)
|
||||
]
|
||||
)
|
||||
assert rows_same(input_df, output_df)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("override_num_blocks", [None, 2])
|
||||
def test_csv_roundtrip(
|
||||
ray_start_regular_shared,
|
||||
tmp_path,
|
||||
override_num_blocks,
|
||||
target_max_block_size_infinite_or_default,
|
||||
):
|
||||
df = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
||||
|
||||
ds = ray.data.from_pandas([df], override_num_blocks=override_num_blocks)
|
||||
ds.write_csv(tmp_path)
|
||||
|
||||
ds2 = ray.data.read_csv(tmp_path)
|
||||
ds2df = ds2.to_pandas()
|
||||
assert rows_same(ds2df, df)
|
||||
for entry in ds2._execute().blocks:
|
||||
# pyrefly: ignore[no-matching-overload]
|
||||
assert (
|
||||
BlockAccessor.for_block(ray.get(entry.ref)).size_bytes()
|
||||
== entry.metadata.size_bytes
|
||||
)
|
||||
|
||||
|
||||
def test_csv_read_invalid_format(ray_start_regular_shared, tmp_path):
|
||||
df = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
||||
|
||||
# Setup: CSV and Parquet files in the same directory.
|
||||
csv_path = os.path.join(tmp_path, "test.csv")
|
||||
df.to_csv(csv_path, index=False)
|
||||
|
||||
table = pa.Table.from_pandas(df)
|
||||
parquet_path = os.path.join(tmp_path, "test.parquet")
|
||||
pq.write_table(table, parquet_path)
|
||||
|
||||
# Test 1: CSV parser should fail on Parquet file.
|
||||
error_message = "Failed to read CSV file"
|
||||
with pytest.raises(ValueError, match=error_message):
|
||||
ray.data.read_csv(parquet_path).materialize()
|
||||
|
||||
# Test 2: CSV parser should fail when directory contains non-CSV files.
|
||||
with pytest.raises(ValueError, match=error_message):
|
||||
ray.data.read_csv(tmp_path).materialize()
|
||||
|
||||
|
||||
def test_csv_read_no_header(ray_start_regular_shared, tmp_path):
|
||||
from pyarrow import csv
|
||||
|
||||
file_path = os.path.join(tmp_path, "test.csv")
|
||||
df = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
||||
df.to_csv(file_path, index=False, header=False)
|
||||
ds = ray.data.read_csv(
|
||||
file_path,
|
||||
read_options=csv.ReadOptions(column_names=["one", "two"]),
|
||||
)
|
||||
out_df = ds.to_pandas()
|
||||
pd.testing.assert_frame_equal(df.astype(out_df.dtypes.to_dict()), out_df)
|
||||
|
||||
|
||||
def test_csv_read_with_column_type_specified(ray_start_regular_shared, tmp_path):
|
||||
from pyarrow import csv
|
||||
|
||||
file_path = os.path.join(tmp_path, "test.csv")
|
||||
df = pd.DataFrame({"one": [1, 2, 3e1], "two": ["a", "b", "c"]})
|
||||
df.to_csv(file_path, index=False)
|
||||
|
||||
# Incorrect to parse scientific notation in int64 as PyArrow represents
|
||||
# it as double.
|
||||
with pytest.raises(ValueError):
|
||||
ray.data.read_csv(
|
||||
file_path,
|
||||
convert_options=csv.ConvertOptions(
|
||||
column_types={"one": "int64", "two": "string"}
|
||||
),
|
||||
).schema()
|
||||
|
||||
# Parsing scientific notation in double should work.
|
||||
ds = ray.data.read_csv(
|
||||
file_path,
|
||||
convert_options=csv.ConvertOptions(
|
||||
column_types={"one": "float64", "two": "string"}
|
||||
),
|
||||
)
|
||||
expected_df = pd.DataFrame({"one": [1.0, 2.0, 30.0], "two": ["a", "b", "c"]})
|
||||
actual_df = ds.to_pandas()
|
||||
pd.testing.assert_frame_equal(
|
||||
expected_df.astype(actual_df.dtypes.to_dict()), actual_df
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
Version(pa.__version__) < Version("7.0.0"),
|
||||
reason="invalid_row_handler was added in pyarrow 7.0.0",
|
||||
)
|
||||
def test_csv_invalid_file_handler(
|
||||
ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
|
||||
):
|
||||
from pyarrow import csv
|
||||
|
||||
invalid_txt = "f1,f2\n2,3\nx\n4,5"
|
||||
invalid_file = os.path.join(tmp_path, "invalid.csv")
|
||||
with open(invalid_file, "wt") as f:
|
||||
f.write(invalid_txt)
|
||||
|
||||
ray.data.read_csv(
|
||||
invalid_file,
|
||||
parse_options=csv.ParseOptions(
|
||||
delimiter=",", invalid_row_handler=lambda i: "skip"
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_read_example_data(ray_start_regular_shared, tmp_path):
|
||||
ds = ray.data.read_csv("example://iris.csv")
|
||||
assert ds.count() == 150
|
||||
assert ds.take(1) == [
|
||||
{
|
||||
"sepal.length": 5.1,
|
||||
"sepal.width": 3.5,
|
||||
"petal.length": 1.4,
|
||||
"petal.width": 0.2,
|
||||
"variety": "Setosa",
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,59 @@
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def ray_start(request):
|
||||
"""Initialize Ray for Daft tests."""
|
||||
import ray
|
||||
|
||||
try:
|
||||
yield ray.init(
|
||||
num_cpus=16,
|
||||
)
|
||||
finally:
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
def test_daft_round_trip(ray_start):
|
||||
import daft
|
||||
|
||||
import ray
|
||||
|
||||
data = {
|
||||
"int_col": list(range(128)),
|
||||
"str_col": [str(i) for i in range(128)],
|
||||
"nested_list_col": [[i] * 3 for i in range(128)],
|
||||
"tensor_col": [np.array([[i] * 3] * 3) for i in range(128)],
|
||||
}
|
||||
df = daft.from_pydict(data)
|
||||
ds = ray.data.from_daft(df)
|
||||
|
||||
# Ray stores data in Arrow format, so to_pandas() returns Arrow-backed
|
||||
# dtypes (e.g. int64[pyarrow]) while Daft may return numpy dtypes.
|
||||
# Compare values only, not dtypes.
|
||||
pd.testing.assert_frame_equal(ds.to_pandas(), df.to_pandas(), check_dtype=False)
|
||||
|
||||
df2 = ds.to_daft()
|
||||
df_pandas = df.to_pandas()
|
||||
df2_pandas = df2.to_pandas()
|
||||
|
||||
for c in data.keys():
|
||||
# NOTE: tensor behavior on round-trip is different because Ray Data provides
|
||||
# Daft with more information about a column being a fixed-shape-tensor.
|
||||
#
|
||||
# Hence the Pandas representation of `df1` is "just" an object column, but
|
||||
# `df2` knows that this is actually a numpy fixed shaped tensor column
|
||||
if c == "tensor_col":
|
||||
original = np.array(list(df_pandas[c]))
|
||||
roundtripped = np.array(list(df2_pandas[c]))
|
||||
np.testing.assert_array_equal(original, roundtripped)
|
||||
else:
|
||||
pd.testing.assert_series_equal(df_pandas[c], df2_pandas[c])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,289 @@
|
||||
"""Unit tests for Databricks credential providers."""
|
||||
|
||||
import os
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
|
||||
from ray.data._internal.datasource.databricks_credentials import (
|
||||
DatabricksCredentialProvider,
|
||||
DatabricksTableCredentialConfig,
|
||||
EnvironmentCredentialProvider,
|
||||
StaticCredentialProvider,
|
||||
UnityCatalogCredentialConfig,
|
||||
resolve_credential_provider,
|
||||
)
|
||||
|
||||
SAMPLE_TOKEN = "dapi_test_token_abc123"
|
||||
SAMPLE_HOST = "https://my-workspace.cloud.databricks.com"
|
||||
SAMPLE_URL = "https://uc-workspace.databricks.com"
|
||||
ALT_TOKEN = "dapi_alt_token_xyz789"
|
||||
ALT_HOST = "https://alt-workspace.databricks.com"
|
||||
|
||||
|
||||
class TestDatabricksCredentialProvider:
|
||||
"""Tests for the abstract DatabricksCredentialProvider base class."""
|
||||
|
||||
def test_cannot_instantiate_abstract_class(self):
|
||||
"""Verify DatabricksCredentialProvider cannot be instantiated directly."""
|
||||
with pytest.raises(TypeError, match="Can't instantiate abstract class"):
|
||||
DatabricksCredentialProvider()
|
||||
|
||||
def test_abstract_methods_defined(self):
|
||||
"""Verify all abstract methods are defined."""
|
||||
abstract_methods = DatabricksCredentialProvider.__abstractmethods__
|
||||
assert "get_token" in abstract_methods
|
||||
assert "get_host" in abstract_methods
|
||||
assert "invalidate" in abstract_methods
|
||||
|
||||
|
||||
class TestStaticCredentialProvider:
|
||||
"""Tests for StaticCredentialProvider."""
|
||||
|
||||
def test_init_with_valid_token_and_host(self):
|
||||
"""Test successful initialization with token and host."""
|
||||
provider = StaticCredentialProvider(token=SAMPLE_TOKEN, host=SAMPLE_HOST)
|
||||
assert provider.get_token() == SAMPLE_TOKEN
|
||||
assert provider.get_host() == SAMPLE_HOST
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"token,host,expected_error",
|
||||
[
|
||||
("", SAMPLE_HOST, "Token cannot be empty"),
|
||||
(None, SAMPLE_HOST, "Token cannot be empty"),
|
||||
(SAMPLE_TOKEN, "", "Host cannot be empty"),
|
||||
(SAMPLE_TOKEN, None, "Host cannot be empty"),
|
||||
],
|
||||
)
|
||||
def test_init_with_invalid_inputs_raises_error(self, token, host, expected_error):
|
||||
"""Test that invalid token or host raises ValueError."""
|
||||
with pytest.raises(ValueError, match=expected_error):
|
||||
StaticCredentialProvider(token=token, host=host)
|
||||
|
||||
def test_invalidate_is_noop(self):
|
||||
"""Test that invalidate doesn't affect the static token."""
|
||||
provider = StaticCredentialProvider(token=SAMPLE_TOKEN, host=SAMPLE_HOST)
|
||||
provider.invalidate()
|
||||
assert provider.get_token() == SAMPLE_TOKEN
|
||||
assert provider.get_host() == SAMPLE_HOST
|
||||
|
||||
def test_get_token_returns_same_value(self):
|
||||
"""Test that get_token always returns the same value."""
|
||||
provider = StaticCredentialProvider(token=SAMPLE_TOKEN, host=SAMPLE_HOST)
|
||||
assert provider.get_token() == SAMPLE_TOKEN
|
||||
assert provider.get_token() == SAMPLE_TOKEN
|
||||
|
||||
|
||||
class TestEnvironmentCredentialProvider:
|
||||
"""Tests for EnvironmentCredentialProvider."""
|
||||
|
||||
def test_get_token_from_env(self):
|
||||
"""Test get_token reads from environment variable."""
|
||||
with mock.patch.dict(
|
||||
os.environ,
|
||||
{"DATABRICKS_TOKEN": SAMPLE_TOKEN, "DATABRICKS_HOST": SAMPLE_HOST},
|
||||
):
|
||||
provider = EnvironmentCredentialProvider()
|
||||
assert provider.get_token() == SAMPLE_TOKEN
|
||||
|
||||
def test_get_host_from_env(self):
|
||||
"""Test get_host reads from environment variable."""
|
||||
with mock.patch.dict(
|
||||
os.environ,
|
||||
{"DATABRICKS_TOKEN": SAMPLE_TOKEN, "DATABRICKS_HOST": SAMPLE_HOST},
|
||||
):
|
||||
provider = EnvironmentCredentialProvider()
|
||||
assert provider.get_host() == SAMPLE_HOST
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"env_vars,expected_error",
|
||||
[
|
||||
({"DATABRICKS_HOST": SAMPLE_HOST}, "DATABRICKS_TOKEN.*not set"),
|
||||
(
|
||||
{"DATABRICKS_TOKEN": SAMPLE_TOKEN},
|
||||
"set environment variable.*DATABRICKS_HOST",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_init_raises_when_env_var_not_set(self, env_vars, expected_error):
|
||||
"""Test __init__ raises ValueError when required env var is not set."""
|
||||
with mock.patch.dict(os.environ, env_vars, clear=True):
|
||||
with pytest.raises(ValueError, match=expected_error):
|
||||
EnvironmentCredentialProvider()
|
||||
|
||||
def test_host_detected_from_databricks_runtime(self):
|
||||
"""Test host is detected from Databricks runtime when env var not set."""
|
||||
detected_host = "detected-host.databricks.com"
|
||||
with (
|
||||
mock.patch.dict(os.environ, {"DATABRICKS_TOKEN": SAMPLE_TOKEN}, clear=True),
|
||||
mock.patch.object(
|
||||
EnvironmentCredentialProvider,
|
||||
"_detect_databricks_host",
|
||||
return_value=detected_host,
|
||||
),
|
||||
):
|
||||
provider = EnvironmentCredentialProvider()
|
||||
assert provider.get_host() == detected_host
|
||||
|
||||
def test_custom_env_var_names(self):
|
||||
"""Test using custom environment variable names."""
|
||||
with mock.patch.dict(
|
||||
os.environ, {"MY_TOKEN": SAMPLE_TOKEN, "MY_HOST": SAMPLE_HOST}
|
||||
):
|
||||
provider = EnvironmentCredentialProvider(
|
||||
token_env_var="MY_TOKEN", host_env_var="MY_HOST"
|
||||
)
|
||||
assert provider.get_token() == SAMPLE_TOKEN
|
||||
assert provider.get_host() == SAMPLE_HOST
|
||||
|
||||
def test_invalidate_refreshes_token_from_env(self):
|
||||
"""Test that invalidate re-reads token from environment."""
|
||||
refreshed_token = "dapi_refreshed_token_456"
|
||||
with mock.patch.dict(
|
||||
os.environ,
|
||||
{"DATABRICKS_TOKEN": SAMPLE_TOKEN, "DATABRICKS_HOST": SAMPLE_HOST},
|
||||
):
|
||||
provider = EnvironmentCredentialProvider()
|
||||
assert provider.get_token() == SAMPLE_TOKEN
|
||||
|
||||
# Simulate external token refresh
|
||||
os.environ["DATABRICKS_TOKEN"] = refreshed_token
|
||||
provider.invalidate()
|
||||
assert provider.get_token() == refreshed_token
|
||||
|
||||
def test_invalidate_keeps_token_if_env_unset(self):
|
||||
"""Test that invalidate keeps existing token if env var is unset."""
|
||||
with mock.patch.dict(
|
||||
os.environ,
|
||||
{"DATABRICKS_TOKEN": SAMPLE_TOKEN, "DATABRICKS_HOST": SAMPLE_HOST},
|
||||
):
|
||||
provider = EnvironmentCredentialProvider()
|
||||
|
||||
# Remove env var after initialization
|
||||
del os.environ["DATABRICKS_TOKEN"]
|
||||
provider.invalidate()
|
||||
# Should keep the old token rather than failing
|
||||
assert provider.get_token() == SAMPLE_TOKEN
|
||||
|
||||
|
||||
class TestDatabricksTableCredentialConfig:
|
||||
"""Tests for DatabricksTableCredentialConfig and resolve_credential_provider."""
|
||||
|
||||
def test_resolve_with_explicit_provider(self):
|
||||
"""Test that explicit credential_provider is returned as-is."""
|
||||
provider = StaticCredentialProvider(token=SAMPLE_TOKEN, host=SAMPLE_HOST)
|
||||
config = DatabricksTableCredentialConfig(credential_provider=provider)
|
||||
result = resolve_credential_provider(config)
|
||||
assert result is provider
|
||||
|
||||
@pytest.mark.parametrize("credential_provider_arg", [None, "no_arg"])
|
||||
def test_resolve_with_none_returns_environment_provider(
|
||||
self, credential_provider_arg
|
||||
):
|
||||
"""Test that EnvironmentCredentialProvider is returned when none provided."""
|
||||
with mock.patch.dict(
|
||||
os.environ,
|
||||
{"DATABRICKS_TOKEN": SAMPLE_TOKEN, "DATABRICKS_HOST": SAMPLE_HOST},
|
||||
):
|
||||
if credential_provider_arg == "no_arg":
|
||||
config = DatabricksTableCredentialConfig()
|
||||
else:
|
||||
config = DatabricksTableCredentialConfig(
|
||||
credential_provider=credential_provider_arg
|
||||
)
|
||||
result = resolve_credential_provider(config)
|
||||
assert isinstance(result, EnvironmentCredentialProvider)
|
||||
|
||||
|
||||
class TestUnityCatalogCredentialConfig:
|
||||
"""Tests for UnityCatalogCredentialConfig and resolve_credential_provider."""
|
||||
|
||||
def test_resolve_with_explicit_provider(self):
|
||||
"""Test that explicit credential_provider is returned as-is."""
|
||||
provider = StaticCredentialProvider(token=SAMPLE_TOKEN, host=SAMPLE_HOST)
|
||||
config = UnityCatalogCredentialConfig(credential_provider=provider)
|
||||
result = resolve_credential_provider(config)
|
||||
assert result is provider
|
||||
|
||||
def test_resolve_with_explicit_provider_ignores_url_and_token(self):
|
||||
"""Test that url/token are ignored when credential_provider is given."""
|
||||
provider = StaticCredentialProvider(token=SAMPLE_TOKEN, host=SAMPLE_HOST)
|
||||
config = UnityCatalogCredentialConfig(
|
||||
credential_provider=provider, url=ALT_HOST, token=ALT_TOKEN
|
||||
)
|
||||
result = resolve_credential_provider(config)
|
||||
assert result is provider
|
||||
|
||||
def test_resolve_with_url_and_token(self):
|
||||
"""Test that url and token create a StaticCredentialProvider."""
|
||||
config = UnityCatalogCredentialConfig(url=SAMPLE_URL, token=SAMPLE_TOKEN)
|
||||
result = resolve_credential_provider(config)
|
||||
assert isinstance(result, StaticCredentialProvider)
|
||||
assert result.get_token() == SAMPLE_TOKEN
|
||||
assert result.get_host() == SAMPLE_URL
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"kwargs",
|
||||
[
|
||||
{},
|
||||
{"url": SAMPLE_URL},
|
||||
{"token": SAMPLE_TOKEN},
|
||||
],
|
||||
ids=["no_args", "only_url", "only_token"],
|
||||
)
|
||||
def test_config_raises_with_incomplete_args(self, kwargs):
|
||||
"""Test that ValueError is raised when args are missing or incomplete."""
|
||||
config = UnityCatalogCredentialConfig(**kwargs)
|
||||
with pytest.raises(ValueError, match="Either 'credential_provider' or both"):
|
||||
resolve_credential_provider(config)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"url,token",
|
||||
[
|
||||
("", SAMPLE_TOKEN),
|
||||
(SAMPLE_URL, ""),
|
||||
],
|
||||
ids=["empty_url", "empty_token"],
|
||||
)
|
||||
def test_resolve_with_empty_string_raises(self, url, token):
|
||||
"""Test that empty strings for url or token raise ValueError."""
|
||||
config = UnityCatalogCredentialConfig(url=url, token=token)
|
||||
with pytest.raises(ValueError):
|
||||
resolve_credential_provider(config)
|
||||
|
||||
|
||||
class TestCredentialProviderSerialization:
|
||||
"""Tests for credential provider serialization (needed for Ray workers)."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"provider_type,expected_token,expected_host",
|
||||
[
|
||||
("static", SAMPLE_TOKEN, SAMPLE_HOST),
|
||||
("environment", SAMPLE_TOKEN, SAMPLE_HOST),
|
||||
],
|
||||
)
|
||||
def test_provider_is_picklable(self, provider_type, expected_token, expected_host):
|
||||
"""Verify credential providers can be pickled and unpickled."""
|
||||
import pickle
|
||||
|
||||
with mock.patch.dict(
|
||||
os.environ,
|
||||
{"DATABRICKS_TOKEN": expected_token, "DATABRICKS_HOST": expected_host},
|
||||
):
|
||||
if provider_type == "static":
|
||||
provider = StaticCredentialProvider(
|
||||
token=expected_token, host=expected_host
|
||||
)
|
||||
else:
|
||||
provider = EnvironmentCredentialProvider()
|
||||
|
||||
pickled = pickle.dumps(provider)
|
||||
unpickled = pickle.loads(pickled)
|
||||
assert unpickled.get_token() == expected_token
|
||||
assert unpickled.get_host() == expected_host
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,700 @@
|
||||
"""Tests for Databricks Unity Catalog datasource."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import tempfile
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from unittest import mock
|
||||
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
import ray.cloudpickle as pickle
|
||||
from ray.data._internal.datasource.databricks_credentials import (
|
||||
DatabricksCredentialProvider,
|
||||
StaticCredentialProvider,
|
||||
)
|
||||
from ray.data._internal.datasource.databricks_uc_datasource import (
|
||||
DatabricksUCDatasource,
|
||||
)
|
||||
from ray.data._internal.util import rows_same
|
||||
from ray.data.tests.datasource.databricks_test_utils import (
|
||||
MockResponse,
|
||||
RefreshableCredentialProvider,
|
||||
)
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
# =============================================================================
|
||||
# Dataclasses for mock objects
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockChunk:
|
||||
"""Mock chunk data for testing."""
|
||||
|
||||
index: int
|
||||
row_count: int
|
||||
byte_count: int
|
||||
data: bytes
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Mock credential providers for testing
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TokenTrackingProvider(DatabricksCredentialProvider):
|
||||
"""A credential provider that returns incrementing tokens to track fetches."""
|
||||
|
||||
def __init__(self):
|
||||
self.token_fetch_count = 0
|
||||
|
||||
def get_token(self) -> str:
|
||||
self.token_fetch_count += 1
|
||||
return f"token_{self.token_fetch_count}"
|
||||
|
||||
def get_host(self) -> str:
|
||||
return "test_host"
|
||||
|
||||
def invalidate(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Pytest fixtures
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def databricks_env():
|
||||
"""Fixture that sets up Databricks environment variables."""
|
||||
with mock.patch.dict(
|
||||
os.environ,
|
||||
{"DATABRICKS_HOST": "test_host", "DATABRICKS_TOKEN": "test_token"},
|
||||
):
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def refreshable_credential_provider():
|
||||
"""Fixture that provides a refreshable credential provider."""
|
||||
return RefreshableCredentialProvider(host="test_host")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def token_tracking_provider():
|
||||
"""Fixture that provides a token tracking credential provider."""
|
||||
return TokenTrackingProvider()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def requests_mocker():
|
||||
"""Fixture that mocks requests.get and requests.post."""
|
||||
with mock.patch("requests.get") as mock_get:
|
||||
with mock.patch("requests.post") as mock_post:
|
||||
yield {"get": mock_get, "post": mock_post}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_data():
|
||||
"""Fixture that provides test DataFrame and configuration."""
|
||||
return {
|
||||
"expected_df": pd.DataFrame(
|
||||
{
|
||||
"c1": range(10000),
|
||||
"c2": [f"str{i}" for i in range(10000)],
|
||||
}
|
||||
),
|
||||
"token": "test_token",
|
||||
"warehouse_id": "test_warehouse_id",
|
||||
"catalog": "catalog1",
|
||||
"schema": "db1",
|
||||
"query": "select * from table1",
|
||||
"rows_per_chunk": 700,
|
||||
}
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Helper functions
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def create_mock_chunks(df: pd.DataFrame, rows_per_chunk: int) -> list[MockChunk]:
|
||||
"""Create mock chunks from a DataFrame."""
|
||||
chunks = []
|
||||
num_rows = len(df)
|
||||
cur_pos = 0
|
||||
index = 0
|
||||
|
||||
while cur_pos < num_rows:
|
||||
chunk_rows = min(rows_per_chunk, num_rows - cur_pos)
|
||||
chunk_df = df[cur_pos : cur_pos + chunk_rows]
|
||||
chunk_pa_table = pa.Table.from_pandas(chunk_df)
|
||||
|
||||
sink = pa.BufferOutputStream()
|
||||
with pa.ipc.new_stream(sink, chunk_pa_table.schema) as writer:
|
||||
writer.write_table(chunk_pa_table)
|
||||
|
||||
chunks.append(
|
||||
MockChunk(
|
||||
index=index,
|
||||
row_count=chunk_rows,
|
||||
byte_count=len(sink.getvalue()),
|
||||
data=sink.getvalue(),
|
||||
)
|
||||
)
|
||||
index += 1
|
||||
cur_pos += rows_per_chunk
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Test classes
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestDatabricksUCDatasourceIntegration:
|
||||
"""Integration tests for DatabricksUCDatasource."""
|
||||
|
||||
_MOCK_ENV_VAR = "RAY_DATABRICKS_UC_DATASOURCE_READ_FN_MOCK_TEST_SETUP_FN_PATH"
|
||||
|
||||
@contextmanager
|
||||
def _setup_mock(self, test_data: dict, mock_chunks: list[MockChunk]):
|
||||
"""Set up mocks for integration tests."""
|
||||
chunk_meta_json = [
|
||||
{
|
||||
"chunk_index": chunk.index,
|
||||
"row_count": chunk.row_count,
|
||||
"byte_count": chunk.byte_count,
|
||||
}
|
||||
for chunk in mock_chunks
|
||||
]
|
||||
chunk_meta_json.reverse()
|
||||
valid_statement_ids = set()
|
||||
|
||||
def request_post_mock(url, data=None, json=None, **kwargs):
|
||||
import json as jsonlib
|
||||
|
||||
headers = kwargs["headers"]
|
||||
|
||||
if url == "https://test_shard/api/2.0/sql/statements/":
|
||||
assert headers == {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {test_data['token']}",
|
||||
}
|
||||
assert jsonlib.loads(data) == {
|
||||
"statement": test_data["query"],
|
||||
"warehouse_id": test_data["warehouse_id"],
|
||||
"wait_timeout": "0s",
|
||||
"disposition": "EXTERNAL_LINKS",
|
||||
"format": "ARROW_STREAM",
|
||||
"catalog": test_data["catalog"],
|
||||
"schema": test_data["schema"],
|
||||
}
|
||||
|
||||
statement_id = uuid.uuid4().hex
|
||||
valid_statement_ids.add(statement_id)
|
||||
|
||||
return MockResponse(
|
||||
status_code=200,
|
||||
content=b"",
|
||||
_json_data={
|
||||
"statement_id": statement_id,
|
||||
"status": {"state": "PENDING"},
|
||||
},
|
||||
)
|
||||
|
||||
assert False, "Invalid request."
|
||||
|
||||
def request_get_mock(url, params=None, **kwargs):
|
||||
headers = kwargs["headers"]
|
||||
|
||||
if match := re.match(
|
||||
r"^https://test_shard/api/2\.0/sql/statements/([^/]*)/$", url
|
||||
):
|
||||
statement_id = match.group(1)
|
||||
assert headers == {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {test_data['token']}",
|
||||
}
|
||||
assert statement_id in valid_statement_ids
|
||||
|
||||
return MockResponse(
|
||||
status_code=200,
|
||||
_json_data={
|
||||
"status": {"state": "SUCCEEDED"},
|
||||
"manifest": {
|
||||
"truncated": False,
|
||||
"chunks": chunk_meta_json,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
if match := re.match(
|
||||
r"^https://test_shard/api/2\.0/sql/"
|
||||
r"statements/([^/]*)/result/chunks/([^/]*)$",
|
||||
url,
|
||||
):
|
||||
assert headers == {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {test_data['token']}",
|
||||
}
|
||||
|
||||
chunk_index = match.group(2)
|
||||
external_link = f"https://test_external_link/{chunk_index}"
|
||||
return MockResponse(
|
||||
status_code=200,
|
||||
_json_data={"external_links": [{"external_link": external_link}]},
|
||||
)
|
||||
|
||||
if match := re.match(r"^https://test_external_link/([^/]*)$", url):
|
||||
assert headers is None
|
||||
chunk_index = int(match.group(1))
|
||||
return MockResponse(
|
||||
status_code=200,
|
||||
content=mock_chunks[chunk_index].data,
|
||||
)
|
||||
|
||||
assert False, "Invalid request."
|
||||
|
||||
with (
|
||||
mock.patch("requests.get", request_get_mock),
|
||||
mock.patch("requests.post", request_post_mock),
|
||||
mock.patch.dict(
|
||||
os.environ,
|
||||
{
|
||||
"DATABRICKS_HOST": "test_shard",
|
||||
"DATABRICKS_TOKEN": test_data["token"],
|
||||
},
|
||||
),
|
||||
):
|
||||
yield
|
||||
|
||||
@contextmanager
|
||||
def _setup_integration_test(self, test_data: dict):
|
||||
"""Set up complete integration test environment with mocks and Ray."""
|
||||
mock_chunks = create_mock_chunks(
|
||||
test_data["expected_df"], test_data["rows_per_chunk"]
|
||||
)
|
||||
|
||||
setup_mock_fn_path = os.path.join(tempfile.mkdtemp(), "setup_mock_fn.pkl")
|
||||
with open(setup_mock_fn_path, "wb") as fp:
|
||||
pickle.dump(lambda: self._setup_mock(test_data, mock_chunks), fp)
|
||||
|
||||
with (
|
||||
self._setup_mock(test_data, mock_chunks),
|
||||
mock.patch.dict(os.environ, {self._MOCK_ENV_VAR: setup_mock_fn_path}),
|
||||
):
|
||||
ray.shutdown()
|
||||
ray.init()
|
||||
yield
|
||||
|
||||
def test_read_with_table_name(self, test_data):
|
||||
"""Test reading data using table name."""
|
||||
with self._setup_integration_test(test_data):
|
||||
result = ray.data.read_databricks_tables(
|
||||
warehouse_id=test_data["warehouse_id"],
|
||||
table="table1",
|
||||
catalog=test_data["catalog"],
|
||||
schema=test_data["schema"],
|
||||
override_num_blocks=5,
|
||||
).to_pandas()
|
||||
|
||||
assert rows_same(result, test_data["expected_df"])
|
||||
|
||||
def test_read_with_sql_query(self, test_data):
|
||||
"""Test reading data using SQL query."""
|
||||
with self._setup_integration_test(test_data):
|
||||
result = ray.data.read_databricks_tables(
|
||||
warehouse_id=test_data["warehouse_id"],
|
||||
query=test_data["query"],
|
||||
catalog=test_data["catalog"],
|
||||
schema=test_data["schema"],
|
||||
override_num_blocks=5,
|
||||
).to_pandas()
|
||||
|
||||
assert rows_same(result, test_data["expected_df"])
|
||||
|
||||
@pytest.mark.parametrize("num_blocks", [5, 100])
|
||||
def test_read_with_different_parallelism(self, test_data, num_blocks):
|
||||
"""Test reading data with different parallelism settings."""
|
||||
with self._setup_integration_test(test_data):
|
||||
result = ray.data.read_databricks_tables(
|
||||
warehouse_id=test_data["warehouse_id"],
|
||||
query=test_data["query"],
|
||||
catalog=test_data["catalog"],
|
||||
schema=test_data["schema"],
|
||||
override_num_blocks=num_blocks,
|
||||
).to_pandas()
|
||||
|
||||
assert rows_same(result, test_data["expected_df"])
|
||||
|
||||
|
||||
class TestDatabricksUCDatasourceCredentials:
|
||||
"""Tests for credential provider handling."""
|
||||
|
||||
def test_schema_name_does_not_shadow_datasource_fields(self, requests_mocker):
|
||||
"""Test that schema name is stored without using the `schema` attribute.
|
||||
This is a regression test for https://github.com/ray-project/ray/issues/46481.
|
||||
"""
|
||||
requests_mocker["post"].return_value = mock.Mock(
|
||||
status_code=200,
|
||||
raise_for_status=lambda: None,
|
||||
json=lambda: {"statement_id": "test_stmt", "status": {"state": "PENDING"}},
|
||||
)
|
||||
requests_mocker["get"].return_value = mock.Mock(
|
||||
status_code=200,
|
||||
raise_for_status=lambda: None,
|
||||
json=lambda: {
|
||||
"status": {"state": "SUCCEEDED"},
|
||||
"manifest": {"truncated": False, "chunks": []},
|
||||
},
|
||||
)
|
||||
|
||||
provider = StaticCredentialProvider(token="my_provider_token", host="test_host")
|
||||
datasource = DatabricksUCDatasource(
|
||||
warehouse_id="test_warehouse",
|
||||
catalog="test_catalog",
|
||||
schema="test_schema",
|
||||
query="SELECT 1",
|
||||
credential_provider=provider,
|
||||
)
|
||||
|
||||
assert datasource.schema_name == "test_schema"
|
||||
assert "schema" not in datasource.__dict__
|
||||
|
||||
call_kwargs = requests_mocker["post"].call_args[1]
|
||||
payload = json.loads(call_kwargs["data"])
|
||||
assert payload["schema"] == "test_schema"
|
||||
|
||||
def test_with_credential_provider(self, requests_mocker):
|
||||
"""Test DatabricksUCDatasource with credential_provider parameter."""
|
||||
requests_mocker["post"].return_value = mock.Mock(
|
||||
status_code=200,
|
||||
raise_for_status=lambda: None,
|
||||
json=lambda: {"statement_id": "test_stmt", "status": {"state": "PENDING"}},
|
||||
)
|
||||
requests_mocker["get"].return_value = mock.Mock(
|
||||
status_code=200,
|
||||
raise_for_status=lambda: None,
|
||||
json=lambda: {
|
||||
"status": {"state": "SUCCEEDED"},
|
||||
"manifest": {"truncated": False},
|
||||
},
|
||||
)
|
||||
|
||||
provider = StaticCredentialProvider(token="my_provider_token", host="test_host")
|
||||
_datasource = DatabricksUCDatasource(
|
||||
warehouse_id="test_warehouse",
|
||||
catalog="test_catalog",
|
||||
schema="test_schema",
|
||||
query="SELECT 1",
|
||||
credential_provider=provider,
|
||||
)
|
||||
|
||||
# Verify the token from provider was used in requests
|
||||
call_kwargs = requests_mocker["post"].call_args[1]
|
||||
assert "Authorization" in call_kwargs["headers"]
|
||||
assert "Bearer my_provider_token" in call_kwargs["headers"]["Authorization"]
|
||||
|
||||
def test_fresh_token_per_request(self, requests_mocker, token_tracking_provider):
|
||||
"""Test that fresh tokens are fetched for each request during polling."""
|
||||
tokens_used = []
|
||||
|
||||
def capture_post(url, *args, **kwargs):
|
||||
tokens_used.append(kwargs["headers"]["Authorization"])
|
||||
return mock.Mock(
|
||||
status_code=200,
|
||||
raise_for_status=lambda: None,
|
||||
json=lambda: {
|
||||
"statement_id": "test_stmt",
|
||||
"status": {"state": "PENDING"},
|
||||
},
|
||||
)
|
||||
|
||||
poll_count = [0]
|
||||
|
||||
def capture_get(url, *args, **kwargs):
|
||||
tokens_used.append(kwargs["headers"]["Authorization"])
|
||||
poll_count[0] += 1
|
||||
state = "PENDING" if poll_count[0] < 3 else "SUCCEEDED"
|
||||
return mock.Mock(
|
||||
status_code=200,
|
||||
raise_for_status=lambda: None,
|
||||
json=lambda: {
|
||||
"status": {"state": state},
|
||||
"manifest": {"truncated": False, "chunks": []},
|
||||
},
|
||||
)
|
||||
|
||||
requests_mocker["post"].side_effect = capture_post
|
||||
requests_mocker["get"].side_effect = capture_get
|
||||
|
||||
DatabricksUCDatasource(
|
||||
warehouse_id="test_warehouse",
|
||||
catalog="test_catalog",
|
||||
schema="test_schema",
|
||||
query="SELECT 1",
|
||||
credential_provider=token_tracking_provider,
|
||||
)
|
||||
|
||||
# Verify fresh token was fetched for each request:
|
||||
# 1 POST (statement creation) + 3 GETs (polling)
|
||||
assert token_tracking_provider.token_fetch_count == 4
|
||||
assert tokens_used == [
|
||||
"Bearer token_1", # POST
|
||||
"Bearer token_2", # GET poll 1
|
||||
"Bearer token_3", # GET poll 2
|
||||
"Bearer token_4", # GET poll 3
|
||||
]
|
||||
|
||||
|
||||
class TestDatabricksUCDatasource401Retry:
|
||||
"""Tests for 401 retry behavior."""
|
||||
|
||||
def test_401_during_initial_post(
|
||||
self, requests_mocker, refreshable_credential_provider
|
||||
):
|
||||
"""Test that 401 during initial POST triggers credential invalidation and retry."""
|
||||
post_call_count = [0]
|
||||
post_headers_captured = []
|
||||
|
||||
def post_side_effect(url, *args, **kwargs):
|
||||
post_call_count[0] += 1
|
||||
headers = kwargs.get("headers", {})
|
||||
post_headers_captured.append(headers.get("Authorization", ""))
|
||||
|
||||
# First POST returns 401
|
||||
if post_call_count[0] == 1:
|
||||
return mock.Mock(status_code=401)
|
||||
# Retry succeeds
|
||||
return mock.Mock(
|
||||
status_code=200,
|
||||
json=lambda: {
|
||||
"statement_id": "test_stmt",
|
||||
"status": {"state": "SUCCEEDED"},
|
||||
"manifest": {"truncated": False, "chunks": []},
|
||||
},
|
||||
)
|
||||
|
||||
requests_mocker["post"].side_effect = post_side_effect
|
||||
|
||||
DatabricksUCDatasource(
|
||||
warehouse_id="test_warehouse",
|
||||
catalog="test_catalog",
|
||||
schema="test_schema",
|
||||
query="SELECT 1",
|
||||
credential_provider=refreshable_credential_provider,
|
||||
)
|
||||
|
||||
# Verify retry occurred
|
||||
assert (
|
||||
post_call_count[0] == 2
|
||||
), "Expected POST to be called twice (initial + retry)"
|
||||
|
||||
# Verify invalidate was called
|
||||
assert refreshable_credential_provider.invalidate_count == 1
|
||||
|
||||
# Verify first request used expired token, retry used refreshed token
|
||||
assert "expired_token" in post_headers_captured[0]
|
||||
assert "refreshed_token" in post_headers_captured[1]
|
||||
|
||||
def test_401_during_polling(self, requests_mocker, refreshable_credential_provider):
|
||||
"""Test that 401 during polling triggers credential invalidation and retry."""
|
||||
poll_call_count = [0]
|
||||
poll_headers_captured = []
|
||||
|
||||
requests_mocker["post"].return_value = mock.Mock(
|
||||
status_code=200,
|
||||
json=lambda: {
|
||||
"statement_id": "test_stmt",
|
||||
"status": {"state": "PENDING"},
|
||||
},
|
||||
)
|
||||
|
||||
def get_side_effect(url, *args, **kwargs):
|
||||
poll_call_count[0] += 1
|
||||
headers = kwargs.get("headers", {})
|
||||
poll_headers_captured.append(headers.get("Authorization", ""))
|
||||
|
||||
# First poll returns 401 with expired token
|
||||
if poll_call_count[0] == 1:
|
||||
return mock.Mock(status_code=401)
|
||||
# Retry succeeds
|
||||
return mock.Mock(
|
||||
status_code=200,
|
||||
json=lambda: {
|
||||
"status": {"state": "SUCCEEDED"},
|
||||
"manifest": {"truncated": False, "chunks": []},
|
||||
},
|
||||
)
|
||||
|
||||
requests_mocker["get"].side_effect = get_side_effect
|
||||
|
||||
DatabricksUCDatasource(
|
||||
warehouse_id="test_warehouse",
|
||||
catalog="test_catalog",
|
||||
schema="test_schema",
|
||||
query="SELECT 1",
|
||||
credential_provider=refreshable_credential_provider,
|
||||
)
|
||||
|
||||
# Verify retry occurred
|
||||
assert (
|
||||
poll_call_count[0] == 2
|
||||
), "Expected GET to be called twice (initial + retry)"
|
||||
|
||||
# Verify invalidate was called once
|
||||
assert refreshable_credential_provider.invalidate_count == 1
|
||||
|
||||
# Verify first request used expired token, retry used refreshed token
|
||||
assert "expired_token" in poll_headers_captured[0]
|
||||
assert "refreshed_token" in poll_headers_captured[1]
|
||||
|
||||
def test_401_during_chunk_fetch(
|
||||
self, requests_mocker, refreshable_credential_provider
|
||||
):
|
||||
"""Test that 401 during chunk fetch triggers credential invalidation and retry."""
|
||||
chunk_fetch_count = [0]
|
||||
chunk_fetch_headers = []
|
||||
|
||||
# Create Arrow data for external URL response
|
||||
table = pa.Table.from_pydict({"col1": [1, 2, 3]})
|
||||
sink = pa.BufferOutputStream()
|
||||
with pa.ipc.new_stream(sink, table.schema) as writer:
|
||||
writer.write_table(table)
|
||||
arrow_data = sink.getvalue().to_pybytes()
|
||||
|
||||
# POST for statement creation succeeds
|
||||
requests_mocker["post"].return_value = mock.Mock(
|
||||
status_code=200,
|
||||
json=lambda: {
|
||||
"statement_id": "test_stmt",
|
||||
"status": {"state": "SUCCEEDED"},
|
||||
"manifest": {
|
||||
"truncated": False,
|
||||
"chunks": [{"chunk_index": 0, "row_count": 10, "byte_count": 100}],
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
def get_side_effect(url, *args, **kwargs):
|
||||
headers = kwargs.get("headers", {})
|
||||
|
||||
# External URL fetch (no auth headers)
|
||||
if url.startswith("https://external/"):
|
||||
return mock.Mock(status_code=200, content=arrow_data)
|
||||
|
||||
if "/result/chunks/" in url:
|
||||
chunk_fetch_count[0] += 1
|
||||
chunk_fetch_headers.append(headers.get("Authorization", ""))
|
||||
|
||||
# First chunk fetch returns 401
|
||||
if chunk_fetch_count[0] == 1:
|
||||
return mock.Mock(status_code=401)
|
||||
# Retry succeeds
|
||||
return mock.Mock(
|
||||
status_code=200,
|
||||
json=lambda: {
|
||||
"external_links": [{"external_link": "https://external/data"}]
|
||||
},
|
||||
)
|
||||
else:
|
||||
# Polling response (already succeeded in POST)
|
||||
return mock.Mock(
|
||||
status_code=200,
|
||||
json=lambda: {
|
||||
"status": {"state": "SUCCEEDED"},
|
||||
"manifest": {
|
||||
"truncated": False,
|
||||
"chunks": [
|
||||
{"chunk_index": 0, "row_count": 10, "byte_count": 100}
|
||||
],
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
requests_mocker["get"].side_effect = get_side_effect
|
||||
|
||||
# Create datasource
|
||||
datasource = DatabricksUCDatasource(
|
||||
warehouse_id="test_warehouse",
|
||||
catalog="test_catalog",
|
||||
schema="test_schema",
|
||||
query="SELECT 1",
|
||||
credential_provider=refreshable_credential_provider,
|
||||
)
|
||||
|
||||
# Get read tasks and execute the read function to trigger chunk fetch
|
||||
read_tasks = datasource.get_read_tasks(parallelism=1)
|
||||
assert len(read_tasks) == 1
|
||||
|
||||
# Execute the read function - this triggers chunk fetch
|
||||
read_fn = read_tasks[0].read_fn
|
||||
results = list(read_fn())
|
||||
|
||||
# Verify chunk fetch retry occurred
|
||||
assert (
|
||||
chunk_fetch_count[0] == 2
|
||||
), "Expected chunk fetch to be called twice (initial + retry)"
|
||||
|
||||
# Verify invalidate was called during chunk fetch
|
||||
assert refreshable_credential_provider.invalidate_count == 1
|
||||
|
||||
# Verify first chunk fetch used expired token, retry used refreshed token
|
||||
assert "expired_token" in chunk_fetch_headers[0]
|
||||
assert "refreshed_token" in chunk_fetch_headers[1]
|
||||
|
||||
# Verify we got results
|
||||
assert len(results) == 1
|
||||
|
||||
|
||||
class TestDatabricksUCDatasourceEmptyResult:
|
||||
"""Tests for empty result handling."""
|
||||
|
||||
def test_empty_result_returns_zero_count(self, requests_mocker, databricks_env):
|
||||
"""Test that empty result returns zero count."""
|
||||
|
||||
def post_mock(url, *args, **kwargs):
|
||||
return MockResponse(
|
||||
status_code=200,
|
||||
_json_data={
|
||||
"statement_id": "test_stmt",
|
||||
"status": {"state": "PENDING"},
|
||||
},
|
||||
)
|
||||
|
||||
def get_mock(url, *args, **kwargs):
|
||||
return MockResponse(
|
||||
status_code=200,
|
||||
_json_data={
|
||||
"status": {"state": "SUCCEEDED"},
|
||||
"manifest": {"truncated": False},
|
||||
},
|
||||
)
|
||||
|
||||
requests_mocker["post"].side_effect = post_mock
|
||||
requests_mocker["get"].side_effect = get_mock
|
||||
|
||||
ds = ray.data.read_databricks_tables(
|
||||
warehouse_id="dummy_warehouse",
|
||||
query="select * from dummy_table",
|
||||
catalog="dummy_catalog",
|
||||
schema="dummy_schema",
|
||||
override_num_blocks=1,
|
||||
)
|
||||
|
||||
assert ds.count() == 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,145 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Iterable, List
|
||||
|
||||
import numpy
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data._internal.execution.interfaces import TaskContext
|
||||
from ray.data.block import Block
|
||||
from ray.data.datasource import Datasink
|
||||
from ray.data.datasource.datasink import DummyOutputDatasink, WriteResult
|
||||
|
||||
|
||||
def test_write_datasink(ray_start_regular_shared):
|
||||
output = DummyOutputDatasink()
|
||||
ds = ray.data.range(10, override_num_blocks=2)
|
||||
ds.write_datasink(output)
|
||||
assert output.num_ok == 1
|
||||
assert output.num_failed == 0
|
||||
assert ray.get(output.data_sink.get_rows_written.remote()) == 10
|
||||
|
||||
output.enabled = False
|
||||
ds = ray.data.range(10, override_num_blocks=2)
|
||||
with pytest.raises(ValueError):
|
||||
ds.write_datasink(output, ray_remote_args={"max_retries": 0})
|
||||
assert output.num_ok == 1
|
||||
assert output.num_failed == 1
|
||||
assert ray.get(output.data_sink.get_rows_written.remote()) == 10
|
||||
|
||||
|
||||
@pytest.mark.parametrize("min_rows_per_write", [25, 50])
|
||||
def test_min_rows_per_write(tmp_path, ray_start_regular_shared, min_rows_per_write):
|
||||
class MockDatasink(Datasink[None]):
|
||||
def __init__(self, min_rows_per_write):
|
||||
self._min_rows_per_write = min_rows_per_write
|
||||
|
||||
def write(self, blocks: Iterable[Block], ctx: TaskContext) -> None:
|
||||
assert sum(len(block) for block in blocks) == self._min_rows_per_write
|
||||
|
||||
@property
|
||||
def min_rows_per_write(self):
|
||||
return self._min_rows_per_write
|
||||
|
||||
ray.data.range(100, override_num_blocks=4).write_datasink(
|
||||
MockDatasink(min_rows_per_write)
|
||||
)
|
||||
|
||||
|
||||
def test_write_result(ray_start_regular_shared):
|
||||
"""Test the write_result argument in `on_write_complete`."""
|
||||
|
||||
@dataclass
|
||||
class CustomWriteResult:
|
||||
|
||||
ids: List[int]
|
||||
|
||||
class CustomDatasink(Datasink[CustomWriteResult]):
|
||||
def __init__(self) -> None:
|
||||
self.ids = []
|
||||
self.num_rows = 0
|
||||
self.size_bytes = 0
|
||||
|
||||
def write(self, blocks: Iterable[Block], ctx: TaskContext):
|
||||
ids = []
|
||||
for b in blocks:
|
||||
ids.extend(b["id"].to_pylist())
|
||||
return CustomWriteResult(ids=ids)
|
||||
|
||||
def on_write_complete(self, write_result: WriteResult[CustomWriteResult]):
|
||||
ids = []
|
||||
for result in write_result.write_returns:
|
||||
ids.extend(result.ids)
|
||||
self.ids = sorted(ids)
|
||||
self.num_rows = write_result.num_rows
|
||||
self.size_bytes = write_result.size_bytes
|
||||
|
||||
num_items = 10
|
||||
size_bytes_per_row = 500
|
||||
|
||||
def map_fn(row):
|
||||
row["data"] = numpy.zeros(size_bytes_per_row, dtype=numpy.int8)
|
||||
return row
|
||||
|
||||
ds = ray.data.range(num_items).map(map_fn)
|
||||
|
||||
datasink = CustomDatasink()
|
||||
ds.write_datasink(datasink)
|
||||
|
||||
assert datasink.ids == list(range(num_items))
|
||||
assert datasink.num_rows == num_items
|
||||
assert datasink.size_bytes == pytest.approx(num_items * size_bytes_per_row, rel=0.1)
|
||||
|
||||
|
||||
class NodeLoggerOutputDatasink(Datasink[None]):
|
||||
"""A writable datasource that logs node IDs of write tasks, for testing."""
|
||||
|
||||
def __init__(self, node_id: str):
|
||||
|
||||
self.num_ok = 0
|
||||
self.num_failed = 0
|
||||
self.node_id = node_id
|
||||
self.num_rows_written = 0
|
||||
|
||||
def write(
|
||||
self,
|
||||
blocks: Iterable[Block],
|
||||
ctx: TaskContext,
|
||||
) -> None:
|
||||
|
||||
node_id = ray.get_runtime_context().get_node_id()
|
||||
assert node_id == self.node_id
|
||||
|
||||
def on_write_complete(self, write_result: WriteResult[None]):
|
||||
self.num_ok += 1
|
||||
self.num_rows_written += write_result.num_rows
|
||||
|
||||
def on_write_failed(self, error: Exception) -> None:
|
||||
self.num_failed += 1
|
||||
|
||||
|
||||
def test_write_datasink_ray_remote_args(ray_start_cluster):
|
||||
ray.shutdown()
|
||||
cluster = ray_start_cluster
|
||||
cluster.add_node(
|
||||
resources={"foo": 100},
|
||||
num_cpus=1,
|
||||
)
|
||||
bar_worker = cluster.add_node(resources={"bar": 100}, num_cpus=1)
|
||||
bar_node_id = bar_worker.node_id
|
||||
|
||||
ray.init(cluster.address)
|
||||
|
||||
output = NodeLoggerOutputDatasink(bar_node_id)
|
||||
ds = ray.data.range(100, override_num_blocks=10)
|
||||
# Pin write tasks to node with "bar" resource.
|
||||
ds.write_datasink(output, ray_remote_args={"resources": {"bar": 1}})
|
||||
assert output.num_ok == 1
|
||||
assert output.num_failed == 0
|
||||
assert output.num_rows_written == 100
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,171 @@
|
||||
import os
|
||||
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from packaging.version import parse as parse_version
|
||||
|
||||
import ray
|
||||
from ray.data import Schema
|
||||
from ray.data._internal.util import rows_same
|
||||
from ray.data._internal.utils.arrow_utils import get_pyarrow_version
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.data.tests.mock_http_server import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
# deltalake's write_deltalake requires pyarrow >= 15 for the Arrow C Stream interface.
|
||||
_pa_version = get_pyarrow_version()
|
||||
assert _pa_version is not None, "pyarrow must be installed to run these tests"
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
_pa_version < parse_version("15.0.0"),
|
||||
reason="deltalake write_deltalake requires pyarrow >= 15.0",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_size",
|
||||
[1, 100],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"write_mode",
|
||||
["append", "overwrite"],
|
||||
)
|
||||
def test_delta_read_basic(tmp_path, batch_size, write_mode):
|
||||
from deltalake import write_deltalake
|
||||
|
||||
# Parse the data path.
|
||||
path = os.path.join(tmp_path, "tmp_test_delta")
|
||||
|
||||
# Create a sample Delta Lake table
|
||||
df = pd.DataFrame(
|
||||
{"x": [42] * batch_size, "y": ["a"] * batch_size, "z": [3.14] * batch_size}
|
||||
)
|
||||
table = pa.Table.from_pandas(df)
|
||||
if write_mode == "append":
|
||||
write_deltalake(path, table, mode=write_mode)
|
||||
write_deltalake(path, table, mode=write_mode)
|
||||
expected = pd.concat([df, df], ignore_index=True)
|
||||
elif write_mode == "overwrite":
|
||||
write_deltalake(path, table, mode=write_mode)
|
||||
expected = df
|
||||
else:
|
||||
raise ValueError(f"Unexpected write_mode: {write_mode}")
|
||||
|
||||
# Read the Delta Lake table
|
||||
ds = ray.data.read_delta(path)
|
||||
|
||||
assert ds.schema() == Schema(
|
||||
pa.schema(
|
||||
{
|
||||
"x": pa.int64(),
|
||||
"y": pa.string(),
|
||||
"z": pa.float64(),
|
||||
}
|
||||
)
|
||||
)
|
||||
assert rows_same(ds.to_pandas(), expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"columns, expected_columns",
|
||||
[
|
||||
(["a", "c"], ["a", "c"]),
|
||||
(["b"], ["b"]),
|
||||
(["a", "b", "c"], ["a", "b", "c"]),
|
||||
],
|
||||
)
|
||||
def test_delta_read_column_selection(tmp_path, columns, expected_columns):
|
||||
from deltalake import write_deltalake
|
||||
|
||||
path = os.path.join(tmp_path, "tmp_test_delta_cols")
|
||||
df = pd.DataFrame({"a": [1, 2, 3], "b": ["x", "y", "z"], "c": [1.0, 2.0, 3.0]})
|
||||
write_deltalake(path, pa.Table.from_pandas(df))
|
||||
|
||||
ds = ray.data.read_delta(path, columns=columns)
|
||||
expected = df[expected_columns]
|
||||
|
||||
assert ds.schema().names == expected_columns
|
||||
assert rows_same(ds.to_pandas(), expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"version, expected_data",
|
||||
[
|
||||
(0, {"x": [1, 2]}),
|
||||
(1, {"x": [3, 4, 5]}),
|
||||
(None, {"x": [3, 4, 5]}),
|
||||
],
|
||||
)
|
||||
def test_delta_read_version(tmp_path, version, expected_data):
|
||||
from deltalake import write_deltalake
|
||||
|
||||
path = os.path.join(tmp_path, "tmp_test_delta_version")
|
||||
write_deltalake(path, pa.table({"x": [1, 2]}))
|
||||
write_deltalake(path, pa.table({"x": [3, 4, 5]}), mode="overwrite")
|
||||
|
||||
ds = ray.data.read_delta(path, version=version)
|
||||
expected = pd.DataFrame(expected_data)
|
||||
|
||||
assert rows_same(ds.to_pandas(), expected)
|
||||
|
||||
|
||||
def test_delta_read_schema_evolution(tmp_path):
|
||||
"""Older files missing newer columns should be null-filled."""
|
||||
from deltalake import write_deltalake
|
||||
|
||||
path = os.path.join(tmp_path, "tmp_test_delta_schema_evo")
|
||||
|
||||
write_deltalake(path, pa.table({"x": [1, 2]}))
|
||||
write_deltalake(
|
||||
path,
|
||||
pa.table({"x": [3, 4], "y": ["a", "b"]}),
|
||||
mode="append",
|
||||
schema_mode="merge", # pyrefly: ignore[unexpected-keyword]
|
||||
)
|
||||
|
||||
ds = ray.data.read_delta(path)
|
||||
expected = pd.DataFrame(
|
||||
{"x": [1, 2, 3, 4], "y": [None, None, "a", "b"]},
|
||||
)
|
||||
# Match the Arrow-backed null sentinel produced by ``to_pandas()``.
|
||||
expected["y"] = expected["y"].astype("string")
|
||||
|
||||
assert rows_same(ds.to_pandas(), expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"storage_options",
|
||||
[{}, None],
|
||||
)
|
||||
def test_delta_read_storage_options(tmp_path, storage_options):
|
||||
"""Verify that storage_options are forwarded to DeltaTable."""
|
||||
from deltalake import write_deltalake
|
||||
|
||||
path = os.path.join(tmp_path, "tmp_test_delta_storage_opts")
|
||||
df = pd.DataFrame({"x": [1, 2, 3]})
|
||||
write_deltalake(path, pa.Table.from_pandas(df))
|
||||
|
||||
ds = ray.data.read_delta(path, storage_options=storage_options)
|
||||
assert rows_same(ds.to_pandas(), df)
|
||||
|
||||
|
||||
def test_delta_read_empty_table(tmp_path):
|
||||
from deltalake import write_deltalake
|
||||
|
||||
path = os.path.join(tmp_path, "tmp_test_delta_empty")
|
||||
write_deltalake(path, pa.table({"x": pa.array([], type=pa.int64())}))
|
||||
|
||||
ds = ray.data.read_delta(path)
|
||||
assert ds.count() == 0
|
||||
|
||||
|
||||
def test_delta_read_rejects_multiple_paths():
|
||||
with pytest.raises(ValueError, match="Only a single Delta Lake table path"):
|
||||
ray.data.read_delta(["path1", "path2"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,294 @@
|
||||
import json
|
||||
import unittest
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
from unittest import mock
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from delta_sharing.protocol import Table
|
||||
from delta_sharing.rest_client import DataSharingRestClient
|
||||
|
||||
from ray.data._internal.datasource.delta_sharing_datasource import (
|
||||
DeltaSharingDatasource,
|
||||
_parse_delta_sharing_url,
|
||||
)
|
||||
from ray.data.block import BlockMetadata
|
||||
from ray.data.dataset import Dataset
|
||||
from ray.data.datasource.datasource import ReadTask
|
||||
from ray.data.read_api import read_delta_sharing_tables
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data.context import DataContext
|
||||
|
||||
|
||||
class TestDeltaSharingDatasource(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.url = "path/to/profile#share.schema.table"
|
||||
self.limit = 1000
|
||||
self.version = 1
|
||||
self.json_predicate_hints = '{"column":"value"}'
|
||||
self.table = Table(name="table", share="share", schema="schema")
|
||||
|
||||
self.mock_rest_client = mock.create_autospec(DataSharingRestClient)
|
||||
self.mock_response = mock.Mock()
|
||||
self.mock_rest_client.list_files_in_table.return_value = self.mock_response
|
||||
|
||||
self.mock_response.add_files = [
|
||||
{"url": "file1", "id": "1"},
|
||||
{"url": "file2", "id": "2"},
|
||||
]
|
||||
self.mock_response.metadata.schema_string = json.dumps(
|
||||
{
|
||||
"type": "struct",
|
||||
"fields": [
|
||||
{
|
||||
"name": "column1",
|
||||
"type": "string",
|
||||
"nullable": True,
|
||||
"metadata": {},
|
||||
}
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
@patch(
|
||||
"ray.data._internal.datasource.delta_sharing_datasource.DeltaSharingDatasource."
|
||||
"setup_delta_sharing_connections"
|
||||
)
|
||||
def test_init(self, mock_setup_delta_sharing_connections):
|
||||
mock_setup_delta_sharing_connections.return_value = (
|
||||
self.table,
|
||||
self.mock_rest_client,
|
||||
)
|
||||
datasource = DeltaSharingDatasource(
|
||||
url=self.url,
|
||||
json_predicate_hints=self.json_predicate_hints,
|
||||
limit=self.limit,
|
||||
version=self.version,
|
||||
timestamp=None,
|
||||
)
|
||||
|
||||
self.assertEqual(datasource._url, self.url)
|
||||
self.assertEqual(datasource._json_predicate_hints, self.json_predicate_hints)
|
||||
self.assertEqual(datasource._limit, self.limit)
|
||||
self.assertEqual(datasource._version, self.version)
|
||||
self.assertEqual(datasource._timestamp, None)
|
||||
|
||||
@patch(
|
||||
"ray.data._internal.datasource.delta_sharing_datasource.DeltaSharingDatasource."
|
||||
"setup_delta_sharing_connections"
|
||||
)
|
||||
def test_get_read_tasks(self, mock_setup_delta_sharing_connections):
|
||||
mock_setup_delta_sharing_connections.return_value = (
|
||||
self.table,
|
||||
self.mock_rest_client,
|
||||
)
|
||||
datasource = DeltaSharingDatasource(
|
||||
url=self.url,
|
||||
json_predicate_hints=self.json_predicate_hints,
|
||||
limit=self.limit,
|
||||
version=self.version,
|
||||
timestamp=None,
|
||||
)
|
||||
|
||||
read_tasks = datasource.get_read_tasks(parallelism=2)
|
||||
self.assertEqual(len(read_tasks), 2)
|
||||
self.assertTrue(all(isinstance(task, ReadTask) for task in read_tasks))
|
||||
|
||||
for task in read_tasks:
|
||||
metadata = task.metadata
|
||||
self.assertIsInstance(metadata, BlockMetadata)
|
||||
self.assertEqual(len(metadata.input_files), 1)
|
||||
self.assertTrue(metadata.input_files[0]["url"] in ["file1", "file2"])
|
||||
self.assertEqual(metadata.num_rows, None)
|
||||
self.assertEqual(metadata.size_bytes, None)
|
||||
self.assertEqual(task.schema, None)
|
||||
self.assertEqual(metadata.exec_stats, None)
|
||||
|
||||
|
||||
class TestParseDeltaSharingUrl(unittest.TestCase):
|
||||
def test_valid_url(self):
|
||||
url = "profile#share.schema.table"
|
||||
expected_result = ("profile", "share", "schema", "table")
|
||||
self.assertEqual(_parse_delta_sharing_url(url), expected_result)
|
||||
|
||||
def test_missing_hash(self):
|
||||
url = "profile-share.schema.table"
|
||||
with self.assertRaises(ValueError) as context:
|
||||
_parse_delta_sharing_url(url)
|
||||
self.assertEqual(str(context.exception), f"Invalid 'url': {url}")
|
||||
|
||||
def test_missing_fragments(self):
|
||||
url = "profile#share.schema"
|
||||
with self.assertRaises(ValueError) as context:
|
||||
_parse_delta_sharing_url(url)
|
||||
self.assertEqual(str(context.exception), f"Invalid 'url': {url}")
|
||||
|
||||
def test_empty_profile(self):
|
||||
url = "#share.schema.table"
|
||||
with self.assertRaises(ValueError) as context:
|
||||
_parse_delta_sharing_url(url)
|
||||
self.assertEqual(str(context.exception), f"Invalid 'url': {url}")
|
||||
|
||||
def test_empty_share(self):
|
||||
url = "profile#.schema.table"
|
||||
with self.assertRaises(ValueError) as context:
|
||||
_parse_delta_sharing_url(url)
|
||||
self.assertEqual(str(context.exception), f"Invalid 'url': {url}")
|
||||
|
||||
def test_empty_schema(self):
|
||||
url = "profile#share..table"
|
||||
with self.assertRaises(ValueError) as context:
|
||||
_parse_delta_sharing_url(url)
|
||||
self.assertEqual(str(context.exception), f"Invalid 'url': {url}")
|
||||
|
||||
def test_empty_table(self):
|
||||
url = "profile#share.schema."
|
||||
with self.assertRaises(ValueError) as context:
|
||||
_parse_delta_sharing_url(url)
|
||||
self.assertEqual(str(context.exception), f"Invalid 'url': {url}")
|
||||
|
||||
|
||||
class MockDeltaSharingDatasource:
|
||||
def __init__(
|
||||
self, url, json_predicate_hints=None, limit=None, version=None, timestamp=None
|
||||
):
|
||||
self._url = url
|
||||
self._json_predicate_hints = json_predicate_hints
|
||||
self._limit = limit
|
||||
self._version = version
|
||||
self._timestamp = timestamp
|
||||
|
||||
def setup_delta_sharing_connections(self, url):
|
||||
# Return mock objects for table and rest_client
|
||||
table = MagicMock()
|
||||
rest_client = MagicMock()
|
||||
|
||||
# Mock the rest_client's list_files_in_table method
|
||||
rest_client.list_files_in_table.return_value = MagicMock(
|
||||
add_files=[
|
||||
{
|
||||
"url": "https://s3-bucket-name.s3.us-west-2.amazonaws.com/delta-exchange-test/table2/date%3D2021-04-28/part-00000-591723a8-6a27-4240-a90e-57426f4736d2.c000.snappy.parquet", # noqa E501
|
||||
"id": "591723a8-6a27-4240-a90e-57426f4736d2",
|
||||
"size": 573,
|
||||
"partitionValues": {"date": "2021-04-28"},
|
||||
"stats": '{"numRecords":1,"minValues":{"eventTime":"2021-04-28T23:33:48.719Z"},"maxValues":{"eventTime":"2021-04-28T23:33:48.719Z"},"nullCount":{"eventTime":0}}', # noqa E501
|
||||
"expirationTimestamp": 1652140800000,
|
||||
}
|
||||
],
|
||||
metadata=MagicMock(
|
||||
schema_string='{"type":"struct","fields":[{"name":"eventTime","type":"timestamp","nullable":true,"metadata":{}},{"name":"date","type":"date","nullable":true,"metadata":{}}]}' # noqa E501
|
||||
),
|
||||
)
|
||||
return table, rest_client
|
||||
|
||||
def get_read_tasks(
|
||||
self,
|
||||
parallelism: int,
|
||||
per_task_row_limit: Optional[int] = None,
|
||||
data_context: Optional["DataContext"] = None,
|
||||
):
|
||||
self._table, self._rest_client = self.setup_delta_sharing_connections(self._url)
|
||||
response = self._rest_client.list_files_in_table(
|
||||
self._table,
|
||||
jsonPredicateHints=self._json_predicate_hints,
|
||||
limitHint=self._limit,
|
||||
version=self._version,
|
||||
timestamp=self._timestamp,
|
||||
)
|
||||
|
||||
read_tasks = []
|
||||
for _ in range(parallelism):
|
||||
read_task = MagicMock()
|
||||
read_task.metadata = MagicMock(
|
||||
num_rows=1,
|
||||
schema=None,
|
||||
input_files=[file["url"] for file in response.add_files],
|
||||
size_bytes=573,
|
||||
exec_stats=None,
|
||||
)
|
||||
read_task.data = MagicMock(
|
||||
return_value=[
|
||||
{
|
||||
"eventTime": "2021-04-28T23:33:48.719Z",
|
||||
"date": "2021-04-28",
|
||||
}
|
||||
]
|
||||
)
|
||||
read_task.per_task_row_limit = per_task_row_limit
|
||||
read_tasks.append(read_task)
|
||||
return read_tasks
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_delta_sharing_datasource(mocker):
|
||||
mock_datasource = mocker.patch(
|
||||
"ray.data._internal.datasource.delta_sharing_datasource.DeltaSharingDatasource",
|
||||
new=MockDeltaSharingDatasource,
|
||||
)
|
||||
return mock_datasource
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_ray_data_read_datasource(mocker):
|
||||
mock_read_datasource = mocker.patch("ray.data.read_datasource")
|
||||
mock_read_datasource.return_value = MagicMock(spec=Dataset)
|
||||
return mock_read_datasource
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def setup_profile_file(tmpdir):
|
||||
profile_content = {
|
||||
"shareCredentialsVersion": 1,
|
||||
"endpoint": "https://sharing.delta.io/delta-sharing/",
|
||||
"bearerToken": "<token>",
|
||||
"expirationTime": "2021-11-12T00:12:29.0Z",
|
||||
}
|
||||
profile_file = tmpdir.join("profile.json")
|
||||
profile_file.write(json.dumps(profile_content))
|
||||
return str(profile_file)
|
||||
|
||||
|
||||
def test_read_delta_sharing_tables(
|
||||
mock_delta_sharing_datasource, mock_ray_data_read_datasource, setup_profile_file
|
||||
):
|
||||
url = f"{setup_profile_file}#share.schema.table"
|
||||
limit = 100
|
||||
version = 1
|
||||
timestamp = "2021-01-01T00:00:00Z"
|
||||
json_predicate_hints = '{"eventTime": "2021-04-28T23:33:48.719Z"}'
|
||||
ray_remote_args = {"num_cpus": 2}
|
||||
concurrency = 4
|
||||
override_num_blocks = 2
|
||||
|
||||
# Call the function under test
|
||||
result = read_delta_sharing_tables(
|
||||
url=url,
|
||||
limit=limit,
|
||||
version=version,
|
||||
timestamp=timestamp,
|
||||
json_predicate_hints=json_predicate_hints,
|
||||
ray_remote_args=ray_remote_args,
|
||||
concurrency=concurrency,
|
||||
override_num_blocks=override_num_blocks,
|
||||
)
|
||||
|
||||
# Assert the result and interactions
|
||||
assert isinstance(result, Dataset)
|
||||
mock_ray_data_read_datasource.assert_called_once()
|
||||
args, kwargs = mock_ray_data_read_datasource.call_args
|
||||
datasource = kwargs["datasource"]
|
||||
assert datasource._url == url
|
||||
assert datasource._json_predicate_hints == json_predicate_hints
|
||||
assert datasource._limit == limit
|
||||
assert datasource._version == version
|
||||
assert datasource._timestamp == timestamp
|
||||
assert kwargs["ray_remote_args"] == ray_remote_args
|
||||
assert kwargs["concurrency"] == concurrency
|
||||
assert kwargs["override_num_blocks"] == override_num_blocks
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,508 @@
|
||||
import os
|
||||
from typing import Any, Dict, Iterator, List
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import pyarrow
|
||||
import pytest
|
||||
from pytest_lazy_fixtures import lf as lazy_fixture
|
||||
|
||||
import ray
|
||||
from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
|
||||
from ray.data.block import Block, BlockAccessor
|
||||
from ray.data.datasource.datasource import ReadTask
|
||||
from ray.data.datasource.file_based_datasource import (
|
||||
FileBasedDatasource,
|
||||
)
|
||||
from ray.data.datasource.partitioning import (
|
||||
Partitioning,
|
||||
PartitionStyle,
|
||||
PathPartitionFilter,
|
||||
)
|
||||
|
||||
|
||||
class MockFileBasedDatasource(FileBasedDatasource):
|
||||
def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
|
||||
builder = DelegatingBlockBuilder()
|
||||
builder.add({"data": f.readall()})
|
||||
yield builder.build()
|
||||
|
||||
|
||||
def execute_read_tasks(tasks: List[ReadTask]) -> List[Dict[str, Any]]:
|
||||
"""Execute the read tasks and return the resulting rows.
|
||||
|
||||
The motivation for this utility function is so that we can test datasources without
|
||||
scheduling Ray tasks.
|
||||
"""
|
||||
builder = DelegatingBlockBuilder()
|
||||
for task in tasks:
|
||||
for block in task():
|
||||
builder.add_block(block)
|
||||
block = builder.build()
|
||||
|
||||
block_accessor = BlockAccessor.for_block(block)
|
||||
rows = list(block_accessor.iter_rows(public_row_format=True))
|
||||
|
||||
return rows
|
||||
|
||||
|
||||
def strip_scheme(uri):
|
||||
"""Remove scheme from a URI, if it exists."""
|
||||
parsed = urlparse(uri)
|
||||
if parsed.scheme:
|
||||
return uri.split("://", 1)[1] # remove scheme
|
||||
return uri # no scheme, return as-is
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"filesystem,dir_path,endpoint_url",
|
||||
[
|
||||
(None, lazy_fixture("local_path"), None),
|
||||
(lazy_fixture("local_fs"), lazy_fixture("local_path"), None),
|
||||
(lazy_fixture("s3_fs"), lazy_fixture("s3_path"), lazy_fixture("s3_server")),
|
||||
(
|
||||
lazy_fixture("s3_fs_with_space"),
|
||||
lazy_fixture("s3_path_with_space"),
|
||||
lazy_fixture("s3_server"),
|
||||
),
|
||||
(
|
||||
lazy_fixture("s3_fs_with_special_chars"),
|
||||
lazy_fixture("s3_path_with_special_chars"),
|
||||
lazy_fixture("s3_server"),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_read_single_file(ray_start_regular_shared, filesystem, dir_path, endpoint_url):
|
||||
# `FileBasedDatasource` should read from the local filesystem if you don't specify
|
||||
# one.
|
||||
write_filesystem = filesystem
|
||||
if write_filesystem is None:
|
||||
write_filesystem = pyarrow.fs.LocalFileSystem()
|
||||
|
||||
file_uri = os.path.join(dir_path, "file.txt")
|
||||
|
||||
# PyArrow filesystems expect paths without schemes. `FileBasedDatasource` handles
|
||||
# this internally, but we need to manually strip the scheme for the test setup.
|
||||
write_path = strip_scheme(file_uri)
|
||||
with write_filesystem.open_output_stream(write_path) as f:
|
||||
f.write(b"spam")
|
||||
|
||||
datasource = MockFileBasedDatasource(file_uri, filesystem=filesystem)
|
||||
tasks = datasource.get_read_tasks(1)
|
||||
rows = execute_read_tasks(tasks)
|
||||
|
||||
assert rows == [{"data": b"spam"}]
|
||||
|
||||
|
||||
def test_read_single_directory(ray_start_regular_shared, tmp_path):
|
||||
dir_path = tmp_path / "dir"
|
||||
dir_path.mkdir()
|
||||
|
||||
p1 = dir_path / "a.txt"
|
||||
p1.write_bytes(b"a")
|
||||
|
||||
p2 = dir_path / "b.txt"
|
||||
p2.write_bytes(b"b")
|
||||
|
||||
datasource = MockFileBasedDatasource(dir_path)
|
||||
rows = execute_read_tasks(datasource.get_read_tasks(1))
|
||||
|
||||
assert sorted(rows, key=lambda r: r["data"]) == [{"data": b"a"}, {"data": b"b"}]
|
||||
|
||||
|
||||
def test_read_dir_and_file_mixed(ray_start_regular_shared, tmp_path):
|
||||
dir_path = tmp_path / "dir"
|
||||
dir_path.mkdir()
|
||||
|
||||
p1 = dir_path / "a.txt"
|
||||
p1.write_bytes(b"a")
|
||||
|
||||
p2 = tmp_path / "c.txt"
|
||||
p2.write_bytes(b"c")
|
||||
|
||||
datasource = MockFileBasedDatasource([str(dir_path), str(p2)])
|
||||
rows = execute_read_tasks(datasource.get_read_tasks(1))
|
||||
|
||||
assert sorted(rows, key=lambda r: r["data"]) == [{"data": b"a"}, {"data": b"c"}]
|
||||
|
||||
|
||||
def test_pathlib_paths(ray_start_regular_shared, tmp_path):
|
||||
"""Test that FileBasedDatasource accepts pathlib.Path objects."""
|
||||
from pathlib import Path
|
||||
|
||||
path = Path(tmp_path) / "test_pathlib"
|
||||
path.mkdir()
|
||||
|
||||
# Create pathlib.Path objects
|
||||
file1 = path / "file1.txt"
|
||||
file2 = path / "file2.txt"
|
||||
|
||||
file1.write_bytes(b"hello")
|
||||
file2.write_bytes(b"world")
|
||||
|
||||
# Verify list of pathlib.Path works
|
||||
datasource = MockFileBasedDatasource([file1, file2])
|
||||
rows = execute_read_tasks(datasource.get_read_tasks(1))
|
||||
assert sorted(rows, key=lambda r: r["data"]) == [
|
||||
{"data": b"hello"},
|
||||
{"data": b"world"},
|
||||
]
|
||||
|
||||
# Verify single pathlib.Path works
|
||||
datasource = MockFileBasedDatasource(file1)
|
||||
rows = execute_read_tasks(datasource.get_read_tasks(1))
|
||||
assert rows == [{"data": b"hello"}]
|
||||
|
||||
|
||||
def test_single_file_infinite_target_max_block_size(
|
||||
ray_start_regular_shared, target_max_block_size_infinite_or_default, tmp_path
|
||||
):
|
||||
path = tmp_path / "file.txt"
|
||||
path.write_bytes(b"spam")
|
||||
|
||||
datasource = MockFileBasedDatasource(path)
|
||||
rows = execute_read_tasks(datasource.get_read_tasks(1))
|
||||
|
||||
assert rows == [{"data": b"spam"}]
|
||||
|
||||
|
||||
def test_partitioning_hive(ray_start_regular_shared, tmp_path):
|
||||
path = os.path.join(tmp_path, "country=us")
|
||||
os.mkdir(path)
|
||||
with open(os.path.join(path, "file.txt"), "wb") as file:
|
||||
file.write(b"")
|
||||
|
||||
datasource = MockFileBasedDatasource(tmp_path, partitioning=Partitioning("hive"))
|
||||
|
||||
tasks = datasource.get_read_tasks(1)
|
||||
rows = execute_read_tasks(tasks)
|
||||
|
||||
assert rows == [{"data": b"", "country": "us"}]
|
||||
|
||||
|
||||
def test_partition_filter_hive(ray_start_regular_shared, tmp_path):
|
||||
for country in ["us", "jp"]:
|
||||
path = os.path.join(tmp_path, f"country={country}")
|
||||
os.mkdir(path)
|
||||
with open(os.path.join(path, "file.txt"), "wb") as file:
|
||||
file.write(b"")
|
||||
|
||||
filter = PathPartitionFilter.of(
|
||||
style=PartitionStyle.HIVE,
|
||||
filter_fn=lambda partitions: partitions["country"] == "us",
|
||||
)
|
||||
datasource = MockFileBasedDatasource(
|
||||
tmp_path, partitioning=Partitioning("hive"), partition_filter=filter
|
||||
)
|
||||
|
||||
tasks = datasource.get_read_tasks(1)
|
||||
rows = execute_read_tasks(tasks)
|
||||
|
||||
assert rows == [{"data": b"", "country": "us"}]
|
||||
|
||||
|
||||
def test_partitioning_dir(ray_start_regular_shared, tmp_path):
|
||||
path = os.path.join(tmp_path, "us")
|
||||
os.mkdir(path)
|
||||
with open(os.path.join(path, "file.txt"), "wb") as file:
|
||||
file.write(b"")
|
||||
|
||||
datasource = MockFileBasedDatasource(
|
||||
tmp_path,
|
||||
partitioning=Partitioning("dir", field_names=["country"], base_dir=tmp_path),
|
||||
)
|
||||
|
||||
tasks = datasource.get_read_tasks(1)
|
||||
rows = execute_read_tasks(tasks)
|
||||
|
||||
assert rows == [{"data": b"", "country": "us"}]
|
||||
|
||||
|
||||
def test_partition_filter_dir(ray_start_regular_shared, tmp_path):
|
||||
for country in ["us", "jp"]:
|
||||
path = os.path.join(tmp_path, country)
|
||||
os.mkdir(path)
|
||||
with open(os.path.join(path, "file.txt"), "wb") as file:
|
||||
file.write(b"")
|
||||
|
||||
filter = PathPartitionFilter.of(
|
||||
style=PartitionStyle.DIRECTORY,
|
||||
base_dir=tmp_path,
|
||||
field_names=["country"],
|
||||
filter_fn=lambda partitions: partitions["country"] == "us",
|
||||
)
|
||||
partitioning = Partitioning("dir", field_names=["country"], base_dir=tmp_path)
|
||||
datasource = MockFileBasedDatasource(
|
||||
tmp_path, partitioning=partitioning, partition_filter=filter
|
||||
)
|
||||
|
||||
tasks = datasource.get_read_tasks(1)
|
||||
rows = execute_read_tasks(tasks)
|
||||
|
||||
assert rows == [{"data": b"", "country": "us"}]
|
||||
|
||||
|
||||
def test_partitioning_raises_on_mismatch(ray_start_regular_shared, tmp_path):
|
||||
"""Test when the partition key already exists in the data."""
|
||||
|
||||
class StubDatasource(FileBasedDatasource):
|
||||
def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
|
||||
builder = DelegatingBlockBuilder()
|
||||
builder.add({"country": f.readall()})
|
||||
yield builder.build()
|
||||
|
||||
path = os.path.join(tmp_path, "country=us")
|
||||
os.mkdir(path)
|
||||
with open(os.path.join(path, "file.txt"), "wb") as file:
|
||||
file.write(b"jp")
|
||||
|
||||
datasource = StubDatasource(tmp_path, partitioning=Partitioning("hive"))
|
||||
|
||||
# The data is `jp`, but the path contains `us`. Since the values are different,
|
||||
# the datasource should raise a ValueError.
|
||||
with pytest.raises(ValueError):
|
||||
tasks = datasource.get_read_tasks(1)
|
||||
execute_read_tasks(tasks)
|
||||
|
||||
|
||||
def test_ignore_missing_paths_true(ray_start_regular_shared, tmp_path):
|
||||
path = os.path.join(tmp_path, "file.txt")
|
||||
with open(path, "wb") as file:
|
||||
file.write(b"")
|
||||
|
||||
datasource = MockFileBasedDatasource(
|
||||
[path, "missing.txt"], ignore_missing_paths=True
|
||||
)
|
||||
|
||||
tasks = datasource.get_read_tasks(1)
|
||||
rows = execute_read_tasks(tasks)
|
||||
|
||||
assert rows == [{"data": b""}]
|
||||
|
||||
|
||||
def test_ignore_missing_paths_false(ray_start_regular_shared, tmp_path):
|
||||
path = os.path.join(tmp_path, "file.txt")
|
||||
with open(path, "wb") as file:
|
||||
file.write(b"")
|
||||
|
||||
with pytest.raises(FileNotFoundError):
|
||||
datasource = MockFileBasedDatasource(
|
||||
[path, "missing.txt"], ignore_missing_paths=False
|
||||
)
|
||||
tasks = datasource.get_read_tasks(1)
|
||||
execute_read_tasks(tasks)
|
||||
|
||||
|
||||
def test_local_paths(ray_start_regular_shared, tmp_path):
|
||||
path = os.path.join(tmp_path, "test.txt")
|
||||
with open(path, "w"):
|
||||
pass
|
||||
|
||||
datasource = MockFileBasedDatasource(path)
|
||||
assert datasource.supports_distributed_reads
|
||||
|
||||
datasource = MockFileBasedDatasource(f"local://{path}")
|
||||
assert not datasource.supports_distributed_reads
|
||||
|
||||
|
||||
def test_local_paths_with_client_raises_error(ray_start_cluster_enabled, tmp_path):
|
||||
ray_start_cluster_enabled.add_node(num_cpus=1)
|
||||
ray_start_cluster_enabled.head_node._ray_params.ray_client_server_port = "10004"
|
||||
ray_start_cluster_enabled.head_node.start_ray_client_server()
|
||||
ray.init("ray://localhost:10004")
|
||||
|
||||
path = os.path.join(tmp_path, "test.txt")
|
||||
with open(path, "w"):
|
||||
pass
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
MockFileBasedDatasource(f"local://{path}")
|
||||
|
||||
|
||||
def test_include_paths(ray_start_regular_shared, tmp_path):
|
||||
path = os.path.join(tmp_path, "test.txt")
|
||||
with open(path, "w"):
|
||||
pass
|
||||
|
||||
datasource = MockFileBasedDatasource(path, include_paths=True)
|
||||
ds = ray.data.read_datasource(datasource)
|
||||
|
||||
paths = [row["path"] for row in ds.take_all()]
|
||||
assert paths == [path]
|
||||
|
||||
|
||||
def test_file_extensions(ray_start_regular_shared, tmp_path):
|
||||
csv_path = os.path.join(tmp_path, "file.csv")
|
||||
with open(csv_path, "w") as file:
|
||||
file.write("spam")
|
||||
|
||||
txt_path = os.path.join(tmp_path, "file.txt")
|
||||
with open(txt_path, "w") as file:
|
||||
file.write("ham")
|
||||
|
||||
datasource = MockFileBasedDatasource([csv_path, txt_path], file_extensions=None)
|
||||
ds = ray.data.read_datasource(datasource)
|
||||
assert sorted(ds.input_files()) == sorted([csv_path, txt_path])
|
||||
|
||||
datasource = MockFileBasedDatasource([csv_path, txt_path], file_extensions=["csv"])
|
||||
ds = ray.data.read_datasource(datasource)
|
||||
assert ds.input_files() == [csv_path]
|
||||
|
||||
|
||||
def test_file_extensions_no_match_raises(ray_start_regular_shared, tmp_path):
|
||||
txt_path = tmp_path / "file.txt"
|
||||
txt_path.write_bytes(b"ham")
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="No input files found to read with the following file extensions",
|
||||
):
|
||||
MockFileBasedDatasource([str(txt_path)], file_extensions=["csv"])
|
||||
|
||||
|
||||
def test_flaky_read_task_retries(ray_start_regular_shared, tmp_path):
|
||||
"""Test that flaky read tasks are retried for both the
|
||||
default set of retried errors and a custom set of retried errors."""
|
||||
csv_path = os.path.join(tmp_path, "file.csv")
|
||||
with open(csv_path, "w") as file:
|
||||
file.write("spam")
|
||||
|
||||
class Counter:
|
||||
def __init__(self):
|
||||
self.value = 0
|
||||
|
||||
def increment(self):
|
||||
self.value += 1
|
||||
return self.value
|
||||
|
||||
default_retried_error = ray.data.context.DEFAULT_RETRIED_IO_ERRORS[0]
|
||||
custom_retried_error = "AWS Error ACCESS_DENIED"
|
||||
|
||||
class FlakyFileBasedDatasource(MockFileBasedDatasource):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
CounterActor = ray.remote(Counter)
|
||||
# This actor ref is shared across all read tasks.
|
||||
self.counter = CounterActor.remote()
|
||||
|
||||
def _read_stream(self, f: "pyarrow.NativeFile", path: str):
|
||||
count = ray.get(self.counter.increment.remote())
|
||||
if count == 1:
|
||||
raise RuntimeError(default_retried_error)
|
||||
elif count == 2:
|
||||
raise RuntimeError(custom_retried_error)
|
||||
else:
|
||||
yield from super()._read_stream(f, path)
|
||||
|
||||
ray.data.DataContext.get_current().retried_io_errors.append(custom_retried_error)
|
||||
|
||||
datasource = FlakyFileBasedDatasource([csv_path])
|
||||
ds = ray.data.read_datasource(datasource)
|
||||
assert len(ds.take()) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"fs",
|
||||
[pyarrow.fs.S3FileSystem(), pyarrow.fs.LocalFileSystem()],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"wrap_with_retries",
|
||||
[True, False],
|
||||
)
|
||||
def test_s3_filesystem_serialization(fs, wrap_with_retries):
|
||||
"""Tests that the S3FileSystem can be serialized and deserialized with
|
||||
the serialization workaround (_S3FileSystemWrapper).
|
||||
|
||||
Also checks that filesystems wrapped with RetryingPyFileSystem are
|
||||
properly unwrapped.
|
||||
"""
|
||||
import ray.cloudpickle as ray_pickle
|
||||
from ray.data._internal.util import RetryingPyFileSystem
|
||||
from ray.data.datasource.file_based_datasource import (
|
||||
_unwrap_s3_serialization_workaround,
|
||||
_wrap_s3_serialization_workaround,
|
||||
)
|
||||
|
||||
orig_fs = fs
|
||||
|
||||
if wrap_with_retries:
|
||||
fs = RetryingPyFileSystem.wrap(fs, retryable_errors=["DUMMY ERROR"])
|
||||
|
||||
wrapped_fs = _wrap_s3_serialization_workaround(fs)
|
||||
unpickled_fs = ray_pickle.loads(ray_pickle.dumps(wrapped_fs))
|
||||
unwrapped_fs = _unwrap_s3_serialization_workaround(unpickled_fs)
|
||||
|
||||
if wrap_with_retries:
|
||||
assert isinstance(unwrapped_fs, RetryingPyFileSystem)
|
||||
assert isinstance(unwrapped_fs.unwrap(), orig_fs.__class__)
|
||||
assert unwrapped_fs.retryable_errors == ["DUMMY ERROR"]
|
||||
else:
|
||||
assert isinstance(unwrapped_fs, orig_fs.__class__)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shuffle", [True, False, "file"])
|
||||
def test_invalid_shuffle_arg_raises_error(ray_start_regular_shared, shuffle):
|
||||
with pytest.raises(ValueError):
|
||||
FileBasedDatasource("example://iris.csv", shuffle=shuffle)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shuffle", [None, "files"])
|
||||
def test_valid_shuffle_arg_does_not_raise_error(ray_start_regular_shared, shuffle):
|
||||
FileBasedDatasource("example://iris.csv", shuffle=shuffle)
|
||||
|
||||
|
||||
def test_shuffle_files_changes_order(ray_start_regular_shared, tmp_path):
|
||||
NUM_FILES = 10
|
||||
NUM_RUNS = 5
|
||||
|
||||
for i in range(NUM_FILES):
|
||||
(tmp_path / f"file_{i:02d}.txt").write_bytes(f"data_{i}".encode())
|
||||
|
||||
datasource = MockFileBasedDatasource(
|
||||
str(tmp_path), shuffle="files", include_paths=True
|
||||
)
|
||||
|
||||
output_paths_list = []
|
||||
# Run NUM_RUNS times to verify shuffle produces different orderings
|
||||
for _ in range(NUM_RUNS):
|
||||
tasks = datasource.get_read_tasks(1)
|
||||
rows = execute_read_tasks(tasks)
|
||||
output_filenames = [os.path.basename(row["path"]) for row in rows]
|
||||
output_paths_list.append(output_filenames)
|
||||
|
||||
expected_order = [f"file_{i:02d}.txt" for i in range(NUM_FILES)]
|
||||
|
||||
# Verify shuffle produces non-deterministic orderings across runs
|
||||
unique_orderings = {tuple(paths) for paths in output_paths_list}
|
||||
assert len(unique_orderings) >= 2
|
||||
|
||||
# Verify all files are present in each run
|
||||
for output_paths in output_paths_list:
|
||||
assert sorted(output_paths) == sorted(expected_order)
|
||||
|
||||
|
||||
def test_read_s3_file_error(shutdown_only, s3_path):
|
||||
from ray.data.datasource.file_meta_provider import _handle_read_os_error
|
||||
|
||||
dummy_path = s3_path + "_dummy"
|
||||
error_message = "Please check that file exists and has properly configured access."
|
||||
with pytest.raises(OSError, match=error_message):
|
||||
ray.data.read_parquet(dummy_path)
|
||||
with pytest.raises(OSError, match=error_message):
|
||||
ray.data.read_binary_files(dummy_path)
|
||||
with pytest.raises(OSError, match=error_message):
|
||||
ray.data.read_csv(dummy_path)
|
||||
with pytest.raises(OSError, match=error_message):
|
||||
ray.data.read_json(dummy_path)
|
||||
with pytest.raises(OSError, match=error_message):
|
||||
error = OSError(
|
||||
f"Error creating dataset. Could not read schema from {dummy_path}: AWS "
|
||||
"Error [code 15]: No response body.. Is this a 'parquet' file?"
|
||||
)
|
||||
_handle_read_os_error(error, dummy_path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,176 @@
|
||||
import os
|
||||
from typing import Any, Dict
|
||||
|
||||
import pyarrow
|
||||
import pytest
|
||||
from pyarrow.fs import LocalFileSystem
|
||||
|
||||
import ray
|
||||
from ray.data.block import BlockAccessor
|
||||
from ray.data.datasource import BlockBasedFileDatasink, RowBasedFileDatasink
|
||||
|
||||
|
||||
class FlakyOutputStream:
|
||||
def __init__(self, stream: pyarrow.NativeFile, num_attempts: int):
|
||||
self._stream = stream
|
||||
self._num_attempts = num_attempts
|
||||
|
||||
def __enter__(self):
|
||||
return self._stream.__enter__()
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
if self._num_attempts < 2:
|
||||
raise RuntimeError("AWS Error NETWORK_CONNECTION")
|
||||
|
||||
self._stream.__exit__(exc_type, exc_value, traceback)
|
||||
|
||||
|
||||
def test_flaky_block_based_open_output_stream(ray_start_regular_shared, tmp_path):
|
||||
class FlakyCSVDatasink(BlockBasedFileDatasink):
|
||||
def __init__(self, path: str):
|
||||
super().__init__(path)
|
||||
self._num_attempts = 0
|
||||
self._filesystem = LocalFileSystem()
|
||||
|
||||
def open_output_stream(self, path: str) -> "pyarrow.NativeFile":
|
||||
stream = self._filesystem.open_output_stream(path)
|
||||
flaky_stream = FlakyOutputStream(stream, self._num_attempts)
|
||||
self._num_attempts += 1
|
||||
return flaky_stream
|
||||
|
||||
def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
|
||||
block.to_pandas().to_csv(file)
|
||||
|
||||
ds = ray.data.range(100)
|
||||
|
||||
ds.write_datasink(FlakyCSVDatasink(tmp_path))
|
||||
|
||||
expected_values = list(range(100))
|
||||
written_values = [row["id"] for row in ray.data.read_csv(tmp_path).take_all()]
|
||||
assert sorted(written_values) == sorted(expected_values)
|
||||
|
||||
|
||||
def test_flaky_row_based_open_output_stream(ray_start_regular_shared, tmp_path):
|
||||
class FlakyTextDatasink(RowBasedFileDatasink):
|
||||
def __init__(self, path: str):
|
||||
super().__init__(path)
|
||||
self._num_attempts = 0
|
||||
self._filesystem = LocalFileSystem()
|
||||
|
||||
def open_output_stream(self, path: str) -> "pyarrow.NativeFile":
|
||||
stream = self._filesystem.open_output_stream(path)
|
||||
flaky_stream = FlakyOutputStream(stream, self._num_attempts)
|
||||
self._num_attempts += 1
|
||||
return flaky_stream
|
||||
|
||||
def write_row_to_file(self, row: Dict[str, Any], file: "pyarrow.NativeFile"):
|
||||
file.write(f"{row['id']}".encode())
|
||||
|
||||
ds = ray.data.range(100)
|
||||
|
||||
ds.write_datasink(FlakyTextDatasink(tmp_path))
|
||||
|
||||
expected_values = [str(i) for i in range(100)]
|
||||
written_values = [row["text"] for row in ray.data.read_text(tmp_path).take_all()]
|
||||
assert sorted(written_values) == sorted(expected_values)
|
||||
|
||||
|
||||
def test_flaky_write_block_to_file(ray_start_regular_shared, tmp_path):
|
||||
class FlakyCSVDatasink(BlockBasedFileDatasink):
|
||||
def __init__(self, path: str):
|
||||
super().__init__(path)
|
||||
self._num_attempts = 0
|
||||
|
||||
def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
|
||||
if self._num_attempts < 2:
|
||||
self._num_attempts += 1
|
||||
raise RuntimeError("AWS Error INTERNAL_FAILURE")
|
||||
|
||||
block.to_pandas().to_csv(file)
|
||||
|
||||
ds = ray.data.range(100)
|
||||
|
||||
ds.write_datasink(FlakyCSVDatasink(tmp_path))
|
||||
|
||||
expected_values = list(range(100))
|
||||
written_values = [row["id"] for row in ray.data.read_csv(tmp_path).take_all()]
|
||||
assert sorted(written_values) == sorted(expected_values)
|
||||
|
||||
|
||||
def test_flaky_write_row_to_file(ray_start_regular_shared, tmp_path):
|
||||
class FlakyTextDatasink(RowBasedFileDatasink):
|
||||
def __init__(self, path: str):
|
||||
super().__init__(path)
|
||||
self._num_attempts = 0
|
||||
|
||||
def write_row_to_file(self, row: Dict[str, Any], file: "pyarrow.NativeFile"):
|
||||
if self._num_attempts < 2:
|
||||
self._num_attempts += 1
|
||||
raise RuntimeError("AWS Error INTERNAL_FAILURE")
|
||||
|
||||
file.write(f"{row['id']}".encode())
|
||||
|
||||
ds = ray.data.range(100)
|
||||
|
||||
ds.write_datasink(FlakyTextDatasink(tmp_path))
|
||||
|
||||
expected_values = [str(i) for i in range(100)]
|
||||
written_values = [row["text"] for row in ray.data.read_text(tmp_path).take_all()]
|
||||
assert sorted(written_values) == sorted(expected_values)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_rows", [0, 1])
|
||||
def test_write_preserves_user_directory(num_rows, tmp_path, ray_start_regular_shared):
|
||||
class MockFileDatasink(BlockBasedFileDatasink):
|
||||
def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
|
||||
file.write(b"")
|
||||
|
||||
ds = ray.data.range(num_rows)
|
||||
path = os.path.join(tmp_path, "test")
|
||||
os.mkdir(path) # User-created directory
|
||||
|
||||
ds.write_datasink(MockFileDatasink(path=path))
|
||||
|
||||
assert os.path.isdir(path)
|
||||
|
||||
|
||||
def test_write_creates_dir(tmp_path, ray_start_regular_shared):
|
||||
class MockFileDatasink(BlockBasedFileDatasink):
|
||||
def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
|
||||
file.write(b"")
|
||||
|
||||
ds = ray.data.range(1)
|
||||
path = os.path.join(tmp_path, "test")
|
||||
|
||||
ds.write_datasink(MockFileDatasink(path=path, try_create_dir=True))
|
||||
|
||||
assert os.path.isdir(path)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("min_rows_per_file", [5, 10, 50])
|
||||
def test_write_min_rows_per_file(tmp_path, ray_start_regular_shared, min_rows_per_file):
|
||||
class MockFileDatasink(BlockBasedFileDatasink):
|
||||
def write_block_to_file(self, block: BlockAccessor, file: "pyarrow.NativeFile"):
|
||||
for _ in range(block.num_rows()):
|
||||
file.write(b"row\n")
|
||||
|
||||
ds = ray.data.range(100, override_num_blocks=20)
|
||||
|
||||
ds.write_datasink(
|
||||
MockFileDatasink(path=tmp_path, min_rows_per_file=min_rows_per_file)
|
||||
)
|
||||
|
||||
num_rows_written_total = 0
|
||||
for filename in os.listdir(tmp_path):
|
||||
with open(os.path.join(tmp_path, filename), "r") as file:
|
||||
num_rows_written = len(file.read().splitlines())
|
||||
assert num_rows_written == min_rows_per_file
|
||||
num_rows_written_total += num_rows_written
|
||||
|
||||
assert num_rows_written_total == 100
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,161 @@
|
||||
import os
|
||||
import zipfile
|
||||
|
||||
import pytest
|
||||
from packaging.version import parse as parse_version
|
||||
from pytest_lazy_fixtures import lf as lazy_fixture
|
||||
|
||||
import ray
|
||||
from ray.data._internal.utils.arrow_utils import get_pyarrow_version
|
||||
from ray.data.datasource.path_util import (
|
||||
_resolve_paths_and_filesystem,
|
||||
_unwrap_protocol,
|
||||
)
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.data.tests.mock_http_server import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
MIN_PYARROW_VERSION_FOR_HUDI = parse_version("11.0.0")
|
||||
PYARROW_VERSION = get_pyarrow_version()
|
||||
PYARROW_VERSION_MEETS_REQUIREMENT = (
|
||||
PYARROW_VERSION and PYARROW_VERSION >= MIN_PYARROW_VERSION_FOR_HUDI
|
||||
)
|
||||
PYARROW_HUDI_TEST_SKIP_REASON = (
|
||||
f"Hudi only supported if pyarrow >= {MIN_PYARROW_VERSION_FOR_HUDI}"
|
||||
)
|
||||
|
||||
|
||||
def _extract_testing_table(fixture_path: str, table_dir: str, target_dir: str) -> str:
|
||||
with zipfile.ZipFile(fixture_path, "r") as zip_ref:
|
||||
zip_ref.extractall(target_dir)
|
||||
return os.path.join(target_dir, table_dir)
|
||||
|
||||
|
||||
def _get_hudi_table_path(fs, data_path, table_name, testing_dir="test_hudi") -> str:
|
||||
setup_data_path = _unwrap_protocol(data_path)
|
||||
target_testing_dir = os.path.join(setup_data_path, testing_dir)
|
||||
fixture_path, _ = _resolve_paths_and_filesystem(
|
||||
f"example://hudi-tables/{table_name}.zip", fs
|
||||
)
|
||||
return _extract_testing_table(fixture_path[0], table_name, target_testing_dir)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not PYARROW_VERSION_MEETS_REQUIREMENT,
|
||||
reason=PYARROW_HUDI_TEST_SKIP_REASON,
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"fs,data_path",
|
||||
[
|
||||
(None, lazy_fixture("local_path")),
|
||||
(lazy_fixture("local_fs"), lazy_fixture("local_path")),
|
||||
],
|
||||
)
|
||||
def test_hudi_snapshot_query_v6_trips_table(ray_start_regular_shared, fs, data_path):
|
||||
table_path = _get_hudi_table_path(fs, data_path, "v6_trips_8i1u")
|
||||
|
||||
ds = ray.data.read_hudi(table_path, filters=[("city", "=", "san_francisco")])
|
||||
|
||||
assert ds.schema().names == [
|
||||
"_hoodie_commit_time",
|
||||
"_hoodie_commit_seqno",
|
||||
"_hoodie_record_key",
|
||||
"_hoodie_partition_path",
|
||||
"_hoodie_file_name",
|
||||
"ts",
|
||||
"uuid",
|
||||
"rider",
|
||||
"driver",
|
||||
"fare",
|
||||
"city",
|
||||
]
|
||||
assert ds.count() == 4
|
||||
rows = (
|
||||
ds.select_columns(["_hoodie_commit_time", "ts", "rider", "fare"])
|
||||
.sort("fare")
|
||||
.take_all()
|
||||
)
|
||||
first_commit = "20250715043008154"
|
||||
second_commit = "20250715043011090"
|
||||
assert rows == [
|
||||
{
|
||||
"_hoodie_commit_time": first_commit,
|
||||
"ts": 1695159649087,
|
||||
"rider": "rider-A",
|
||||
"fare": 19.10,
|
||||
},
|
||||
{
|
||||
"_hoodie_commit_time": second_commit,
|
||||
"ts": 1695046462179,
|
||||
"rider": "rider-D",
|
||||
"fare": 25.0,
|
||||
},
|
||||
{
|
||||
"_hoodie_commit_time": first_commit,
|
||||
"ts": 1695091554788,
|
||||
"rider": "rider-C",
|
||||
"fare": 27.70,
|
||||
},
|
||||
{
|
||||
"_hoodie_commit_time": first_commit,
|
||||
"ts": 1695332066204,
|
||||
"rider": "rider-E",
|
||||
"fare": 93.50,
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not PYARROW_VERSION_MEETS_REQUIREMENT,
|
||||
reason=PYARROW_HUDI_TEST_SKIP_REASON,
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"fs,data_path",
|
||||
[
|
||||
(None, lazy_fixture("local_path")),
|
||||
(lazy_fixture("local_fs"), lazy_fixture("local_path")),
|
||||
],
|
||||
)
|
||||
def test_hudi_incremental_query_v6_trips_table(ray_start_regular_shared, fs, data_path):
|
||||
table_path = _get_hudi_table_path(fs, data_path, "v6_trips_8i1u")
|
||||
|
||||
first_commit = "20250715043008154"
|
||||
second_commit = "20250715043011090"
|
||||
ds = ray.data.read_hudi(
|
||||
table_path,
|
||||
query_type="incremental",
|
||||
hudi_options={
|
||||
"hoodie.read.file_group.start_timestamp": first_commit,
|
||||
"hoodie.read.file_group.end_timestamp": second_commit,
|
||||
},
|
||||
)
|
||||
|
||||
assert ds.schema().names == [
|
||||
"_hoodie_commit_time",
|
||||
"_hoodie_commit_seqno",
|
||||
"_hoodie_record_key",
|
||||
"_hoodie_partition_path",
|
||||
"_hoodie_file_name",
|
||||
"ts",
|
||||
"uuid",
|
||||
"rider",
|
||||
"driver",
|
||||
"fare",
|
||||
"city",
|
||||
]
|
||||
assert ds.count() == 1
|
||||
rows = ds.select_columns(["_hoodie_commit_time", "ts", "rider", "fare"]).take_all()
|
||||
assert rows == [
|
||||
{
|
||||
"_hoodie_commit_time": second_commit,
|
||||
"ts": 1695046462179,
|
||||
"rider": "rider-D",
|
||||
"fare": 25.0,
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,464 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import datasets
|
||||
import pyarrow
|
||||
import pytest
|
||||
import requests
|
||||
from packaging.version import Version
|
||||
|
||||
import ray
|
||||
from ray.data.dataset import Dataset, MaterializedDataset
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_hf_dataset():
|
||||
"""Create a mock HuggingFace dataset for testing."""
|
||||
texts = [
|
||||
"Climate change is a serious threat to our planet",
|
||||
"We need to take action on global warming",
|
||||
"Renewable energy is the future",
|
||||
"Fossil fuels are destroying the environment",
|
||||
"Solar power is becoming more affordable",
|
||||
"Wind energy is growing rapidly",
|
||||
"Electric vehicles are the way forward",
|
||||
"Carbon emissions must be reduced",
|
||||
"Green technology is advancing quickly",
|
||||
"Sustainability is important for future generations",
|
||||
"Climate science is well established",
|
||||
"Ocean levels are rising due to warming",
|
||||
"Extreme weather events are increasing",
|
||||
"Biodiversity loss is accelerating",
|
||||
"Deforestation contributes to climate change",
|
||||
"Clean energy jobs are growing",
|
||||
"Energy efficiency saves money",
|
||||
"Public transportation reduces emissions",
|
||||
"Plant-based diets help the environment",
|
||||
"Recycling is essential for sustainability",
|
||||
]
|
||||
|
||||
# Create labels array with exactly the same length as texts
|
||||
labels = [i % 2 for i in range(len(texts))] # Alternating 0s and 1s
|
||||
|
||||
return datasets.Dataset.from_dict(
|
||||
{
|
||||
"text": texts,
|
||||
"label": labels,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_hf_dataset_dict(mock_hf_dataset):
|
||||
"""Create a mock HuggingFace DatasetDict for testing."""
|
||||
return datasets.DatasetDict({"train": mock_hf_dataset})
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_hf_iterable_dataset():
|
||||
"""Create a mock HuggingFace IterableDataset for testing."""
|
||||
texts = [
|
||||
"Streaming climate tweet 1: The planet is warming",
|
||||
"Streaming climate tweet 2: Renewable energy is key",
|
||||
"Streaming climate tweet 3: We must act now",
|
||||
"Streaming climate tweet 4: Solar panels everywhere",
|
||||
"Streaming climate tweet 5: Wind turbines are beautiful",
|
||||
"Streaming climate tweet 6: Electric cars are the future",
|
||||
"Streaming climate tweet 7: Carbon neutral by 2050",
|
||||
"Streaming climate tweet 8: Green energy revolution",
|
||||
"Streaming climate tweet 9: Climate action needed",
|
||||
"Streaming climate tweet 10: Sustainable development",
|
||||
"Streaming climate tweet 11: Ocean conservation",
|
||||
"Streaming climate tweet 12: Forest protection",
|
||||
"Streaming climate tweet 13: Clean air matters",
|
||||
"Streaming climate tweet 14: Water conservation",
|
||||
"Streaming climate tweet 15: Biodiversity protection",
|
||||
]
|
||||
|
||||
labels = [1, 0, 1, 0, 0, 0, 1, 0, 1, 0, 1, 1, 0, 0, 1]
|
||||
|
||||
dataset = datasets.Dataset.from_dict(
|
||||
{
|
||||
"text": texts,
|
||||
"label": labels,
|
||||
}
|
||||
)
|
||||
iterable_dataset = dataset.to_iterable_dataset()
|
||||
iterable_dataset.expected_count = len(texts)
|
||||
return iterable_dataset
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_parquet_urls():
|
||||
"""Fixture providing mock parquet URLs for testing."""
|
||||
return [
|
||||
"https://huggingface.co/datasets/test/parquet/train-00000-of-00001.parquet",
|
||||
"https://huggingface.co/datasets/test/parquet/train-00001-of-00001.parquet",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_resolved_urls():
|
||||
"""Fixture providing mock resolved URLs (after HTTP redirects) for testing."""
|
||||
return [
|
||||
"https://cdn-lfs.huggingface.co/datasets/test/parquet/train-00000-of-00001.parquet",
|
||||
"https://cdn-lfs.huggingface.co/datasets/test/parquet/train-00001-of-00001.parquet",
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_ray_dataset(mock_hf_dataset):
|
||||
"""Fixture providing a mock Ray dataset that matches the mock HuggingFace dataset."""
|
||||
return ray.data.from_items(
|
||||
[
|
||||
{"text": text, "label": label}
|
||||
for text, label in zip(mock_hf_dataset["text"], mock_hf_dataset["label"])
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_successful_http_responses(mock_parquet_urls):
|
||||
"""Fixture providing mock successful HTTP responses for URL resolution."""
|
||||
mock_responses = []
|
||||
for url in mock_parquet_urls:
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.url = url
|
||||
mock_responses.append(mock_response)
|
||||
return mock_responses
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_redirected_http_responses(mock_parquet_urls, mock_resolved_urls):
|
||||
"""Fixture providing mock HTTP responses that simulate redirects."""
|
||||
mock_responses = []
|
||||
for original_url, resolved_url in zip(mock_parquet_urls, mock_resolved_urls):
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.url = resolved_url
|
||||
mock_responses.append(mock_response)
|
||||
return mock_responses
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_huggingface_datasource():
|
||||
"""Fixture providing the HuggingFaceDatasource class for mocking."""
|
||||
from ray.data._internal.datasource.huggingface_datasource import (
|
||||
HuggingFaceDatasource,
|
||||
)
|
||||
|
||||
return HuggingFaceDatasource
|
||||
|
||||
|
||||
def verify_http_requests(mock_requests_head, expected_urls):
|
||||
"""Verify that HTTP requests were made correctly."""
|
||||
assert mock_requests_head.call_count == len(expected_urls)
|
||||
|
||||
for i, url in enumerate(expected_urls):
|
||||
call_args = mock_requests_head.call_args_list[i]
|
||||
assert call_args[0][0] == url
|
||||
assert call_args[1]["allow_redirects"] is True
|
||||
assert call_args[1]["timeout"] == 5
|
||||
|
||||
|
||||
def verify_read_parquet_call(mock_read_parquet, expected_urls):
|
||||
"""Verify that read_parquet was called with correct parameters."""
|
||||
mock_read_parquet.assert_called_once()
|
||||
call_args = mock_read_parquet.call_args
|
||||
|
||||
# Check that the parquet URLs were passed
|
||||
assert call_args[0][0] == expected_urls
|
||||
|
||||
# Check that the filesystem is HTTPFileSystem
|
||||
assert "filesystem" in call_args[1]
|
||||
assert "HTTPFileSystem" in str(type(call_args[1]["filesystem"]))
|
||||
|
||||
# Check that retry_exceptions includes FileNotFoundError and ClientResponseError
|
||||
assert "ray_remote_args" in call_args[1]
|
||||
assert FileNotFoundError in call_args[1]["ray_remote_args"]["retry_exceptions"]
|
||||
|
||||
|
||||
def verify_dataset_creation(ds, mock_hf_dataset):
|
||||
"""Verify that the dataset was created successfully."""
|
||||
assert isinstance(ds, MaterializedDataset)
|
||||
assert ds.count() == mock_hf_dataset.num_rows
|
||||
|
||||
|
||||
def setup_parquet_mocks(
|
||||
mock_huggingface_datasource,
|
||||
mock_parquet_urls,
|
||||
mock_http_responses,
|
||||
mock_ray_dataset,
|
||||
):
|
||||
"""Setup common mocking pattern for parquet-based tests."""
|
||||
patches = []
|
||||
|
||||
# Mock the list_parquet_urls_from_dataset method
|
||||
datasource_patch = patch.object(
|
||||
mock_huggingface_datasource,
|
||||
"list_parquet_urls_from_dataset",
|
||||
return_value=mock_parquet_urls,
|
||||
)
|
||||
patches.append(datasource_patch)
|
||||
|
||||
# Mock the requests.head calls
|
||||
requests_patch = patch("requests.head")
|
||||
patches.append(requests_patch)
|
||||
|
||||
# Mock the read_parquet function
|
||||
read_parquet_patch = patch("ray.data.read_api.read_parquet")
|
||||
patches.append(read_parquet_patch)
|
||||
|
||||
# Start all patches
|
||||
datasource_mock = datasource_patch.start()
|
||||
requests_mock = requests_patch.start()
|
||||
read_parquet_mock = read_parquet_patch.start()
|
||||
|
||||
# Configure mocks
|
||||
requests_mock.side_effect = mock_http_responses
|
||||
read_parquet_mock.return_value = mock_ray_dataset
|
||||
|
||||
return datasource_mock, requests_mock, read_parquet_mock, patches
|
||||
|
||||
|
||||
def hfds_assert_equals(hfds: datasets.Dataset, ds: Dataset):
|
||||
hfds_table = hfds.data.table
|
||||
ds_table = pyarrow.concat_tables([ray.get(tbl) for tbl in ds.to_arrow_refs()])
|
||||
|
||||
sorting = [(name, "descending") for name in hfds_table.column_names]
|
||||
hfds_table = hfds_table.sort_by(sorting)
|
||||
ds_table = ds_table.sort_by(sorting)
|
||||
|
||||
assert hfds_table.equals(ds_table)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_par", [1, 4])
|
||||
def test_from_huggingface(mock_hf_dataset_dict, ray_start_regular_shared, num_par):
|
||||
# Check that DatasetDict is not directly supported.
|
||||
assert isinstance(mock_hf_dataset_dict, datasets.DatasetDict)
|
||||
with pytest.raises(
|
||||
DeprecationWarning,
|
||||
match="You provided a Hugging Face DatasetDict",
|
||||
):
|
||||
ray.data.from_huggingface(mock_hf_dataset_dict)
|
||||
|
||||
ray_datasets = {
|
||||
"train": ray.data.from_huggingface(
|
||||
mock_hf_dataset_dict["train"], override_num_blocks=num_par
|
||||
),
|
||||
}
|
||||
|
||||
assert isinstance(ray_datasets["train"], ray.data.Dataset)
|
||||
hfds_assert_equals(mock_hf_dataset_dict["train"], ray_datasets["train"])
|
||||
|
||||
# Test reading in a split Hugging Face dataset yields correct individual datasets
|
||||
base_hf_dataset = mock_hf_dataset_dict["train"]
|
||||
hf_dataset_split = base_hf_dataset.train_test_split(test_size=0.2)
|
||||
ray_dataset_split_train = ray.data.from_huggingface(hf_dataset_split["train"])
|
||||
assert ray_dataset_split_train.count() == hf_dataset_split["train"].num_rows
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
datasets.Version(datasets.__version__) < datasets.Version("2.8.0"),
|
||||
reason="IterableDataset.iter() added in 2.8.0",
|
||||
)
|
||||
@pytest.mark.skipif(
|
||||
Version(pyarrow.__version__) < Version("8.0.0"),
|
||||
reason="pyarrow.Table.to_reader() added in 8.0.0",
|
||||
)
|
||||
# Note, pandas is excluded here because IterableDatasets do not support pandas format.
|
||||
@pytest.mark.parametrize(
|
||||
"batch_format",
|
||||
[None, "numpy", "arrow", "torch", "tensorflow", "jax"],
|
||||
)
|
||||
def test_from_huggingface_streaming(
|
||||
mock_hf_iterable_dataset, batch_format, ray_start_regular_shared
|
||||
):
|
||||
hfds = mock_hf_iterable_dataset.with_format(batch_format)
|
||||
assert isinstance(hfds, datasets.IterableDataset)
|
||||
|
||||
ds = ray.data.from_huggingface(hfds)
|
||||
expected_count = mock_hf_iterable_dataset.expected_count
|
||||
assert ds.count() == expected_count
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
datasets.Version(datasets.__version__) < datasets.Version("2.8.0"),
|
||||
reason="IterableDataset.iter() added in 2.8.0",
|
||||
)
|
||||
def test_from_huggingface_dynamic_generated(ray_start_regular_shared):
|
||||
# https://github.com/ray-project/ray/issues/49529
|
||||
# Mock the dynamic dataset loading
|
||||
mock_dataset = datasets.Dataset.from_dict(
|
||||
{
|
||||
"text": [
|
||||
"dynamic tweet 1",
|
||||
"dynamic tweet 2",
|
||||
"dynamic tweet 3",
|
||||
"dynamic tweet 4",
|
||||
"dynamic tweet 5",
|
||||
],
|
||||
"label": [0, 1, 0, 1, 0],
|
||||
}
|
||||
)
|
||||
mock_iterable = mock_dataset.to_iterable_dataset()
|
||||
|
||||
with patch("datasets.load_dataset", return_value=mock_iterable):
|
||||
hfds = datasets.load_dataset(
|
||||
"dataset-org/dream",
|
||||
split="test",
|
||||
streaming=True,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
ds = ray.data.from_huggingface(hfds)
|
||||
ds.take(1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("override_num_blocks", [1, 2, 4, 8])
|
||||
def test_from_huggingface_override_num_blocks(
|
||||
mock_hf_dataset, ray_start_regular_shared, override_num_blocks
|
||||
):
|
||||
"""Test that override_num_blocks works correctly with HuggingFace datasets."""
|
||||
hf_train = mock_hf_dataset
|
||||
|
||||
ds_subset = ray.data.from_huggingface(
|
||||
hf_train, override_num_blocks=override_num_blocks
|
||||
)
|
||||
|
||||
assert isinstance(ds_subset, MaterializedDataset)
|
||||
|
||||
# Verify number of blocks
|
||||
assert ds_subset.num_blocks() == override_num_blocks
|
||||
|
||||
# Verify data integrity
|
||||
assert ds_subset.count() == hf_train.num_rows
|
||||
hfds_assert_equals(hf_train, ds_subset)
|
||||
|
||||
# Test with a smaller subset to test edge cases
|
||||
small_size = max(override_num_blocks * 3, 10)
|
||||
hf_small = hf_train.select(range(min(small_size, hf_train.num_rows)))
|
||||
ds_small = ray.data.from_huggingface(
|
||||
hf_small, override_num_blocks=override_num_blocks
|
||||
)
|
||||
|
||||
# Verify number of blocks
|
||||
assert ds_small.num_blocks() == override_num_blocks
|
||||
|
||||
# Verify data integrity
|
||||
assert ds_small.count() == hf_small.num_rows
|
||||
hfds_assert_equals(hf_small, ds_small)
|
||||
|
||||
|
||||
def test_from_huggingface_with_parquet_files(
|
||||
mock_hf_dataset,
|
||||
ray_start_regular_shared,
|
||||
mock_parquet_urls,
|
||||
mock_ray_dataset,
|
||||
mock_successful_http_responses,
|
||||
mock_huggingface_datasource,
|
||||
):
|
||||
"""Test the distributed read path when parquet file URLs are available."""
|
||||
datasource_mock, requests_mock, read_parquet_mock, patches = setup_parquet_mocks(
|
||||
mock_huggingface_datasource,
|
||||
mock_parquet_urls,
|
||||
mock_successful_http_responses,
|
||||
mock_ray_dataset,
|
||||
)
|
||||
|
||||
try:
|
||||
ds = ray.data.from_huggingface(mock_hf_dataset)
|
||||
|
||||
# Verify HTTP requests
|
||||
verify_http_requests(requests_mock, mock_parquet_urls)
|
||||
|
||||
# Verify read_parquet call
|
||||
verify_read_parquet_call(read_parquet_mock, mock_parquet_urls)
|
||||
|
||||
# Verify dataset creation
|
||||
verify_dataset_creation(ds, mock_hf_dataset)
|
||||
|
||||
finally:
|
||||
# Stop all patches
|
||||
for patch_obj in patches:
|
||||
patch_obj.stop()
|
||||
|
||||
|
||||
def test_from_huggingface_with_resolved_urls(
|
||||
mock_hf_dataset,
|
||||
ray_start_regular_shared,
|
||||
mock_parquet_urls,
|
||||
mock_resolved_urls,
|
||||
mock_ray_dataset,
|
||||
mock_redirected_http_responses,
|
||||
mock_huggingface_datasource,
|
||||
):
|
||||
"""Test the URL resolution logic when HTTP redirects are encountered."""
|
||||
datasource_mock, requests_mock, read_parquet_mock, patches = setup_parquet_mocks(
|
||||
mock_huggingface_datasource,
|
||||
mock_parquet_urls,
|
||||
mock_redirected_http_responses,
|
||||
mock_ray_dataset,
|
||||
)
|
||||
|
||||
try:
|
||||
ds = ray.data.from_huggingface(mock_hf_dataset)
|
||||
|
||||
# Verify HTTP requests
|
||||
verify_http_requests(requests_mock, mock_parquet_urls)
|
||||
|
||||
# Verify read_parquet call with resolved URLs
|
||||
verify_read_parquet_call(read_parquet_mock, mock_resolved_urls)
|
||||
|
||||
# Verify dataset creation
|
||||
verify_dataset_creation(ds, mock_hf_dataset)
|
||||
|
||||
finally:
|
||||
# Stop all patches
|
||||
for patch_obj in patches:
|
||||
patch_obj.stop()
|
||||
|
||||
|
||||
def test_from_huggingface_url_resolution_failures(
|
||||
mock_hf_dataset,
|
||||
ray_start_regular_shared,
|
||||
mock_parquet_urls,
|
||||
mock_ray_dataset,
|
||||
mock_huggingface_datasource,
|
||||
):
|
||||
"""Test URL resolution failures fall back to single node read."""
|
||||
# Convert the mock dataset to an IterableDataset so it uses the read_datasource fallback
|
||||
mock_iterable_dataset = mock_hf_dataset.to_iterable_dataset()
|
||||
|
||||
with patch.object(
|
||||
mock_huggingface_datasource,
|
||||
"list_parquet_urls_from_dataset",
|
||||
return_value=mock_parquet_urls,
|
||||
):
|
||||
# Mock the requests.head calls to simulate failures
|
||||
with patch("requests.head") as mock_requests_head:
|
||||
# Configure mock to raise an exception for all URLs
|
||||
mock_requests_head.side_effect = requests.RequestException(
|
||||
"Connection failed"
|
||||
)
|
||||
|
||||
# Mock the fallback path
|
||||
with patch("ray.data.read_api.read_datasource") as mock_read_datasource:
|
||||
mock_read_datasource.return_value = mock_ray_dataset
|
||||
|
||||
ds = ray.data.from_huggingface(mock_iterable_dataset)
|
||||
|
||||
# Verify that requests.head was called for each URL
|
||||
assert mock_requests_head.call_count == len(mock_parquet_urls)
|
||||
|
||||
# Verify that the fallback read_datasource was called
|
||||
mock_read_datasource.assert_called_once()
|
||||
|
||||
# Verify the dataset was created successfully via fallback
|
||||
verify_dataset_creation(ds, mock_hf_dataset)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,198 @@
|
||||
import os
|
||||
from typing import Dict
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from fsspec.implementations.local import LocalFileSystem
|
||||
from PIL import Image
|
||||
|
||||
import ray
|
||||
from ray.data._internal.datasource.image_datasource import (
|
||||
ImageDatasource,
|
||||
ImageFileMetadataProvider,
|
||||
)
|
||||
from ray.data._internal.tensor_extensions.arrow import (
|
||||
get_arrow_extension_fixed_shape_tensor_types,
|
||||
)
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
|
||||
class TestReadImages:
|
||||
def test_basic(self, ray_start_regular_shared):
|
||||
# "simple" contains three 32x32 RGB images.
|
||||
ds = ray.data.read_images("example://image-datasets/simple")
|
||||
assert ds.schema().names == ["image"]
|
||||
column_type = ds.schema().types[0]
|
||||
assert isinstance(column_type, get_arrow_extension_fixed_shape_tensor_types())
|
||||
assert all(record["image"].shape == (32, 32, 3) for record in ds.take())
|
||||
|
||||
@pytest.mark.parametrize("num_threads", [-1, 0, 1, 2, 4])
|
||||
def test_multi_threading(self, ray_start_regular_shared, num_threads, monkeypatch):
|
||||
monkeypatch.setattr(
|
||||
ray.data._internal.datasource.image_datasource.ImageDatasource,
|
||||
"_NUM_THREADS_PER_TASK",
|
||||
num_threads,
|
||||
)
|
||||
ds = ray.data.read_images(
|
||||
"example://image-datasets/simple",
|
||||
override_num_blocks=1,
|
||||
include_paths=True,
|
||||
)
|
||||
paths = [item["path"][-len("image1.jpg") :] for item in ds.take_all()]
|
||||
if num_threads > 1:
|
||||
# If there are more than 1 threads, the order is not guaranteed.
|
||||
paths = sorted(paths)
|
||||
expected_paths = ["image1.jpg", "image2.jpg", "image3.jpg"]
|
||||
assert paths == expected_paths
|
||||
|
||||
def test_size(self, ray_start_regular_shared):
|
||||
# "different-sizes" contains RGB images with different heights and widths.
|
||||
ds = ray.data.read_images(
|
||||
"example://image-datasets/different-sizes", size=(32, 32)
|
||||
)
|
||||
assert all(record["image"].shape == (32, 32, 3) for record in ds.take())
|
||||
|
||||
def test_different_sizes(self, ray_start_regular_shared):
|
||||
ds = ray.data.read_images("example://image-datasets/different-sizes")
|
||||
assert sorted(record["image"].shape for record in ds.take()) == [
|
||||
(16, 16, 3),
|
||||
(32, 32, 3),
|
||||
(64, 64, 3),
|
||||
]
|
||||
|
||||
@pytest.mark.parametrize("size", [(-32, 32), (32, -32), (-32, -32)])
|
||||
def test_invalid_size(self, ray_start_regular_shared, size):
|
||||
with pytest.raises(ValueError):
|
||||
ray.data.read_images("example://image-datasets/simple", size=size)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"mode, expected_shape", [("L", (32, 32)), ("RGB", (32, 32, 3))]
|
||||
)
|
||||
def test_mode(
|
||||
self,
|
||||
mode,
|
||||
expected_shape,
|
||||
ray_start_regular_shared,
|
||||
):
|
||||
# "different-modes" contains 32x32 images with modes "CMYK", "L", and "RGB"
|
||||
ds = ray.data.read_images("example://image-datasets/different-modes", mode=mode)
|
||||
assert all([record["image"].shape == expected_shape for record in ds.take()])
|
||||
|
||||
def test_e2e_prediction(self, shutdown_only):
|
||||
import torch
|
||||
from torchvision import transforms
|
||||
from torchvision.models import resnet18
|
||||
|
||||
ray.shutdown()
|
||||
ray.init(num_cpus=2)
|
||||
|
||||
dataset = ray.data.read_images("example://image-datasets/simple")
|
||||
transform = transforms.ToTensor()
|
||||
|
||||
def preprocess(batch: Dict[str, np.ndarray]):
|
||||
return {"out": np.stack([transform(image) for image in batch["image"]])}
|
||||
|
||||
dataset = dataset.map_batches(preprocess, batch_format="numpy")
|
||||
|
||||
class Predictor:
|
||||
def __init__(self):
|
||||
self.model = resnet18(pretrained=True)
|
||||
|
||||
def __call__(self, batch: Dict[str, np.ndarray]):
|
||||
with torch.inference_mode():
|
||||
torch_tensor = torch.as_tensor(batch["out"])
|
||||
return {"prediction": self.model(torch_tensor)}
|
||||
|
||||
predictions = dataset.map_batches(
|
||||
Predictor, compute=ray.data.ActorPoolStrategy(min_size=1), batch_size=4096
|
||||
)
|
||||
|
||||
for _ in predictions.iter_batches():
|
||||
pass
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"image_size,image_mode,expected_size,expected_ratio",
|
||||
[(64, "RGB", 30000, 4), (32, "L", 3500, 0.5), (256, "RGBA", 750000, 85)],
|
||||
)
|
||||
def test_data_size_estimate(
|
||||
self,
|
||||
ray_start_regular_shared,
|
||||
image_size,
|
||||
image_mode,
|
||||
expected_size,
|
||||
expected_ratio,
|
||||
):
|
||||
root = "example://image-datasets/different-sizes"
|
||||
ds = ray.data.read_images(
|
||||
root, size=(image_size, image_size), mode=image_mode, override_num_blocks=1
|
||||
)
|
||||
|
||||
data_size = ds.size_bytes()
|
||||
assert data_size >= 0, "estimated data size is out of expected bound"
|
||||
data_size = ds.materialize().size_bytes()
|
||||
assert data_size >= 0, "actual data size is out of expected bound"
|
||||
|
||||
datasource = ImageDatasource(
|
||||
paths=[root],
|
||||
size=(image_size, image_size),
|
||||
mode=image_mode,
|
||||
filesystem=LocalFileSystem(),
|
||||
partitioning=None,
|
||||
meta_provider=ImageFileMetadataProvider(),
|
||||
)
|
||||
assert (
|
||||
datasource._encoding_ratio >= expected_ratio
|
||||
and datasource._encoding_ratio <= expected_ratio * 1.5
|
||||
), "encoding ratio is out of expected bound"
|
||||
data_size = datasource.estimate_inmemory_data_size()
|
||||
assert data_size >= 0, "estimated data size is out of expected bound"
|
||||
|
||||
def test_dynamic_block_split(ray_start_regular_shared):
|
||||
ctx = ray.data.context.DataContext.get_current()
|
||||
target_max_block_size = ctx.target_max_block_size
|
||||
# Reduce target max block size to trigger block splitting on small input.
|
||||
# Otherwise we have to generate big input files, which is unnecessary.
|
||||
ctx.target_max_block_size = 1
|
||||
try:
|
||||
root = "example://image-datasets/simple"
|
||||
ds = ray.data.read_images(root, override_num_blocks=1)
|
||||
assert ds._logical_plan.initial_num_blocks() == 1
|
||||
ds = ds.materialize()
|
||||
# Verify dynamic block splitting taking effect to generate more blocks.
|
||||
assert ds._logical_plan.initial_num_blocks() == 3
|
||||
|
||||
# Test union of same datasets
|
||||
union_ds = ds.union(ds, ds, ds).materialize()
|
||||
assert union_ds._logical_plan.initial_num_blocks() == 12
|
||||
finally:
|
||||
ctx.target_max_block_size = target_max_block_size
|
||||
|
||||
def test_unidentified_image_error(ray_start_regular_shared, tmp_path):
|
||||
path = str(tmp_path / "invalid.png")
|
||||
with open(path, "wb") as file:
|
||||
file.write(b"spam") # Invalid bytes for a PNG file
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
ray.data.read_images(paths=file.name).materialize()
|
||||
|
||||
|
||||
class TestWriteImages:
|
||||
def test_write_images(ray_start_regular_shared, tmp_path):
|
||||
ds = ray.data.read_images("example://image-datasets/simple")
|
||||
ds.write_images(
|
||||
path=tmp_path,
|
||||
column="image",
|
||||
)
|
||||
|
||||
assert len(os.listdir(tmp_path)) == ds.count()
|
||||
|
||||
for filename in os.listdir(tmp_path):
|
||||
path = os.path.join(tmp_path, filename)
|
||||
Image.open(path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,470 @@
|
||||
import gzip
|
||||
import json
|
||||
import os
|
||||
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pyarrow.fs as fs
|
||||
import pyarrow.json as pajson
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data import Schema
|
||||
from ray.data._internal.datasource.json_datasource import PandasJSONDatasource
|
||||
from ray.data._internal.pandas_block import PandasBlockBuilder
|
||||
from ray.data._internal.util import rows_same
|
||||
from ray.data.block import BlockAccessor
|
||||
from ray.data.datasource.file_based_datasource import (
|
||||
FILE_SIZE_FETCH_PARALLELIZATION_THRESHOLD,
|
||||
)
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
# Set the test timeout to 6 minutes
|
||||
pytestmark = pytest.mark.timeout(360)
|
||||
|
||||
|
||||
def test_json_read(
|
||||
ray_start_regular_shared, target_max_block_size_infinite_or_default, tmp_path
|
||||
):
|
||||
df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
||||
path1 = os.path.join(tmp_path, "test1.json")
|
||||
df1.to_json(path1, orient="records", lines=True)
|
||||
ds = ray.data.read_json(path1)
|
||||
dsdf = ds.to_pandas()
|
||||
pd.testing.assert_frame_equal(df1.astype(dsdf.dtypes.to_dict()), dsdf)
|
||||
# Metadata ops.
|
||||
assert ds.count() == 3
|
||||
assert ds.input_files() == [path1]
|
||||
assert ds.schema() == Schema(pa.schema([("one", pa.int64()), ("two", pa.string())]))
|
||||
|
||||
|
||||
def test_zipped_json_read(
|
||||
ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
|
||||
):
|
||||
df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
||||
path1 = os.path.join(tmp_path, "test1.json.gz")
|
||||
df1.to_json(path1, compression="gzip", orient="records", lines=True)
|
||||
ds = ray.data.read_json(path1)
|
||||
dsdf = ds.to_pandas()
|
||||
pd.testing.assert_frame_equal(df1.astype(dsdf.dtypes.to_dict()), dsdf)
|
||||
# Metadata ops.
|
||||
assert ds.count() == 3
|
||||
assert ds.input_files() == [path1]
|
||||
|
||||
|
||||
def test_read_json_fallback_from_pyarrow_failure(
|
||||
ray_start_regular_shared, local_path, target_max_block_size_infinite_or_default
|
||||
):
|
||||
# Try to read this with read_json() to trigger fallback logic
|
||||
# to read bytes with json.load().
|
||||
data = [{"one": [1]}, {"one": [1, 2]}]
|
||||
path1 = os.path.join(local_path, "test1.json")
|
||||
with open(path1, "w") as f:
|
||||
json.dump(data, f)
|
||||
|
||||
# pyarrow.json cannot read JSONs containing arrays of different lengths.
|
||||
from pyarrow import ArrowInvalid
|
||||
|
||||
with pytest.raises(ArrowInvalid):
|
||||
pajson.read_json(path1)
|
||||
|
||||
# Ray Data successfully reads this in by
|
||||
# falling back to json.load() when pyarrow fails.
|
||||
ds = ray.data.read_json(path1)
|
||||
assert ds.take_all() == data
|
||||
|
||||
|
||||
def test_json_read_with_read_options(
|
||||
ray_start_regular_shared,
|
||||
tmp_path,
|
||||
target_max_block_size_infinite_or_default,
|
||||
):
|
||||
# Arrow's JSON ReadOptions isn't serializable in pyarrow < 8.0.0, so this test
|
||||
# covers our custom ReadOptions serializer.
|
||||
# TODO(Clark): Remove this test and our custom serializer once we require
|
||||
# pyarrow >= 8.0.0.
|
||||
|
||||
df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
||||
path1 = os.path.join(tmp_path, "test1.json")
|
||||
df1.to_json(path1, orient="records", lines=True)
|
||||
ds = ray.data.read_json(
|
||||
path1,
|
||||
read_options=pajson.ReadOptions(use_threads=False, block_size=2**30),
|
||||
)
|
||||
dsdf = ds.to_pandas()
|
||||
pd.testing.assert_frame_equal(df1.astype(dsdf.dtypes.to_dict()), dsdf)
|
||||
# Test metadata ops.
|
||||
assert ds.count() == 3
|
||||
assert ds.input_files() == [path1]
|
||||
assert ds.schema() == Schema(pa.schema([("one", pa.int64()), ("two", pa.string())]))
|
||||
|
||||
|
||||
def test_json_read_with_parse_options(
|
||||
ray_start_regular_shared,
|
||||
tmp_path,
|
||||
target_max_block_size_infinite_or_default,
|
||||
):
|
||||
# Arrow's JSON ParseOptions isn't serializable in pyarrow < 8.0.0, so this test
|
||||
# covers our custom ParseOptions serializer, similar to ReadOptions in above test.
|
||||
# TODO(chengsu): Remove this test and our custom serializer once we require
|
||||
# pyarrow >= 8.0.0.
|
||||
|
||||
df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
||||
path1 = os.path.join(tmp_path, "test1.json")
|
||||
df1.to_json(path1, orient="records", lines=True)
|
||||
ds = ray.data.read_json(
|
||||
path1,
|
||||
parse_options=pajson.ParseOptions(
|
||||
explicit_schema=pa.schema([("two", pa.string())]),
|
||||
unexpected_field_behavior="ignore",
|
||||
),
|
||||
)
|
||||
dsdf = ds.to_pandas()
|
||||
assert len(dsdf.columns) == 1
|
||||
pd.testing.assert_series_equal(df1["two"].astype(dsdf["two"].dtype), dsdf["two"])
|
||||
# Test metadata ops.
|
||||
assert ds.count() == 3
|
||||
assert ds.input_files() == [path1]
|
||||
assert ds.schema() == Schema(pa.schema([("two", pa.string())]))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("override_num_blocks", [None, 1, 3])
|
||||
def test_jsonl_lists(
|
||||
ray_start_regular_shared,
|
||||
tmp_path,
|
||||
override_num_blocks,
|
||||
target_max_block_size_infinite_or_default,
|
||||
):
|
||||
"""Test JSONL with mixed types and schemas."""
|
||||
data = [
|
||||
["ray", "rocks", "hello"],
|
||||
["oh", "no"],
|
||||
["rocking", "with", "ray"],
|
||||
]
|
||||
|
||||
path = os.path.join(tmp_path, "test.jsonl")
|
||||
with open(path, "w") as f:
|
||||
for record in data:
|
||||
json.dump(record, f)
|
||||
f.write("\n")
|
||||
|
||||
ds = ray.data.read_json(path, lines=True, override_num_blocks=override_num_blocks)
|
||||
result = ds.take_all()
|
||||
|
||||
assert result[0] == {"0": "ray", "1": "rocks", "2": "hello"}
|
||||
assert result[1] == {"0": "oh", "1": "no", "2": None}
|
||||
assert result[2] == {"0": "rocking", "1": "with", "2": "ray"}
|
||||
|
||||
|
||||
def test_jsonl_mixed_types(
|
||||
ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
|
||||
):
|
||||
"""Test JSONL with mixed types and schemas."""
|
||||
data = [
|
||||
{"a": 1, "b": {"c": 2}}, # Nested dict
|
||||
{"a": 1, "b": {"c": 3}}, # Nested dict
|
||||
{"a": 1, "b": {"c": {"hello": "world"}}}, # Mixed Schema
|
||||
]
|
||||
|
||||
path = os.path.join(tmp_path, "test.jsonl")
|
||||
with open(path, "w") as f:
|
||||
for record in data:
|
||||
json.dump(record, f)
|
||||
f.write("\n")
|
||||
|
||||
ds = ray.data.read_json(path, lines=True)
|
||||
result = ds.take_all()
|
||||
|
||||
assert result[0] == data[0] # Dict stays as is
|
||||
assert result[1] == data[1]
|
||||
assert result[2] == data[2]
|
||||
|
||||
|
||||
def test_json_write(
|
||||
ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
|
||||
):
|
||||
input_df = pd.DataFrame({"id": [0]})
|
||||
ds = ray.data.from_blocks([input_df])
|
||||
|
||||
ds.write_json(tmp_path)
|
||||
|
||||
output_df = pd.concat(
|
||||
[
|
||||
pd.read_json(os.path.join(tmp_path, filename), lines=True)
|
||||
for filename in os.listdir(tmp_path)
|
||||
]
|
||||
)
|
||||
|
||||
assert rows_same(input_df, output_df)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("override_num_blocks", [None, 2])
|
||||
def test_json_roundtrip(
|
||||
ray_start_regular_shared,
|
||||
tmp_path,
|
||||
override_num_blocks,
|
||||
target_max_block_size_infinite_or_default,
|
||||
):
|
||||
df = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
||||
|
||||
ds = ray.data.from_pandas([df], override_num_blocks=override_num_blocks)
|
||||
ds.write_json(tmp_path)
|
||||
|
||||
ds2 = ray.data.read_json(tmp_path)
|
||||
ds2df = ds2.to_pandas()
|
||||
assert rows_same(ds2df, df)
|
||||
for entry in ds2._execute().blocks:
|
||||
assert (
|
||||
# pyrefly: ignore[no-matching-overload]
|
||||
BlockAccessor.for_block(ray.get(entry.ref)).size_bytes()
|
||||
== entry.metadata.size_bytes
|
||||
)
|
||||
|
||||
|
||||
def test_json_read_small_file_unit_block_size(
|
||||
ray_start_regular_shared,
|
||||
tmp_path,
|
||||
target_max_block_size_infinite_or_default,
|
||||
):
|
||||
"""Test reading a small JSON file with unit block_size."""
|
||||
|
||||
df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
||||
path1 = os.path.join(tmp_path, "test1.json")
|
||||
df1.to_json(path1, orient="records", lines=True)
|
||||
ds = ray.data.read_json(path1, read_options=pajson.ReadOptions(block_size=1))
|
||||
dsdf = ds.to_pandas()
|
||||
pd.testing.assert_frame_equal(df1.astype(dsdf.dtypes.to_dict()), dsdf)
|
||||
# Test metadata ops.
|
||||
assert ds.count() == 3
|
||||
assert ds.input_files() == [path1]
|
||||
assert ds.schema() == Schema(pa.schema([("one", pa.int64()), ("two", pa.string())]))
|
||||
|
||||
|
||||
def test_json_read_file_larger_than_block_size(
|
||||
ray_start_regular_shared,
|
||||
tmp_path,
|
||||
target_max_block_size_infinite_or_default,
|
||||
):
|
||||
"""Test reading a JSON file larger than the block size."""
|
||||
block_size = 1024
|
||||
num_chars = 2500
|
||||
num_rows = 3
|
||||
df2 = pd.DataFrame(
|
||||
{
|
||||
"one": ["a" * num_chars for _ in range(num_rows)],
|
||||
"two": ["b" * num_chars for _ in range(num_rows)],
|
||||
}
|
||||
)
|
||||
path2 = os.path.join(tmp_path, "test2.json")
|
||||
df2.to_json(path2, orient="records", lines=True)
|
||||
ds = ray.data.read_json(
|
||||
path2, read_options=pajson.ReadOptions(block_size=block_size)
|
||||
)
|
||||
dsdf = ds.to_pandas()
|
||||
pd.testing.assert_frame_equal(df2.astype(dsdf.dtypes.to_dict()), dsdf)
|
||||
# Test metadata ops.
|
||||
assert ds.count() == num_rows
|
||||
assert ds.input_files() == [path2]
|
||||
assert ds.schema() == Schema(
|
||||
pa.schema([("one", pa.string()), ("two", pa.string())])
|
||||
)
|
||||
|
||||
|
||||
def test_json_read_negative_block_size_fallback(
|
||||
ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
|
||||
):
|
||||
"""Test reading JSON with negative block_size triggers fallback to json.load()."""
|
||||
|
||||
df3 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
||||
path3 = os.path.join(tmp_path, "test3.json")
|
||||
df3.to_json(path3, orient="records", lines=True)
|
||||
|
||||
# Negative Buffer Size, fails with arrow but succeeds in fallback to json.load()
|
||||
ds = ray.data.read_json(path3, read_options=pajson.ReadOptions(block_size=-1))
|
||||
dsdf = ds.to_pandas()
|
||||
pd.testing.assert_frame_equal(df3.astype(dsdf.dtypes.to_dict()), dsdf)
|
||||
|
||||
|
||||
def test_json_read_zero_block_size_failure(
|
||||
ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
|
||||
):
|
||||
"""Test reading JSON with zero block_size fails in both arrow and fallback."""
|
||||
|
||||
df3 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
||||
path3 = os.path.join(tmp_path, "test3.json")
|
||||
df3.to_json(path3, orient="records", lines=True)
|
||||
|
||||
# Zero Buffer Size, fails with arrow and fails in fallback to json.load()
|
||||
with pytest.raises(json.decoder.JSONDecodeError, match="Extra data"):
|
||||
ds = ray.data.read_json(path3, read_options=pajson.ReadOptions(block_size=0))
|
||||
dsdf = ds.to_pandas()
|
||||
assert dsdf.equals(df3)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("min_rows_per_file", [5, 10, 50])
|
||||
def test_write_min_rows_per_file(
|
||||
tmp_path,
|
||||
ray_start_regular_shared,
|
||||
min_rows_per_file,
|
||||
target_max_block_size_infinite_or_default,
|
||||
):
|
||||
ray.data.range(100, override_num_blocks=20).write_json(
|
||||
tmp_path, min_rows_per_file=min_rows_per_file
|
||||
)
|
||||
|
||||
for filename in os.listdir(tmp_path):
|
||||
with open(os.path.join(tmp_path, filename), "r") as file:
|
||||
num_rows_written = len(file.read().splitlines())
|
||||
assert num_rows_written == min_rows_per_file
|
||||
|
||||
|
||||
def test_mixed_gzipped_json_files(
|
||||
ray_start_regular_shared, tmp_path, target_max_block_size_infinite_or_default
|
||||
):
|
||||
# Create a non-empty gzipped JSON file
|
||||
non_empty_file_path = os.path.join(tmp_path, "non_empty.json.gz")
|
||||
data = [{"col1": "value1", "col2": "value2", "col3": "value3"}]
|
||||
with gzip.open(non_empty_file_path, "wt", encoding="utf-8") as f:
|
||||
for record in data:
|
||||
json.dump(record, f)
|
||||
f.write("\n")
|
||||
|
||||
# Create an empty gzipped JSON file
|
||||
empty_file_path = os.path.join(tmp_path, "empty.json.gz")
|
||||
with gzip.open(empty_file_path, "wt", encoding="utf-8"):
|
||||
pass # Write nothing to create an empty file
|
||||
|
||||
# Attempt to read both files with Ray
|
||||
ds = ray.data.read_json(
|
||||
[non_empty_file_path, empty_file_path],
|
||||
arrow_open_stream_args={"compression": "gzip"},
|
||||
)
|
||||
|
||||
# The dataset should only contain data from the non-empty file
|
||||
assert ds.count() == 1
|
||||
# Iterate through each row in the dataset and compare with the expected data
|
||||
for row in ds.iter_rows():
|
||||
assert row == data[0], f"Row {row} does not match expected {data[0]}"
|
||||
|
||||
# Verify the data content using take
|
||||
retrieved_data = ds.take(1)[0]
|
||||
assert (
|
||||
retrieved_data == data[0]
|
||||
), f"Retrieved data {retrieved_data} does not match expected {data[0]}."
|
||||
|
||||
|
||||
def test_json_with_http_path_parallelization(
|
||||
ray_start_regular_shared, httpserver, target_max_block_size_infinite_or_default
|
||||
):
|
||||
num_files = FILE_SIZE_FETCH_PARALLELIZATION_THRESHOLD
|
||||
urls = []
|
||||
for i in range(num_files):
|
||||
httpserver.expect_request(f"/file{i}.json").respond_with_json({"id": i})
|
||||
urls.append(httpserver.url_for(f"/file{i}.json"))
|
||||
|
||||
ds = ray.data.read_json(urls)
|
||||
actual_rows = ds.take_all()
|
||||
|
||||
expected_rows = [{"id": i} for i in range(num_files)]
|
||||
assert sorted(actual_rows, key=lambda row: row["id"]) == sorted(
|
||||
expected_rows, key=lambda row: row["id"]
|
||||
)
|
||||
|
||||
|
||||
class TestPandasJSONDatasource:
|
||||
@pytest.mark.parametrize(
|
||||
"data",
|
||||
[{"a": []}, {"a": [1]}, {"a": [1, 2, 3]}],
|
||||
ids=["empty", "single", "multiple"],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"compression,filename",
|
||||
[("gzip", "test.json.gz"), ("infer", "test.json")], # infer = default
|
||||
)
|
||||
def test_read_stream(
|
||||
self,
|
||||
data,
|
||||
tmp_path,
|
||||
compression,
|
||||
filename,
|
||||
target_max_block_size_infinite_or_default,
|
||||
):
|
||||
# Setup test file.
|
||||
df = pd.DataFrame(data)
|
||||
path = os.path.join(tmp_path, filename)
|
||||
df.to_json(path, orient="records", lines=True, compression=compression)
|
||||
|
||||
# Setup datasource.
|
||||
local_filesystem = fs.LocalFileSystem()
|
||||
source = PandasJSONDatasource(
|
||||
path, target_output_size_bytes=1, filesystem=local_filesystem
|
||||
)
|
||||
|
||||
# Read stream.
|
||||
block_builder = PandasBlockBuilder()
|
||||
with source._open_input_source(local_filesystem, path) as f:
|
||||
for block in source._read_stream(f, path):
|
||||
block_builder.add_block(block)
|
||||
block = block_builder.build()
|
||||
|
||||
# Verify.
|
||||
assert rows_same(block, df)
|
||||
|
||||
def test_read_stream_with_target_output_size_bytes(
|
||||
self, tmp_path, target_max_block_size_infinite_or_default
|
||||
):
|
||||
# Setup test file. It contains 16 lines, each line is 8 MiB.
|
||||
df = pd.DataFrame({"data": ["a" * 8 * 1024 * 1024] * 16})
|
||||
path = os.path.join(tmp_path, "test.json")
|
||||
df.to_json(path, orient="records", lines=True)
|
||||
|
||||
# Setup datasource. It should read 32 MiB (4 lines) per output.
|
||||
local_filesystem = fs.LocalFileSystem()
|
||||
source = PandasJSONDatasource(
|
||||
path,
|
||||
target_output_size_bytes=32 * 1024 * 1024,
|
||||
filesystem=local_filesystem,
|
||||
)
|
||||
|
||||
# Read stream.
|
||||
block_builder = PandasBlockBuilder()
|
||||
with source._open_input_source(local_filesystem, path) as f:
|
||||
for block in source._read_stream(f, path):
|
||||
assert len(block) == 4
|
||||
block_builder.add_block(block)
|
||||
block = block_builder.build()
|
||||
|
||||
# Verify.
|
||||
assert rows_same(block, df)
|
||||
|
||||
def test_read_stream_with_advanced_file_pointer(
|
||||
self, tmp_path, target_max_block_size_infinite_or_default
|
||||
):
|
||||
# Setup test file.
|
||||
df = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
||||
path = os.path.join(tmp_path, "test.json")
|
||||
df.to_json(path, orient="records", lines=True)
|
||||
|
||||
# Setup datasource.
|
||||
local_filesystem = fs.LocalFileSystem()
|
||||
source = PandasJSONDatasource(
|
||||
path, target_output_size_bytes=1, filesystem=local_filesystem
|
||||
)
|
||||
|
||||
# Simulate retrying a stream read on a file handle that was already consumed.
|
||||
block_builder = PandasBlockBuilder()
|
||||
with source._open_input_source(local_filesystem, path) as f:
|
||||
f.read(1)
|
||||
for block in source._read_stream(f, path):
|
||||
block_builder.add_block(block)
|
||||
block = block_builder.build()
|
||||
|
||||
# Verify.
|
||||
assert rows_same(block, df)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,390 @@
|
||||
import os
|
||||
|
||||
import lance
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from packaging.version import Version
|
||||
from pytest_lazy_fixtures import lf as lazy_fixture
|
||||
|
||||
import ray
|
||||
from ray._common.test_utils import wait_for_condition
|
||||
from ray.data import Schema
|
||||
from ray.data._internal.datasource.lance_datasink import (
|
||||
_WRITE_LANCE_FRAGMENTS_DESCRIPTION,
|
||||
LanceDatasink,
|
||||
_write_fragment,
|
||||
)
|
||||
from ray.data.datasource import SaveMode
|
||||
from ray.data.datasource.path_util import _unwrap_protocol
|
||||
|
||||
# Skip tests for older pylance versions (<=0.3.19) due to incompatible lance API changes with pyarrow v9.0.0
|
||||
pytestmark = pytest.mark.skipif(
|
||||
Version(lance.__version__) <= Version("0.3.19"),
|
||||
reason=f"pylance {lance.__version__} <= 0.3.19; API incompatible",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"fs,data_path",
|
||||
[
|
||||
(None, lazy_fixture("local_path")),
|
||||
(lazy_fixture("local_fs"), lazy_fixture("local_path")),
|
||||
(lazy_fixture("s3_fs"), lazy_fixture("s3_path")),
|
||||
(
|
||||
lazy_fixture("s3_fs_with_space"),
|
||||
lazy_fixture("s3_path_with_space"),
|
||||
), # Path contains space.
|
||||
(
|
||||
lazy_fixture("s3_fs_with_anonymous_crendential"),
|
||||
lazy_fixture("s3_path_with_anonymous_crendential"),
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"batch_size",
|
||||
[None, 100],
|
||||
)
|
||||
def test_lance_read_basic(fs, data_path, batch_size):
|
||||
df1 = pa.table({"one": [2, 1, 3, 4, 6, 5], "two": ["b", "a", "c", "e", "g", "f"]})
|
||||
setup_data_path = _unwrap_protocol(data_path)
|
||||
path = os.path.join(setup_data_path, "test.lance")
|
||||
lance.write_dataset(df1, path)
|
||||
|
||||
ds_lance = lance.dataset(path)
|
||||
assert ds_lance is not None
|
||||
df2 = pa.table(
|
||||
{
|
||||
"one": [1, 2, 3, 4, 5, 6],
|
||||
"three": [4, 5, 8, 9, 12, 13],
|
||||
"four": ["u", "v", "w", "x", "y", "z"],
|
||||
}
|
||||
)
|
||||
ds_lance.merge(df2, "one")
|
||||
|
||||
if batch_size is None:
|
||||
ds = ray.data.read_lance(path)
|
||||
else:
|
||||
ds = ray.data.read_lance(path, scanner_options={"batch_size": batch_size})
|
||||
|
||||
# Test metadata-only ops.
|
||||
assert ds.count() == 6
|
||||
assert ds.schema() == Schema(
|
||||
pa.schema(
|
||||
{
|
||||
"one": pa.int64(),
|
||||
"two": pa.string(),
|
||||
"three": pa.int64(),
|
||||
"four": pa.string(),
|
||||
}
|
||||
)
|
||||
)
|
||||
|
||||
# Test read.
|
||||
values = [[s["one"], s["two"]] for s in ds.take_all()]
|
||||
assert sorted(values) == [
|
||||
[1, "a"],
|
||||
[2, "b"],
|
||||
[3, "c"],
|
||||
[4, "e"],
|
||||
[5, "f"],
|
||||
[6, "g"],
|
||||
]
|
||||
|
||||
# Test column projection.
|
||||
ds = ray.data.read_lance(path, columns=["one"])
|
||||
values = [s["one"] for s in ds.take_all()]
|
||||
assert sorted(values) == [1, 2, 3, 4, 5, 6]
|
||||
assert ds.schema().names == ["one"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
|
||||
def test_lance_read_with_scanner_fragments(data_path):
|
||||
table = pa.table({"one": [2, 1, 3, 4, 6, 5], "two": ["b", "a", "c", "e", "g", "f"]})
|
||||
setup_data_path = _unwrap_protocol(data_path)
|
||||
path = os.path.join(setup_data_path, "test.lance")
|
||||
dataset = lance.write_dataset(table, path, max_rows_per_file=2)
|
||||
assert dataset is not None
|
||||
|
||||
fragments = dataset.get_fragments()
|
||||
ds = ray.data.read_lance(path, scanner_options={"fragments": fragments[:1]})
|
||||
values = [[s["one"], s["two"]] for s in ds.take_all()]
|
||||
assert values == [
|
||||
[2, "b"],
|
||||
[1, "a"],
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
|
||||
def test_lance_read_many_files(data_path):
|
||||
setup_data_path = _unwrap_protocol(data_path)
|
||||
path = os.path.join(setup_data_path, "test.lance")
|
||||
num_rows = 1024
|
||||
data = pa.table({"id": pa.array(range(num_rows))})
|
||||
lance.write_dataset(data, path, max_rows_per_file=1)
|
||||
|
||||
def test_lance():
|
||||
ds = ray.data.read_lance(path)
|
||||
return ds.count() == num_rows
|
||||
|
||||
wait_for_condition(test_lance, timeout=10)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
|
||||
def test_lance_write(data_path):
|
||||
schema = pa.schema([pa.field("id", pa.int64()), pa.field("str", pa.string())])
|
||||
|
||||
ray.data.range(10).map(
|
||||
lambda x: {"id": x["id"], "str": f"str-{x['id']}"}
|
||||
).write_lance(data_path, schema=schema)
|
||||
|
||||
ds = lance.dataset(data_path)
|
||||
assert ds is not None
|
||||
ds.count_rows() == 10
|
||||
assert ds.schema.names == schema.names
|
||||
# The schema is platform-dependent, because numpy uses int32 on Windows.
|
||||
# So we observe the schema that is written and use that.
|
||||
schema = ds.schema
|
||||
|
||||
tbl = ds.to_table()
|
||||
assert sorted(tbl["id"].to_pylist()) == list(range(10))
|
||||
assert set(tbl["str"].to_pylist()) == {f"str-{i}" for i in range(10)}
|
||||
|
||||
ray.data.range(10).map(
|
||||
lambda x: {"id": x["id"] + 10, "str": f"str-{x['id'] + 10}"}
|
||||
).write_lance(data_path, mode=SaveMode.APPEND)
|
||||
|
||||
ds = lance.dataset(data_path)
|
||||
assert ds is not None
|
||||
ds.count_rows() == 20
|
||||
tbl = ds.to_table()
|
||||
assert sorted(tbl["id"].to_pylist()) == list(range(20))
|
||||
assert set(tbl["str"].to_pylist()) == {f"str-{i}" for i in range(20)}
|
||||
|
||||
ray.data.range(10).map(
|
||||
lambda x: {"id": x["id"], "str": f"str-{x['id']}"}
|
||||
).write_lance(data_path, schema=schema, mode=SaveMode.OVERWRITE)
|
||||
|
||||
ds = lance.dataset(data_path)
|
||||
assert ds is not None
|
||||
ds.count_rows() == 10
|
||||
assert ds.schema == schema
|
||||
|
||||
|
||||
@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
|
||||
def test_lance_write_create_errors_if_exists(data_path):
|
||||
table_path = os.path.join(data_path, "my_table")
|
||||
ds = ray.data.range(10)
|
||||
|
||||
# First CREATE succeeds on an empty destination.
|
||||
ds.write_lance(table_path, mode=SaveMode.CREATE)
|
||||
assert lance.dataset(table_path).count_rows() == 10
|
||||
|
||||
# A second CREATE must error instead of silently overwriting.
|
||||
with pytest.raises(ValueError, match="already exists"):
|
||||
ray.data.range(5).write_lance(table_path, mode=SaveMode.CREATE)
|
||||
|
||||
# Existing data is untouched.
|
||||
assert lance.dataset(table_path).count_rows() == 10
|
||||
|
||||
# CREATE is also the default mode, so it must guard too.
|
||||
with pytest.raises(ValueError, match="already exists"):
|
||||
ray.data.range(5).write_lance(table_path)
|
||||
assert lance.dataset(table_path).count_rows() == 10
|
||||
|
||||
# OVERWRITE replaces the existing data.
|
||||
ray.data.range(5).write_lance(table_path, mode=SaveMode.OVERWRITE)
|
||||
assert lance.dataset(table_path).count_rows() == 5
|
||||
|
||||
|
||||
@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
|
||||
def test_lance_write_append_errors_if_missing(data_path):
|
||||
table_path = os.path.join(data_path, "missing_table")
|
||||
# APPEND surfaces Lance's own "not found" error. We don't pin the message,
|
||||
# since it can change across Lance versions.
|
||||
expected_errors: tuple[type[Exception], ...] = (
|
||||
ValueError,
|
||||
OSError,
|
||||
FileNotFoundError,
|
||||
)
|
||||
with pytest.raises(expected_errors):
|
||||
ray.data.range(5).write_lance(table_path, mode=SaveMode.APPEND)
|
||||
assert not os.path.exists(table_path)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
|
||||
def test_lance_write_min_rows_per_file(data_path):
|
||||
schema = pa.schema([pa.field("id", pa.int64()), pa.field("str", pa.string())])
|
||||
|
||||
ray.data.range(10).map(
|
||||
lambda x: {"id": x["id"], "str": f"str-{x['id']}"}
|
||||
).write_lance(data_path, schema=schema, min_rows_per_file=100)
|
||||
|
||||
ds = lance.dataset(data_path)
|
||||
assert ds is not None
|
||||
assert ds.count_rows() == 10
|
||||
assert ds.schema == schema
|
||||
|
||||
assert len(ds.get_fragments()) == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
|
||||
def test_lance_write_max_rows_per_file(data_path):
|
||||
schema = pa.schema([pa.field("id", pa.int64()), pa.field("str", pa.string())])
|
||||
|
||||
ray.data.range(10).map(
|
||||
lambda x: {"id": x["id"], "str": f"str-{x['id']}"}
|
||||
).write_lance(data_path, schema=schema, max_rows_per_file=1)
|
||||
|
||||
ds = lance.dataset(data_path)
|
||||
assert ds is not None
|
||||
assert ds.count_rows() == 10
|
||||
assert ds.schema == schema
|
||||
|
||||
assert len(ds.get_fragments()) == 10
|
||||
|
||||
|
||||
@pytest.mark.parametrize("data_path", [lazy_fixture("local_path")])
|
||||
def test_lance_read_with_version(data_path):
|
||||
# Write an initial dataset (version 1)
|
||||
df1 = pa.table({"one": [2, 1, 3, 4, 6, 5], "two": ["b", "a", "c", "e", "g", "f"]})
|
||||
setup_data_path = _unwrap_protocol(data_path)
|
||||
path = os.path.join(setup_data_path, "test_version.lance")
|
||||
lance.write_dataset(df1, path)
|
||||
|
||||
# Merge new data to create a later version (latest)
|
||||
ds_lance = lance.dataset(path)
|
||||
assert ds_lance is not None
|
||||
# Get the initial version
|
||||
initial_version = ds_lance.version
|
||||
|
||||
df2 = pa.table(
|
||||
{
|
||||
"one": [1, 2, 3, 4, 5, 6],
|
||||
"three": [4, 5, 8, 9, 12, 13],
|
||||
"four": ["u", "v", "w", "x", "y", "z"],
|
||||
}
|
||||
)
|
||||
ds_lance.merge(df2, "one")
|
||||
|
||||
# Default read should return the latest (merged) dataset.
|
||||
ds_latest = ray.data.read_lance(path)
|
||||
|
||||
assert ds_latest.count() == 6
|
||||
# Latest dataset should contain merged columns
|
||||
assert "three" in ds_latest.schema().names
|
||||
|
||||
# Read the initial version and ensure it contains the original columns
|
||||
ds_prev = ray.data.read_lance(path, version=initial_version)
|
||||
assert ds_prev.count() == 6
|
||||
assert ds_prev.schema().names == ["one", "two"]
|
||||
|
||||
values_prev = [[s["one"], s["two"]] for s in ds_prev.take_all()]
|
||||
assert sorted(values_prev) == [
|
||||
[1, "a"],
|
||||
[2, "b"],
|
||||
[3, "c"],
|
||||
[4, "e"],
|
||||
[5, "f"],
|
||||
[6, "g"],
|
||||
]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_lance_write(monkeypatch):
|
||||
captured = {}
|
||||
|
||||
class _FakeLanceDatasink:
|
||||
def __init__(self, path, **kwargs):
|
||||
captured["path"] = path
|
||||
captured["kwargs"] = kwargs
|
||||
|
||||
def _fake_write_datasink(self, datasink, **kwargs):
|
||||
captured["datasink"] = datasink
|
||||
captured["write_kwargs"] = kwargs
|
||||
|
||||
monkeypatch.setattr(ray.data.dataset, "LanceDatasink", _FakeLanceDatasink)
|
||||
monkeypatch.setattr(ray.data.Dataset, "write_datasink", _fake_write_datasink)
|
||||
|
||||
return captured, _FakeLanceDatasink
|
||||
|
||||
|
||||
def test_write_lance_passes_namespace_args(mock_lance_write):
|
||||
captured, fake_lance_datasink_cls = mock_lance_write
|
||||
table_id = ["db", "table"]
|
||||
namespace_impl = "dir"
|
||||
namespace_properties = {"path": "/tmp/ns"}
|
||||
|
||||
ds = ray.data.range(1)
|
||||
ds.write_lance(
|
||||
"/tmp/lance-namespace-test",
|
||||
table_id=table_id,
|
||||
namespace_impl=namespace_impl,
|
||||
namespace_properties=namespace_properties,
|
||||
)
|
||||
|
||||
assert captured["path"] == "/tmp/lance-namespace-test"
|
||||
assert captured["kwargs"]["table_id"] == table_id
|
||||
assert captured["kwargs"]["namespace_impl"] == namespace_impl
|
||||
assert captured["kwargs"]["namespace_properties"] == namespace_properties
|
||||
assert isinstance(captured["datasink"], fake_lance_datasink_cls)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", [SaveMode.APPEND, SaveMode.OVERWRITE])
|
||||
def test_lance_namespace_write_rejects_non_create_mode(monkeypatch, mode):
|
||||
class _FakeNamespace:
|
||||
pass
|
||||
|
||||
monkeypatch.setattr(
|
||||
"ray.data._internal.datasource.lance_datasink.get_or_create_namespace",
|
||||
lambda namespace_impl, namespace_properties: _FakeNamespace(),
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="Namespace writes currently only support"):
|
||||
LanceDatasink(
|
||||
uri="/tmp/lance-namespace-test",
|
||||
mode=mode,
|
||||
table_id=["db", "table"],
|
||||
namespace_impl="dir",
|
||||
namespace_properties={"path": "/tmp/ns"},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"max_attempts,expected_blocks_consumed_before_write",
|
||||
[(1, 1), (2, 3)],
|
||||
)
|
||||
def test_write_fragment_only_materializes_stream_when_retrying(
|
||||
monkeypatch, max_attempts, expected_blocks_consumed_before_write
|
||||
):
|
||||
import lance.fragment
|
||||
|
||||
consumed = {"count": 0}
|
||||
blocks = [pa.table({"id": [i]}) for i in range(3)]
|
||||
|
||||
def block_stream():
|
||||
for block in blocks:
|
||||
consumed["count"] += 1
|
||||
yield block
|
||||
|
||||
def fake_write_fragments(reader, uri, **kwargs):
|
||||
assert consumed["count"] == expected_blocks_consumed_before_write
|
||||
return []
|
||||
|
||||
monkeypatch.setattr(lance.fragment, "write_fragments", fake_write_fragments)
|
||||
|
||||
_write_fragment(
|
||||
block_stream(),
|
||||
"/tmp/lance-materialization-test",
|
||||
retry_params={
|
||||
"description": _WRITE_LANCE_FRAGMENTS_DESCRIPTION,
|
||||
"match": [],
|
||||
"max_attempts": max_attempts,
|
||||
"max_backoff_s": 0,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,394 @@
|
||||
import importlib.util
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data.datasource.path_util import _unwrap_protocol
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
# Skip all tests if mcap is not available
|
||||
MCAP_AVAILABLE = importlib.util.find_spec("mcap") is not None
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not MCAP_AVAILABLE,
|
||||
reason="mcap module not available. Install with: pip install mcap",
|
||||
)
|
||||
|
||||
|
||||
def create_test_mcap_file(file_path: str, messages: list) -> None:
|
||||
"""Create a test MCAP file with given messages."""
|
||||
from mcap.writer import Writer
|
||||
|
||||
with open(file_path, "wb") as stream:
|
||||
writer = Writer(stream)
|
||||
writer.start(profile="", library="ray-test")
|
||||
|
||||
# Register schema
|
||||
schema_id = writer.register_schema(
|
||||
name="test_schema",
|
||||
encoding="jsonschema",
|
||||
data=json.dumps(
|
||||
{
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"value": {"type": "number"},
|
||||
"name": {"type": "string"},
|
||||
},
|
||||
}
|
||||
).encode(),
|
||||
)
|
||||
|
||||
# Register channels and write messages
|
||||
channels = {}
|
||||
for msg in messages:
|
||||
topic = msg["topic"]
|
||||
if topic not in channels:
|
||||
channels[topic] = writer.register_channel(
|
||||
schema_id=schema_id,
|
||||
topic=topic,
|
||||
message_encoding="json",
|
||||
)
|
||||
|
||||
writer.add_message(
|
||||
channel_id=channels[topic],
|
||||
log_time=msg["log_time"],
|
||||
publish_time=msg.get("publish_time", msg["log_time"]),
|
||||
data=json.dumps(msg["data"]).encode(),
|
||||
)
|
||||
|
||||
writer.finish()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def simple_mcap_file(tmp_path):
|
||||
"""Fixture providing a simple MCAP file with one message."""
|
||||
path = os.path.join(tmp_path, "test.mcap")
|
||||
messages = [
|
||||
{
|
||||
"topic": "/test",
|
||||
"data": {"value": 1},
|
||||
"log_time": 1000000000,
|
||||
}
|
||||
]
|
||||
create_test_mcap_file(path, messages)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def basic_mcap_file(tmp_path):
|
||||
"""Fixture providing a basic MCAP file with two different topics."""
|
||||
path = os.path.join(tmp_path, "test.mcap")
|
||||
messages = [
|
||||
{
|
||||
"topic": "/camera/image",
|
||||
"data": {"frame_id": 1, "timestamp": 1000},
|
||||
"log_time": 1000000000,
|
||||
},
|
||||
{
|
||||
"topic": "/lidar/points",
|
||||
"data": {"point_count": 1024, "timestamp": 2000},
|
||||
"log_time": 2000000000,
|
||||
},
|
||||
]
|
||||
create_test_mcap_file(path, messages)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def multi_topic_mcap_file(tmp_path):
|
||||
"""Fixture providing an MCAP file with 9 messages across 3 topics."""
|
||||
path = os.path.join(tmp_path, "multi_topic.mcap")
|
||||
base_time = 1000000000
|
||||
messages = []
|
||||
for i in range(9):
|
||||
topics = ["/topic_a", "/topic_b", "/topic_c"]
|
||||
topic = topics[i % 3]
|
||||
messages.append(
|
||||
{
|
||||
"topic": topic,
|
||||
"data": {"seq": i, "topic": topic},
|
||||
"log_time": base_time + i * 1000000,
|
||||
}
|
||||
)
|
||||
create_test_mcap_file(path, messages)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def time_series_mcap_file(tmp_path):
|
||||
"""Fixture providing an MCAP file with 10 time-sequenced messages."""
|
||||
path = os.path.join(tmp_path, "time_test.mcap")
|
||||
base_time = 1000000000
|
||||
messages = [
|
||||
{
|
||||
"topic": "/test_topic",
|
||||
"data": {"seq": i},
|
||||
"log_time": base_time + i * 1000000,
|
||||
}
|
||||
for i in range(10)
|
||||
]
|
||||
create_test_mcap_file(path, messages)
|
||||
return path, base_time
|
||||
|
||||
|
||||
def test_read_mcap_basic(ray_start_regular_shared, basic_mcap_file):
|
||||
"""Test basic MCAP file reading."""
|
||||
ds = ray.data.read_mcap(basic_mcap_file)
|
||||
|
||||
# Test metadata operations
|
||||
assert ds.count() == 2
|
||||
assert ds.input_files() == [_unwrap_protocol(basic_mcap_file)]
|
||||
|
||||
# Verify basic fields are present
|
||||
rows = ds.take_all()
|
||||
for row in rows:
|
||||
assert "data" in row
|
||||
assert "topic" in row
|
||||
assert "log_time" in row
|
||||
assert "publish_time" in row
|
||||
|
||||
|
||||
def test_read_mcap_multiple_files(ray_start_regular_shared, tmp_path):
|
||||
"""Test reading multiple MCAP files."""
|
||||
paths = []
|
||||
for i in range(2):
|
||||
path = os.path.join(tmp_path, f"test_{i}.mcap")
|
||||
messages = [
|
||||
{
|
||||
"topic": f"/test_{i}",
|
||||
"data": {"file_id": i},
|
||||
"log_time": 1000000000 + i * 1000000,
|
||||
}
|
||||
]
|
||||
create_test_mcap_file(path, messages)
|
||||
paths.append(path)
|
||||
|
||||
ds = ray.data.read_mcap(paths)
|
||||
assert ds.count() == 2
|
||||
assert set(ds.input_files()) == {_unwrap_protocol(p) for p in paths}
|
||||
|
||||
rows = ds.take_all()
|
||||
file_ids = {row["data"]["file_id"] for row in rows}
|
||||
assert file_ids == {0, 1}
|
||||
|
||||
|
||||
def test_read_mcap_directory(ray_start_regular_shared, tmp_path):
|
||||
"""Test reading MCAP files from a directory."""
|
||||
# Create MCAP files in directory
|
||||
for i in range(2):
|
||||
path = os.path.join(tmp_path, f"data_{i}.mcap")
|
||||
messages = [
|
||||
{
|
||||
"topic": f"/dir_test_{i}",
|
||||
"data": {"index": i},
|
||||
"log_time": 1000000000 + i * 1000000,
|
||||
}
|
||||
]
|
||||
create_test_mcap_file(path, messages)
|
||||
|
||||
ds = ray.data.read_mcap(tmp_path)
|
||||
assert ds.count() == 2
|
||||
|
||||
|
||||
def test_read_mcap_topic_filtering(ray_start_regular_shared, multi_topic_mcap_file):
|
||||
"""Test filtering by topics."""
|
||||
# Test topic filtering
|
||||
topics = {"/topic_a", "/topic_b"}
|
||||
ds = ray.data.read_mcap(multi_topic_mcap_file, topics=topics)
|
||||
|
||||
rows = ds.take_all()
|
||||
actual_topics = {row["topic"] for row in rows}
|
||||
assert actual_topics.issubset(topics)
|
||||
assert len(rows) == 6 # 2/3 of messages
|
||||
|
||||
|
||||
def test_read_mcap_time_range_filtering(
|
||||
ray_start_regular_shared, time_series_mcap_file
|
||||
):
|
||||
"""Test filtering by time range."""
|
||||
path, base_time = time_series_mcap_file
|
||||
|
||||
# Filter to first 5 messages
|
||||
time_range = (base_time, base_time + 5000000)
|
||||
ds = ray.data.read_mcap(path, time_range=time_range)
|
||||
|
||||
rows = ds.take_all()
|
||||
assert len(rows) <= 5
|
||||
for row in rows:
|
||||
assert base_time <= row["log_time"] <= base_time + 5000000
|
||||
|
||||
|
||||
def test_read_mcap_message_type_filtering(ray_start_regular_shared, simple_mcap_file):
|
||||
"""Test filtering by message types."""
|
||||
# Filter with existing schema
|
||||
ds = ray.data.read_mcap(simple_mcap_file, message_types={"test_schema"})
|
||||
assert ds.count() == 1
|
||||
|
||||
# Filter with non-existent schema
|
||||
ds = ray.data.read_mcap(simple_mcap_file, message_types={"nonexistent"})
|
||||
assert ds.count() == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("include_metadata", [True, False])
|
||||
def test_read_mcap_include_metadata(
|
||||
ray_start_regular_shared, simple_mcap_file, include_metadata
|
||||
):
|
||||
"""Test include_metadata option."""
|
||||
ds = ray.data.read_mcap(simple_mcap_file, include_metadata=include_metadata)
|
||||
rows = ds.take_all()
|
||||
|
||||
if include_metadata:
|
||||
assert "schema_name" in rows[0]
|
||||
assert "channel_id" in rows[0]
|
||||
else:
|
||||
assert "schema_name" not in rows[0]
|
||||
assert "channel_id" not in rows[0]
|
||||
|
||||
|
||||
def test_read_mcap_include_paths(ray_start_regular_shared, simple_mcap_file):
|
||||
"""Test include_paths option."""
|
||||
ds = ray.data.read_mcap(simple_mcap_file, include_paths=True)
|
||||
rows = ds.take_all()
|
||||
|
||||
for row in rows:
|
||||
assert "path" in row
|
||||
assert simple_mcap_file in row["path"]
|
||||
|
||||
|
||||
def test_read_mcap_invalid_time_range(ray_start_regular_shared, simple_mcap_file):
|
||||
"""Test validation of time range parameters."""
|
||||
# Start time >= end time
|
||||
with pytest.raises(ValueError, match="start_time must be less than end_time"):
|
||||
ray.data.read_mcap(simple_mcap_file, time_range=(2000, 1000))
|
||||
|
||||
# Negative times
|
||||
with pytest.raises(ValueError, match="time values must be non-negative"):
|
||||
ray.data.read_mcap(simple_mcap_file, time_range=(-1000, 2000))
|
||||
|
||||
|
||||
def test_read_mcap_missing_dependency(ray_start_regular_shared, simple_mcap_file):
|
||||
"""Test graceful failure when mcap library is missing."""
|
||||
from unittest.mock import patch
|
||||
|
||||
with patch.dict("sys.modules", {"mcap": None}):
|
||||
with pytest.raises(ImportError, match="MCAPDatasource.*depends on 'mcap'"):
|
||||
ray.data.read_mcap(simple_mcap_file)
|
||||
|
||||
|
||||
def test_read_mcap_nonexistent_file(ray_start_regular_shared):
|
||||
"""Test handling of nonexistent files."""
|
||||
with pytest.raises(Exception): # FileNotFoundError or similar
|
||||
ds = ray.data.read_mcap("/nonexistent/file.mcap")
|
||||
ds.materialize() # Force execution
|
||||
|
||||
|
||||
@pytest.mark.parametrize("override_num_blocks", [1, 2])
|
||||
def test_read_mcap_override_num_blocks(
|
||||
ray_start_regular_shared, tmp_path, override_num_blocks
|
||||
):
|
||||
"""Test override_num_blocks parameter."""
|
||||
path = os.path.join(tmp_path, "blocks_test.mcap")
|
||||
messages = [
|
||||
{
|
||||
"topic": "/test",
|
||||
"data": {"seq": i},
|
||||
"log_time": 1000000000 + i * 1000000,
|
||||
}
|
||||
for i in range(3)
|
||||
]
|
||||
create_test_mcap_file(path, messages)
|
||||
|
||||
ds = ray.data.read_mcap(path, override_num_blocks=override_num_blocks)
|
||||
|
||||
# Should still read all the data
|
||||
assert ds.count() == 3
|
||||
rows = ds.take_all()
|
||||
assert len(rows) == 3
|
||||
|
||||
|
||||
def test_read_mcap_file_extensions(ray_start_regular_shared, tmp_path):
|
||||
"""Test file extension filtering."""
|
||||
# Create MCAP file
|
||||
mcap_path = os.path.join(tmp_path, "data.mcap")
|
||||
messages = [
|
||||
{
|
||||
"topic": "/test",
|
||||
"data": {"test": "mcap_data"},
|
||||
"log_time": 1000000000,
|
||||
}
|
||||
]
|
||||
create_test_mcap_file(mcap_path, messages)
|
||||
|
||||
# Create non-MCAP file
|
||||
other_path = os.path.join(tmp_path, "data.txt")
|
||||
with open(other_path, "w") as f:
|
||||
f.write("not mcap data")
|
||||
|
||||
# Should only read .mcap files by default
|
||||
ds = ray.data.read_mcap(tmp_path)
|
||||
assert ds.count() == 1
|
||||
rows = ds.take_all()
|
||||
assert rows[0]["data"]["test"] == "mcap_data"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("ignore_missing_paths", [True, False])
|
||||
def test_read_mcap_ignore_missing_paths(
|
||||
ray_start_regular_shared, simple_mcap_file, ignore_missing_paths
|
||||
):
|
||||
"""Test ignore_missing_paths parameter."""
|
||||
paths = [simple_mcap_file, "/nonexistent/missing.mcap"]
|
||||
|
||||
if ignore_missing_paths:
|
||||
ds = ray.data.read_mcap(paths, ignore_missing_paths=ignore_missing_paths)
|
||||
assert ds.count() == 1
|
||||
assert ds.input_files() == [_unwrap_protocol(simple_mcap_file)]
|
||||
else:
|
||||
with pytest.raises(Exception): # FileNotFoundError or similar
|
||||
ds = ray.data.read_mcap(paths, ignore_missing_paths=ignore_missing_paths)
|
||||
ds.materialize()
|
||||
|
||||
|
||||
def test_read_mcap_json_decoding(ray_start_regular_shared, tmp_path):
|
||||
"""Test that JSON-encoded messages are properly decoded."""
|
||||
path = os.path.join(tmp_path, "json_test.mcap")
|
||||
|
||||
# Test data with nested JSON structure
|
||||
test_data = {
|
||||
"sensor_data": {
|
||||
"temperature": 23.5,
|
||||
"humidity": 45.0,
|
||||
"readings": [1, 2, 3, 4, 5],
|
||||
},
|
||||
"metadata": {"device_id": "sensor_001", "location": "room_a"},
|
||||
}
|
||||
|
||||
messages = [
|
||||
{
|
||||
"topic": "/sensor/data",
|
||||
"data": test_data,
|
||||
"log_time": 1000000000,
|
||||
}
|
||||
]
|
||||
|
||||
create_test_mcap_file(path, messages)
|
||||
assert os.path.exists(path), f"Test MCAP file was not created at {path}"
|
||||
|
||||
ds = ray.data.read_mcap(path)
|
||||
rows = ds.take_all()
|
||||
|
||||
assert len(rows) == 1, f"Expected 1 row, got {len(rows)}"
|
||||
row = rows[0]
|
||||
|
||||
# Verify the data field is properly decoded as a Python dict, not bytes
|
||||
assert isinstance(row["data"], dict), f"Expected dict, got {type(row['data'])}"
|
||||
assert row["data"]["sensor_data"]["temperature"] == 23.5
|
||||
assert row["data"]["metadata"]["device_id"] == "sensor_001"
|
||||
assert row["data"]["sensor_data"]["readings"] == [1, 2, 3, 4, 5]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,301 @@
|
||||
import shutil
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.data.tests.mock_http_server import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
# To run tests locally, make sure you install mongodb-org and have mongod
|
||||
# available on your PATH. Started directly since mongodb-org has no SysV init
|
||||
# script. See https://hub.docker.com/_/mongo
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def start_mongo():
|
||||
import pymongo
|
||||
import pymongo.errors
|
||||
|
||||
dbpath = tempfile.mkdtemp(prefix="mongod_test_")
|
||||
proc = subprocess.Popen(
|
||||
["mongod", "--dbpath", dbpath, "--bind_ip", "127.0.0.1"],
|
||||
stdout=subprocess.DEVNULL,
|
||||
stderr=subprocess.DEVNULL,
|
||||
)
|
||||
# Wait for mongod to accept connections.
|
||||
mongo_url = "mongodb://localhost:27017"
|
||||
for _ in range(30):
|
||||
if proc.poll() is not None:
|
||||
raise RuntimeError(
|
||||
f"mongod exited unexpectedly (returncode={proc.returncode})"
|
||||
)
|
||||
try:
|
||||
client = pymongo.MongoClient(mongo_url, serverSelectionTimeoutMS=1000)
|
||||
client.admin.command("ping")
|
||||
break
|
||||
except pymongo.errors.PyMongoError:
|
||||
time.sleep(0.5)
|
||||
else:
|
||||
proc.kill()
|
||||
raise RuntimeError("mongod failed to start")
|
||||
|
||||
# Make sure a clean slate for each test by dropping
|
||||
# previously created ones (if any).
|
||||
for db in client.list_database_names():
|
||||
# Keep the MongoDB default databases.
|
||||
if db not in ("admin", "local", "config"):
|
||||
client.drop_database(db)
|
||||
yield client, mongo_url
|
||||
|
||||
proc.terminate()
|
||||
proc.wait(timeout=10)
|
||||
shutil.rmtree(dbpath)
|
||||
|
||||
|
||||
def test_read_write_mongo(ray_start_regular_shared, start_mongo):
|
||||
from pymongo.errors import ServerSelectionTimeoutError
|
||||
from pymongoarrow.api import Schema
|
||||
|
||||
client, mongo_url = start_mongo
|
||||
foo_db = "foo-db"
|
||||
foo_collection = "foo-collection"
|
||||
foo = client[foo_db][foo_collection]
|
||||
foo.delete_many({})
|
||||
|
||||
# Read nonexistent URI.
|
||||
with pytest.raises(ServerSelectionTimeoutError):
|
||||
ds = ray.data.read_mongo(
|
||||
uri="nonexistent-uri",
|
||||
database=foo_db,
|
||||
collection=foo_collection,
|
||||
)
|
||||
# Read nonexistent database.
|
||||
with pytest.raises(ValueError):
|
||||
ds = ray.data.read_mongo(
|
||||
uri=mongo_url,
|
||||
database="nonexistent-db",
|
||||
collection=foo_collection,
|
||||
)
|
||||
# Read nonexistent collection.
|
||||
with pytest.raises(ValueError):
|
||||
ds = ray.data.read_mongo(
|
||||
uri=mongo_url,
|
||||
database=foo_db,
|
||||
collection="nonexistent-collection",
|
||||
)
|
||||
|
||||
# Inject 5 test docs.
|
||||
docs = [{"float_field": 2.0 * val, "int_field": val} for val in range(5)]
|
||||
df = pd.DataFrame(docs).astype({"int_field": "int32"})
|
||||
foo.insert_many(docs)
|
||||
|
||||
# Read a non-empty database, with schema specified.
|
||||
schema = Schema({"float_field": pa.float64(), "int_field": pa.int32()})
|
||||
ds = ray.data.read_mongo(
|
||||
uri=mongo_url,
|
||||
database=foo_db,
|
||||
collection=foo_collection,
|
||||
schema=schema,
|
||||
override_num_blocks=2,
|
||||
)
|
||||
assert ds._block_num_rows() == [3, 2]
|
||||
ds_schema = ds.schema()
|
||||
assert ds_schema.names == ["float_field", "int_field"]
|
||||
assert ds_schema.types == [pa.float64(), pa.int32()]
|
||||
result = ds.to_pandas()
|
||||
pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
|
||||
|
||||
# Read with schema inference, which will read all columns (including the auto
|
||||
# generated internal column "_id").
|
||||
ds = ray.data.read_mongo(
|
||||
uri=mongo_url,
|
||||
database=foo_db,
|
||||
collection=foo_collection,
|
||||
override_num_blocks=2,
|
||||
)
|
||||
assert ds._block_num_rows() == [3, 2]
|
||||
assert ds.count() == 5
|
||||
assert ds.schema().names == ["_id", "float_field", "int_field"]
|
||||
# We are not testing the datatype of _id here, because it varies per platform
|
||||
assert ds.schema().types[1:] == [
|
||||
pa.float64(),
|
||||
pa.int32(),
|
||||
]
|
||||
result = ds.drop_columns(["_id"]).to_pandas()
|
||||
pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
|
||||
|
||||
# Read a subset of the collection.
|
||||
ds = ray.data.read_mongo(
|
||||
uri=mongo_url,
|
||||
database=foo_db,
|
||||
collection=foo_collection,
|
||||
pipeline=[{"$match": {"int_field": {"$gte": 0, "$lt": 3}}}],
|
||||
override_num_blocks=2,
|
||||
)
|
||||
assert ds._block_num_rows() == [2, 1]
|
||||
assert ds.count() == 3
|
||||
assert ds.schema().names == ["_id", "float_field", "int_field"]
|
||||
df[df["int_field"] < 3].equals(ds.drop_columns(["_id"]).to_pandas())
|
||||
|
||||
# Read with auto-tuned parallelism.
|
||||
ds = ray.data.read_mongo(
|
||||
uri=mongo_url,
|
||||
database=foo_db,
|
||||
collection=foo_collection,
|
||||
)
|
||||
|
||||
assert ds.count() == 5
|
||||
assert ds.schema().names == ["_id", "float_field", "int_field"]
|
||||
# We are not testing the datatype of _id here, because it varies per platform
|
||||
assert ds.schema().types[1:] == [
|
||||
pa.float64(),
|
||||
pa.int32(),
|
||||
]
|
||||
result = ds.drop_columns(["_id"]).to_pandas()
|
||||
pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
|
||||
|
||||
# Read with a parallelism larger than number of rows.
|
||||
ds = ray.data.read_mongo(
|
||||
uri=mongo_url,
|
||||
database=foo_db,
|
||||
collection=foo_collection,
|
||||
override_num_blocks=1000,
|
||||
)
|
||||
|
||||
assert ds.count() == 5
|
||||
assert ds.schema().names == ["_id", "float_field", "int_field"]
|
||||
# We are not testing the datatype of _id here, because it varies per platform
|
||||
assert ds.schema().types[1:] == [
|
||||
pa.float64(),
|
||||
pa.int32(),
|
||||
]
|
||||
result = ds.drop_columns(["_id"]).to_pandas()
|
||||
pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
|
||||
|
||||
# Add a column and then write back to MongoDB.
|
||||
# Inject 2 more test docs.
|
||||
new_docs = [{"float_field": 2.0 * val, "int_field": val} for val in range(5, 7)]
|
||||
new_df = pd.DataFrame(new_docs).astype({"int_field": "int32"})
|
||||
ds2 = ray.data.from_pandas(new_df)
|
||||
ds2.write_mongo(uri=mongo_url, database=foo_db, collection=foo_collection)
|
||||
|
||||
# Read again to verify the content.
|
||||
expected_ds = ds.drop_columns(["_id"]).union(ds2)
|
||||
ds3 = ray.data.read_mongo(
|
||||
uri=mongo_url,
|
||||
database=foo_db,
|
||||
collection=foo_collection,
|
||||
)
|
||||
ds3.drop_columns(["_id"]).to_pandas().equals(expected_ds.to_pandas())
|
||||
|
||||
# Destination database doesn't exist.
|
||||
with pytest.raises(ValueError):
|
||||
ray.data.range(10).write_mongo(
|
||||
uri=mongo_url, database="nonexistent-db", collection=foo_collection
|
||||
)
|
||||
# Destination collection doesn't exist.
|
||||
with pytest.raises(ValueError):
|
||||
ray.data.range(10).write_mongo(
|
||||
uri=mongo_url, database=foo_db, collection="nonexistent-collection"
|
||||
)
|
||||
|
||||
|
||||
def test_mongo_datasource(ray_start_regular_shared, start_mongo):
|
||||
from pymongoarrow.api import Schema
|
||||
|
||||
client, mongo_url = start_mongo
|
||||
foo_db = "foo-db"
|
||||
foo_collection = "foo-collection"
|
||||
foo = client[foo_db][foo_collection]
|
||||
foo.delete_many({})
|
||||
|
||||
# Inject 5 test docs.
|
||||
docs = [{"float_field": 2.0 * key, "int_field": key} for key in range(5)]
|
||||
df = pd.DataFrame(docs).astype({"int_field": "int32"})
|
||||
foo.insert_many(docs)
|
||||
|
||||
# Read non-empty datasource with a specified schema.
|
||||
schema = Schema({"float_field": pa.float64(), "int_field": pa.int32()})
|
||||
ds = ray.data.read_mongo(
|
||||
uri=mongo_url,
|
||||
database=foo_db,
|
||||
collection=foo_collection,
|
||||
schema=schema,
|
||||
override_num_blocks=2,
|
||||
).materialize()
|
||||
assert ds._block_num_rows() == [3, 2]
|
||||
assert ds.num_blocks() == 2
|
||||
assert ds.count() == 5
|
||||
ds_schema = ds.schema()
|
||||
assert ds_schema.names == ["float_field", "int_field"]
|
||||
assert ds_schema.types == [pa.float64(), pa.int32()]
|
||||
result = ds.to_pandas()
|
||||
pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
|
||||
|
||||
# Read with schema inference, which will read all columns (including the auto
|
||||
# generated internal column "_id").
|
||||
ds = ray.data.read_mongo(
|
||||
uri=mongo_url,
|
||||
database=foo_db,
|
||||
collection=foo_collection,
|
||||
override_num_blocks=2,
|
||||
).materialize()
|
||||
assert ds._block_num_rows() == [3, 2]
|
||||
assert ds.num_blocks() == 2
|
||||
assert ds.count() == 5
|
||||
ds_schema = ds.schema()
|
||||
assert ds_schema.names == ["_id", "float_field", "int_field"]
|
||||
assert ds_schema.types[1:] == [pa.float64(), pa.int32()]
|
||||
result = ds.drop_columns(["_id"]).to_pandas()
|
||||
pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
|
||||
|
||||
# Read with auto-tuned parallelism.
|
||||
ds = ray.data.read_mongo(
|
||||
uri=mongo_url,
|
||||
database=foo_db,
|
||||
collection=foo_collection,
|
||||
).materialize()
|
||||
assert ds.num_blocks() == 2
|
||||
assert ds.count() == 5
|
||||
ds_schema = ds.schema()
|
||||
assert ds_schema.names == ["_id", "float_field", "int_field"]
|
||||
assert ds_schema.types[1:] == [pa.float64(), pa.int32()]
|
||||
result = ds.drop_columns(["_id"]).to_pandas()
|
||||
pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
|
||||
|
||||
# Read with a parallelism larger than number of rows.
|
||||
ds = ray.data.read_mongo(
|
||||
uri=mongo_url,
|
||||
database=foo_db,
|
||||
collection=foo_collection,
|
||||
override_num_blocks=1000,
|
||||
)
|
||||
assert ds.schema(fetch_if_missing=False) is None
|
||||
result = ds.drop_columns(["_id"]).to_pandas()
|
||||
pd.testing.assert_frame_equal(df.astype(result.dtypes.to_dict()), result)
|
||||
|
||||
# Read a subset of the collection.
|
||||
ds = ray.data.read_mongo(
|
||||
uri=mongo_url,
|
||||
database=foo_db,
|
||||
collection=foo_collection,
|
||||
pipeline=[{"$match": {"int_field": {"$gte": 0, "$lt": 3}}}],
|
||||
override_num_blocks=2,
|
||||
)
|
||||
assert ds._block_num_rows() == [2, 1]
|
||||
ds_schema = ds.schema()
|
||||
assert ds_schema.names == ["_id", "float_field", "int_field"]
|
||||
assert ds_schema.types[1:] == [pa.float64(), pa.int32()]
|
||||
df[df["int_field"] < 3].equals(ds.drop_columns(["_id"]).to_pandas())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,123 @@
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data.context import DataContext
|
||||
from ray.data.dataset import Schema
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.data.tests.util import extract_values
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
|
||||
@pytest.mark.parametrize("from_ref", [False, True])
|
||||
def test_from_numpy(ray_start_regular_shared, from_ref):
|
||||
arr1 = np.expand_dims(np.arange(0, 4), axis=1)
|
||||
arr2 = np.expand_dims(np.arange(4, 8), axis=1)
|
||||
arrs = [arr1, arr2]
|
||||
if from_ref:
|
||||
ds = ray.data.from_numpy_refs([ray.put(arr) for arr in arrs])
|
||||
else:
|
||||
ds = ray.data.from_numpy(arrs)
|
||||
values = np.stack(extract_values("data", ds.take(8)))
|
||||
np.testing.assert_array_equal(values, np.concatenate((arr1, arr2)))
|
||||
# Check that conversion task is included in stats.
|
||||
assert "FromNumpy" in ds.stats()
|
||||
|
||||
# Test from single NumPy ndarray.
|
||||
if from_ref:
|
||||
ds = ray.data.from_numpy_refs(ray.put(arr1))
|
||||
else:
|
||||
ds = ray.data.from_numpy(arr1)
|
||||
values = np.stack(extract_values("data", ds.take(4)))
|
||||
np.testing.assert_array_equal(values, arr1)
|
||||
# Check that conversion task is included in stats.
|
||||
assert "FromNumpy" in ds.stats()
|
||||
|
||||
|
||||
def test_from_numpy_variable_shaped(ray_start_regular_shared):
|
||||
arr = np.array([np.ones((2, 2)), np.ones((3, 3))], dtype=object)
|
||||
ds = ray.data.from_numpy(arr)
|
||||
values = np.array(extract_values("data", ds.take(2)), dtype=object)
|
||||
|
||||
def recursive_to_list(a):
|
||||
if not isinstance(a, (list, np.ndarray)):
|
||||
return a
|
||||
return [recursive_to_list(e) for e in a]
|
||||
|
||||
# Convert to a nested Python list in order to circumvent failed comparisons on
|
||||
# ndarray raggedness.
|
||||
np.testing.assert_equal(recursive_to_list(values), recursive_to_list(arr))
|
||||
|
||||
|
||||
def test_to_numpy_refs(ray_start_regular_shared):
|
||||
# Tensor Dataset
|
||||
ds = ray.data.range_tensor(10, override_num_blocks=2)
|
||||
arr = np.concatenate(extract_values("data", ray.get(ds.to_numpy_refs())))
|
||||
np.testing.assert_equal(arr, np.expand_dims(np.arange(0, 10), 1))
|
||||
|
||||
# Table Dataset
|
||||
ds = ray.data.range(10)
|
||||
arr = np.concatenate([t["id"] for t in ray.get(ds.to_numpy_refs())])
|
||||
np.testing.assert_equal(arr, np.arange(0, 10))
|
||||
|
||||
# Test multi-column Arrow dataset.
|
||||
ds = ray.data.from_arrow(pa.table({"a": [1, 2, 3], "b": [4, 5, 6]}))
|
||||
arrs = ray.get(ds.to_numpy_refs())
|
||||
np.testing.assert_equal(
|
||||
arrs, [{"a": np.array([1, 2, 3]), "b": np.array([4, 5, 6])}]
|
||||
)
|
||||
|
||||
# Test multi-column Pandas dataset.
|
||||
ds = ray.data.from_pandas(pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}))
|
||||
arrs = ray.get(ds.to_numpy_refs())
|
||||
np.testing.assert_equal(
|
||||
arrs, [{"a": np.array([1, 2, 3]), "b": np.array([4, 5, 6])}]
|
||||
)
|
||||
|
||||
|
||||
def test_numpy_roundtrip(ray_start_regular_shared, tmp_path):
|
||||
tensor_type = DataContext.get_current().arrow_fixed_shape_tensor_format.to_type()
|
||||
|
||||
ds = ray.data.range_tensor(10, override_num_blocks=2)
|
||||
ds.write_numpy(tmp_path, column="data")
|
||||
ds = ray.data.read_numpy(tmp_path)
|
||||
assert ds.count() == 10
|
||||
assert ds.schema() == Schema(pa.schema([("data", tensor_type((1,), pa.int64()))]))
|
||||
assert sorted(ds.take_all(), key=lambda row: row["data"]) == [
|
||||
{"data": np.array([i])} for i in range(10)
|
||||
]
|
||||
|
||||
|
||||
def test_numpy_read_x(ray_start_regular_shared, tmp_path):
|
||||
tensor_type = DataContext.get_current().arrow_fixed_shape_tensor_format.to_type()
|
||||
|
||||
path = os.path.join(tmp_path, "test_np_dir")
|
||||
os.mkdir(path)
|
||||
np.save(os.path.join(path, "test.npy"), np.expand_dims(np.arange(0, 10), 1))
|
||||
ds = ray.data.read_numpy(path, override_num_blocks=1)
|
||||
assert ds.count() == 10
|
||||
assert ds.schema() == Schema(pa.schema([("data", tensor_type((1,), pa.int64()))]))
|
||||
np.testing.assert_equal(
|
||||
extract_values("data", ds.take(2)), [np.array([0]), np.array([1])]
|
||||
)
|
||||
|
||||
|
||||
def test_numpy_write(ray_start_regular_shared, tmp_path):
|
||||
ds = ray.data.range_tensor(1)
|
||||
|
||||
ds.write_numpy(tmp_path, column="data")
|
||||
|
||||
actual_array = np.concatenate(
|
||||
[np.load(os.path.join(tmp_path, filename)) for filename in os.listdir(tmp_path)]
|
||||
)
|
||||
assert actual_array == np.array((0,))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,223 @@
|
||||
from typing import Iterator
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data._internal.execution.interfaces.ref_bundle import RefBundle
|
||||
from ray.data._internal.tensor_extensions.arrow import (
|
||||
ArrowTensorArray,
|
||||
get_arrow_extension_fixed_shape_tensor_types,
|
||||
)
|
||||
from ray.data.block import Block
|
||||
from ray.data.extensions import TensorDtype
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.data.tests.mock_http_server import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
from ray.types import ObjectRef
|
||||
|
||||
|
||||
def _get_first_block(bundles: Iterator[RefBundle]) -> ObjectRef[Block]:
|
||||
return next(bundles).block_refs[0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_pandas_block", [False, True])
|
||||
def test_from_pandas(ray_start_regular_shared, enable_pandas_block):
|
||||
ctx = ray.data.context.DataContext.get_current()
|
||||
old_enable_pandas_block = ctx.enable_pandas_block
|
||||
ctx.enable_pandas_block = enable_pandas_block
|
||||
try:
|
||||
df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
||||
df2 = pd.DataFrame({"one": [4, 5, 6], "two": ["e", "f", "g"]})
|
||||
ds = ray.data.from_pandas([df1, df2])
|
||||
block = ray.get(_get_first_block(ds.iter_internal_ref_bundles()))
|
||||
assert (
|
||||
isinstance(block, pd.DataFrame)
|
||||
if enable_pandas_block
|
||||
else isinstance(block, pa.Table)
|
||||
)
|
||||
values = [(r["one"], r["two"]) for r in ds.take(6)]
|
||||
rows = [(r.one, r.two) for _, r in pd.concat([df1, df2]).iterrows()]
|
||||
assert values == rows
|
||||
# Check that metadata fetch is included in stats.
|
||||
assert "FromPandas" in ds.stats()
|
||||
|
||||
# test from single pandas dataframe
|
||||
ds = ray.data.from_pandas(df1)
|
||||
block = ray.get(_get_first_block(ds.iter_internal_ref_bundles()))
|
||||
assert (
|
||||
isinstance(block, pd.DataFrame)
|
||||
if enable_pandas_block
|
||||
else isinstance(block, pa.Table)
|
||||
)
|
||||
values = [(r["one"], r["two"]) for r in ds.take(3)]
|
||||
rows = [(r.one, r.two) for _, r in df1.iterrows()]
|
||||
assert values == rows
|
||||
# Check that metadata fetch is included in stats.
|
||||
assert "FromPandas" in ds.stats()
|
||||
finally:
|
||||
ctx.enable_pandas_block = old_enable_pandas_block
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_inputs", [1, 2])
|
||||
def test_from_pandas_override_num_blocks(num_inputs, ray_start_regular_shared):
|
||||
df = pd.DataFrame({"number": [0]})
|
||||
|
||||
ds = ray.data.from_pandas([df] * num_inputs, override_num_blocks=2)
|
||||
|
||||
assert ds.materialize().num_blocks() == 2
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_pandas_block", [False, True])
|
||||
def test_from_pandas_refs(ray_start_regular_shared, enable_pandas_block):
|
||||
ctx = ray.data.context.DataContext.get_current()
|
||||
old_enable_pandas_block = ctx.enable_pandas_block
|
||||
ctx.enable_pandas_block = enable_pandas_block
|
||||
try:
|
||||
df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
||||
df2 = pd.DataFrame({"one": [4, 5, 6], "two": ["e", "f", "g"]})
|
||||
ds = ray.data.from_pandas_refs([ray.put(df1), ray.put(df2)])
|
||||
block = ray.get(_get_first_block(ds.iter_internal_ref_bundles()))
|
||||
assert (
|
||||
isinstance(block, pd.DataFrame)
|
||||
if enable_pandas_block
|
||||
else isinstance(block, pa.Table)
|
||||
)
|
||||
values = [(r["one"], r["two"]) for r in ds.take(6)]
|
||||
rows = [(r.one, r.two) for _, r in pd.concat([df1, df2]).iterrows()]
|
||||
assert values == rows
|
||||
# Check that metadata fetch is included in stats.
|
||||
assert "FromPandas" in ds.stats()
|
||||
|
||||
# test from single pandas dataframe ref
|
||||
ds = ray.data.from_pandas_refs(ray.put(df1))
|
||||
block = ray.get(_get_first_block(ds.iter_internal_ref_bundles()))
|
||||
assert (
|
||||
isinstance(block, pd.DataFrame)
|
||||
if enable_pandas_block
|
||||
else isinstance(block, pa.Table)
|
||||
)
|
||||
values = [(r["one"], r["two"]) for r in ds.take(3)]
|
||||
rows = [(r.one, r.two) for _, r in df1.iterrows()]
|
||||
assert values == rows
|
||||
# Check that metadata fetch is included in stats.
|
||||
assert "FromPandas" in ds.stats()
|
||||
finally:
|
||||
ctx.enable_pandas_block = old_enable_pandas_block
|
||||
|
||||
|
||||
def test_to_pandas(ray_start_regular_shared):
|
||||
n = 5
|
||||
df = pd.DataFrame({"id": list(range(n))})
|
||||
ds = ray.data.range(n)
|
||||
dfds = ds.to_pandas()
|
||||
pd.testing.assert_frame_equal(df.astype(dfds.dtypes.to_dict()), dfds)
|
||||
|
||||
# Test limit.
|
||||
with pytest.raises(ValueError):
|
||||
dfds = ds.to_pandas(limit=3)
|
||||
|
||||
# Test limit greater than number of rows.
|
||||
dfds = ds.to_pandas(limit=6)
|
||||
pd.testing.assert_frame_equal(df.astype(dfds.dtypes.to_dict()), dfds)
|
||||
|
||||
|
||||
def test_to_pandas_different_block_types(ray_start_regular_shared):
|
||||
# Test for https://github.com/ray-project/ray/issues/48575.
|
||||
df = pd.DataFrame({"a": [0]})
|
||||
ds1 = ray.data.from_pandas(df)
|
||||
|
||||
table = pa.Table.from_pandas(df)
|
||||
ds2 = ray.data.from_arrow(table)
|
||||
|
||||
actual_df = ds1.union(ds2).to_pandas()
|
||||
|
||||
expected_df = pd.DataFrame({"a": [0, 0]}).astype(actual_df.dtypes.to_dict())
|
||||
pd.testing.assert_frame_equal(actual_df, expected_df)
|
||||
|
||||
|
||||
def test_to_pandas_refs(ray_start_regular_shared):
|
||||
n = 5
|
||||
df = pd.DataFrame({"id": list(range(n))})
|
||||
ds = ray.data.range(n)
|
||||
dfds = pd.concat(ray.get(ds.to_pandas_refs()), ignore_index=True)
|
||||
pd.testing.assert_frame_equal(df.astype(dfds.dtypes.to_dict()), dfds)
|
||||
|
||||
|
||||
def test_pandas_roundtrip(ray_start_regular_shared, tmp_path):
|
||||
df1 = pd.DataFrame({"one": [1, 2, 3], "two": ["a", "b", "c"]})
|
||||
df2 = pd.DataFrame({"one": [4, 5, 6], "two": ["e", "f", "g"]})
|
||||
ds = ray.data.from_pandas([df1, df2], override_num_blocks=2)
|
||||
dfds = ds.to_pandas()
|
||||
expected = pd.concat([df1, df2], ignore_index=True)
|
||||
pd.testing.assert_frame_equal(expected.astype(dfds.dtypes.to_dict()), dfds)
|
||||
|
||||
|
||||
def test_to_pandas_tensor_column_cast_pandas(ray_start_regular_shared):
|
||||
# Check that tensor column casting occurs when converting a Dataset to a Pandas
|
||||
# DataFrame.
|
||||
data = np.arange(12).reshape((3, 2, 2))
|
||||
ctx = ray.data.context.DataContext.get_current()
|
||||
original = ctx.enable_tensor_extension_casting
|
||||
try:
|
||||
ctx.enable_tensor_extension_casting = True
|
||||
in_df = pd.DataFrame({"a": [data]})
|
||||
ds = ray.data.from_pandas(in_df)
|
||||
dtypes = ds.schema().base_schema.types
|
||||
assert len(dtypes) == 1
|
||||
# Tensor column should be automatically cast to Tensor extension.
|
||||
assert isinstance(dtypes[0], TensorDtype)
|
||||
# Original df should not be changed.
|
||||
assert not isinstance(in_df.dtypes[0], TensorDtype)
|
||||
out_df = ds.to_pandas()
|
||||
# Column should be cast back to object dtype when returning back to user.
|
||||
assert out_df["a"].dtype.type is np.object_
|
||||
expected_df = pd.DataFrame({"a": [data]})
|
||||
pd.testing.assert_frame_equal(out_df, expected_df)
|
||||
finally:
|
||||
ctx.enable_tensor_extension_casting = original
|
||||
|
||||
|
||||
def test_to_pandas_tensor_column_cast_arrow(ray_start_regular_shared):
|
||||
# Check that tensor column casting occurs when converting a Dataset to a Pandas
|
||||
# DataFrame.
|
||||
data = np.arange(12).reshape((3, 2, 2))
|
||||
ctx = ray.data.context.DataContext.get_current()
|
||||
original = ctx.enable_tensor_extension_casting
|
||||
try:
|
||||
ctx.enable_tensor_extension_casting = True
|
||||
in_table = pa.table({"a": ArrowTensorArray.from_numpy(data)})
|
||||
ds = ray.data.from_arrow(in_table)
|
||||
dtype = ds.schema().base_schema.field(0).type
|
||||
assert isinstance(dtype, get_arrow_extension_fixed_shape_tensor_types())
|
||||
out_df = ds.to_pandas()
|
||||
assert out_df["a"].dtype.type is np.object_
|
||||
expected_df = pd.DataFrame({"a": list(data)})
|
||||
pd.testing.assert_frame_equal(out_df, expected_df)
|
||||
finally:
|
||||
ctx.enable_tensor_extension_casting = original
|
||||
|
||||
|
||||
def test_read_pandas_data_array_column(ray_start_regular_shared):
|
||||
df = pd.DataFrame(
|
||||
{
|
||||
"one": [1, 2, 3],
|
||||
"array": [
|
||||
np.array([1, 1, 1]),
|
||||
np.array([2, 2, 2]),
|
||||
np.array([3, 3, 3]),
|
||||
],
|
||||
}
|
||||
)
|
||||
ds = ray.data.from_pandas(df)
|
||||
row = ds.take(1)[0]
|
||||
assert row["one"] == 1
|
||||
assert all(row["array"] == [1, 1, 1])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,71 @@
|
||||
import pandas
|
||||
import pytest
|
||||
import raydp
|
||||
|
||||
import ray
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.data.tests.test_util import _check_usage_record
|
||||
|
||||
|
||||
# RayDP tests require Ray Java. Make sure ray jar is built before running this test.
|
||||
@pytest.fixture(scope="function")
|
||||
def spark(request):
|
||||
ray.init(num_cpus=2, include_dashboard=False)
|
||||
spark_session = raydp.init_spark("test", 1, 1, "500M")
|
||||
|
||||
def stop_all():
|
||||
raydp.stop_spark()
|
||||
ray.shutdown()
|
||||
|
||||
request.addfinalizer(stop_all)
|
||||
return spark_session
|
||||
|
||||
|
||||
def test_raydp_roundtrip(spark):
|
||||
spark_df = spark.createDataFrame([(1, "a"), (2, "b"), (3, "c")], ["one", "two"])
|
||||
rows = [(r.one, r.two) for r in spark_df.take(3)]
|
||||
ds = ray.data.from_spark(spark_df)
|
||||
values = [(r["one"], r["two"]) for r in ds.take(6)]
|
||||
assert values == rows
|
||||
df = ds.to_spark(spark)
|
||||
rows_2 = [(r.one, r.two) for r in df.take(3)]
|
||||
assert values == rows_2
|
||||
|
||||
|
||||
def test_raydp_to_spark(spark):
|
||||
n = 5
|
||||
ds = ray.data.range(n)
|
||||
values = [r["id"] for r in ds.take(5)]
|
||||
df = ds.to_spark(spark)
|
||||
rows = [r.id for r in df.take(5)]
|
||||
assert values == rows
|
||||
|
||||
|
||||
def test_from_spark_e2e(spark):
|
||||
spark_df = spark.createDataFrame([(1, "a"), (2, "b"), (3, "c")], ["one", "two"])
|
||||
|
||||
rows = [(r.one, r.two) for r in spark_df.take(3)]
|
||||
ds = ray.data.from_spark(spark_df)
|
||||
assert len(ds.take_all()) == len(rows)
|
||||
values = [(r["one"], r["two"]) for r in ds.take(6)]
|
||||
assert values == rows
|
||||
|
||||
# Check that metadata fetch is included in stats.
|
||||
assert "FromArrow" in ds.stats()
|
||||
# Underlying implementation uses `FromArrow` operator
|
||||
assert ds._logical_plan.dag.name == "FromArrow"
|
||||
_check_usage_record(["FromArrow"])
|
||||
|
||||
|
||||
def test_to_pandas(spark):
|
||||
df = spark.range(100)
|
||||
ds = ray.data.from_spark(df)
|
||||
pdf = ds.to_pandas()
|
||||
pdf2 = df.toPandas().astype(pdf.dtypes.to_dict())
|
||||
pandas.testing.assert_frame_equal(pdf, pdf2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,214 @@
|
||||
"""Integration-ish tests for ``read_parquet()`` on the DataSourceV2 path.
|
||||
|
||||
These tests exercise planning-time behavior: schema inference,
|
||||
``ListFiles → ReadFiles`` attachment to the logical plan, and
|
||||
unsupported-option gating. They call ``ray.data.read_parquet`` which
|
||||
triggers Ray auto-init, so they live alongside the other datasource
|
||||
integration tests rather than under ``tests/unit/``.
|
||||
"""
|
||||
import pyarrow as pa
|
||||
import pyarrow.parquet as pq
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data._internal.datasource_v2.partitioners.round_robin_partitioner import (
|
||||
RoundRobinPartitioner,
|
||||
)
|
||||
from ray.data._internal.datasource_v2.scanners.parquet_scanner import ParquetScanner
|
||||
from ray.data._internal.logical.operators import ListFiles, ReadFiles
|
||||
from ray.data.context import DataContext
|
||||
|
||||
|
||||
def _write(path, table):
|
||||
pq.write_table(table, str(path))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def restore_ctx():
|
||||
ctx = DataContext.get_current()
|
||||
original = ctx.use_datasource_v2
|
||||
try:
|
||||
yield ctx
|
||||
finally:
|
||||
ctx.use_datasource_v2 = original
|
||||
|
||||
|
||||
def test_v2_flag_default():
|
||||
# The default is driven by ``DEFAULT_USE_DATASOURCE_V2``. Asserting
|
||||
# either direction here would be brittle, so just check that the
|
||||
# default is a bool.
|
||||
ctx = DataContext()
|
||||
assert isinstance(ctx.use_datasource_v2, bool)
|
||||
|
||||
|
||||
def test_read_parquet_builds_list_files_read_files_chain(tmp_path, restore_ctx):
|
||||
f = tmp_path / "data.parquet"
|
||||
_write(f, pa.table({"a": [1, 2, 3], "b": ["x", "y", "z"]}))
|
||||
|
||||
restore_ctx.use_datasource_v2 = True
|
||||
ds = ray.data.read_parquet(str(tmp_path))
|
||||
|
||||
assert isinstance(ds._logical_plan.dag, ReadFiles)
|
||||
assert isinstance(ds._logical_plan.dag.input_dependencies[0], ListFiles)
|
||||
schema = ds.schema()
|
||||
assert schema is not None
|
||||
assert "a" in schema.names
|
||||
assert "b" in schema.names
|
||||
|
||||
|
||||
def test_read_parquet_v2_hive_partitioned(tmp_path, restore_ctx):
|
||||
for p in ["a", "b"]:
|
||||
d = tmp_path / f"color={p}"
|
||||
d.mkdir()
|
||||
_write(d / "data.parquet", pa.table({"x": [1, 2]}))
|
||||
|
||||
restore_ctx.use_datasource_v2 = True
|
||||
ds = ray.data.read_parquet(str(tmp_path))
|
||||
schema = ds.schema()
|
||||
assert "x" in schema.names
|
||||
assert "color" in schema.names
|
||||
|
||||
|
||||
def test_read_parquet_v2_include_paths(tmp_path, restore_ctx):
|
||||
_write(tmp_path / "data.parquet", pa.table({"a": [1]}))
|
||||
|
||||
restore_ctx.use_datasource_v2 = True
|
||||
ds = ray.data.read_parquet(str(tmp_path), include_paths=True)
|
||||
schema = ds.schema()
|
||||
assert "path" in schema.names
|
||||
|
||||
|
||||
def test_read_parquet_v2_include_row_hash(tmp_path, restore_ctx):
|
||||
_write(tmp_path / "data.parquet", pa.table({"a": [1, 2, 3]}))
|
||||
|
||||
restore_ctx.use_datasource_v2 = True
|
||||
ds = ray.data.read_parquet(str(tmp_path), include_row_hash=True)
|
||||
schema = ds.schema()
|
||||
assert schema is not None
|
||||
assert "row_hash" in schema.names
|
||||
assert schema.types[schema.names.index("row_hash")] == pa.uint64()
|
||||
|
||||
|
||||
def test_read_parquet_v2_columns_applies_select_columns(tmp_path, restore_ctx):
|
||||
from ray.data._internal.logical.operators.map_operator import Project
|
||||
|
||||
_write(tmp_path / "data.parquet", pa.table({"a": [1], "b": [2]}))
|
||||
|
||||
restore_ctx.use_datasource_v2 = True
|
||||
with pytest.warns(DeprecationWarning, match="`columns=` on `read_parquet`"):
|
||||
ds = ray.data.read_parquet(str(tmp_path), columns=["a"])
|
||||
|
||||
# ``columns=`` is applied via ``ds.select_columns([...])``, which
|
||||
# wraps the ReadFiles op in a Project node.
|
||||
dag = ds._logical_plan.dag
|
||||
assert isinstance(dag, Project)
|
||||
assert [expr.name for expr in dag.exprs] == ["a"]
|
||||
assert isinstance(dag.input_dependencies[0], ReadFiles)
|
||||
|
||||
|
||||
def test_read_parquet_v2_columns_with_include_paths_preserves_path(
|
||||
tmp_path, restore_ctx
|
||||
):
|
||||
from ray.data._internal.logical.operators.map_operator import Project
|
||||
|
||||
_write(tmp_path / "data.parquet", pa.table({"a": [1], "b": [2]}))
|
||||
|
||||
restore_ctx.use_datasource_v2 = True
|
||||
with pytest.warns(DeprecationWarning, match="`columns=` on `read_parquet`"):
|
||||
ds = ray.data.read_parquet(str(tmp_path), columns=["a"], include_paths=True)
|
||||
|
||||
dag = ds._logical_plan.dag
|
||||
assert isinstance(dag, Project)
|
||||
# V1 ``columns=[...]`` retained ``"path"`` implicitly when
|
||||
# ``include_paths=True``; the V2 path appends it to keep that
|
||||
# behavior.
|
||||
assert [expr.name for expr in dag.exprs] == ["a", "path"]
|
||||
|
||||
|
||||
def test_read_parquet_v2_override_num_blocks_drives_partitioner(tmp_path, restore_ctx):
|
||||
_write(tmp_path / "data.parquet", pa.table({"a": [1, 2, 3]}))
|
||||
|
||||
restore_ctx.use_datasource_v2 = True
|
||||
original = restore_ctx.read_op_min_num_blocks
|
||||
ds = ray.data.read_parquet(str(tmp_path), override_num_blocks=7)
|
||||
|
||||
# The override should drive the ListFiles partitioner's bucket count
|
||||
# for this read only — the global DataContext must not be mutated.
|
||||
list_files_op = ds._logical_plan.dag.input_dependencies[0]
|
||||
assert isinstance(list_files_op, ListFiles)
|
||||
assert isinstance(list_files_op.file_partitioner, RoundRobinPartitioner)
|
||||
assert list_files_op.file_partitioner.num_buckets == 7
|
||||
assert restore_ctx.read_op_min_num_blocks == original
|
||||
|
||||
|
||||
def test_read_parquet_v2_filter_raises(tmp_path, restore_ctx):
|
||||
import pyarrow.dataset as pds
|
||||
|
||||
_write(tmp_path / "data.parquet", pa.table({"a": [1, 2, 3]}))
|
||||
|
||||
restore_ctx.use_datasource_v2 = True
|
||||
with pytest.raises(ValueError, match="`filter=` on `read_parquet`"):
|
||||
ray.data.read_parquet(str(tmp_path), filter=pds.field("a") > 1)
|
||||
|
||||
|
||||
def test_read_parquet_v2_dataset_kwargs_rejects_partitioning(tmp_path, restore_ctx):
|
||||
_write(tmp_path / "data.parquet", pa.table({"a": [1]}))
|
||||
|
||||
restore_ctx.use_datasource_v2 = True
|
||||
with pytest.warns(DeprecationWarning, match="`dataset_kwargs`"):
|
||||
with pytest.raises(
|
||||
ValueError, match="'partitioning' parameter isn't supported"
|
||||
):
|
||||
ray.data.read_parquet(
|
||||
str(tmp_path), dataset_kwargs={"partitioning": "hive"}
|
||||
)
|
||||
|
||||
|
||||
def test_read_parquet_v2_dataset_kwargs_rejects_filters(tmp_path, restore_ctx):
|
||||
_write(tmp_path / "data.parquet", pa.table({"a": [1]}))
|
||||
|
||||
restore_ctx.use_datasource_v2 = True
|
||||
with pytest.warns(DeprecationWarning, match="`dataset_kwargs`"):
|
||||
with pytest.raises(ValueError, match="Row filtering via 'filters'"):
|
||||
ray.data.read_parquet(
|
||||
str(tmp_path), dataset_kwargs={"filters": [("a", ">", 0)]}
|
||||
)
|
||||
|
||||
|
||||
def test_read_parquet_v2_dataset_kwargs_threads_through_to_scanner(
|
||||
tmp_path, restore_ctx
|
||||
):
|
||||
_write(tmp_path / "data.parquet", pa.table({"a": [1, 2, 3]}))
|
||||
|
||||
restore_ctx.use_datasource_v2 = True
|
||||
with pytest.warns(DeprecationWarning, match="`dataset_kwargs`"):
|
||||
ds = ray.data.read_parquet(
|
||||
str(tmp_path),
|
||||
dataset_kwargs={
|
||||
"coerce_int96_timestamp_unit": "ms",
|
||||
"read_dictionary": ["a"],
|
||||
},
|
||||
)
|
||||
|
||||
# ``read_dictionary`` is renamed to ``dictionary_columns`` to match
|
||||
# ``pds.ParquetFileFormat``; ``coerce_int96_timestamp_unit`` passes
|
||||
# through unchanged.
|
||||
read_files_op = ds._logical_plan.dag
|
||||
assert isinstance(read_files_op, ReadFiles)
|
||||
assert isinstance(read_files_op.scanner, ParquetScanner)
|
||||
assert read_files_op.scanner.parquet_format_kwargs == {
|
||||
"coerce_int96_timestamp_unit": "ms",
|
||||
"dictionary_columns": ["a"],
|
||||
}
|
||||
|
||||
|
||||
def test_read_parquet_v2_empty_dir_raises(tmp_path, restore_ctx):
|
||||
restore_ctx.use_datasource_v2 = True
|
||||
with pytest.raises(ValueError, match="no files found"):
|
||||
ray.data.read_parquet(str(tmp_path))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-xvs"]))
|
||||
@@ -0,0 +1,120 @@
|
||||
import base64
|
||||
import os
|
||||
import random
|
||||
import string
|
||||
from typing import Any, Dict, List, Tuple
|
||||
|
||||
import pytest
|
||||
from snowflake.connector import connect
|
||||
|
||||
import ray
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
# Note: Snowflake secrets are only used in postmerge authenticated tests.
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def connection_parameters():
|
||||
private_key_b64 = os.getenv("SNOWFLAKE_PRIVATE_KEY")
|
||||
private_key_bytes = base64.b64decode(private_key_b64)
|
||||
parameters = {
|
||||
"user": os.getenv("SNOWFLAKE_USER"),
|
||||
"account": os.getenv("SNOWFLAKE_ACCOUNT"),
|
||||
"database": os.getenv("SNOWFLAKE_DATABASE"),
|
||||
"schema": os.getenv("SNOWFLAKE_SCHEMA"),
|
||||
"warehouse": os.getenv("SNOWFLAKE_WAREHOUSE"),
|
||||
"private_key": private_key_bytes,
|
||||
}
|
||||
|
||||
yield parameters
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temp_table(connection_parameters):
|
||||
table_name = "".join([random.choice(string.ascii_uppercase) for _ in range(8)])
|
||||
|
||||
yield table_name
|
||||
|
||||
with connect(**connection_parameters) as connection, connection.cursor() as cursor:
|
||||
cursor.execute(f"DROP TABLE IF EXISTS {table_name}")
|
||||
connection.commit()
|
||||
|
||||
|
||||
@pytest.mark.needs_credentials
|
||||
def test_read(ray_start_regular_shared, connection_parameters):
|
||||
# This query fetches a small dataset with a variety of column types.
|
||||
query = "SELECT * FROM SNOWFLAKE_SAMPLE_DATA.TPCDS_SF100TCL.CALL_CENTER"
|
||||
|
||||
# Read the data and check contents.
|
||||
dataset = ray.data.read_snowflake(query, connection_parameters)
|
||||
actual_column_names = dataset.schema().names
|
||||
actual_rows = [tuple(row.values()) for row in dataset.take_all()]
|
||||
expected_column_names, expected_rows = execute(query, connection_parameters)
|
||||
|
||||
assert actual_column_names == expected_column_names
|
||||
assert sorted(actual_rows) == sorted(expected_rows)
|
||||
|
||||
|
||||
@pytest.mark.needs_credentials
|
||||
def test_write(ray_start_regular_shared, temp_table, connection_parameters):
|
||||
expected_column_names = ["title", "year", "score"]
|
||||
expected_rows = [
|
||||
("Monty Python and the Holy Grail", 1975, 8.2),
|
||||
("And Now for Something Completely Different", 1971, 7.5),
|
||||
]
|
||||
|
||||
# Create the table first
|
||||
create_table_sql = f"""
|
||||
CREATE TABLE IF NOT EXISTS {temp_table} (
|
||||
"title" VARCHAR(255),
|
||||
"year" INTEGER,
|
||||
"score" FLOAT
|
||||
)
|
||||
"""
|
||||
execute(create_table_sql, connection_parameters)
|
||||
|
||||
items = [dict(zip(expected_column_names, row)) for row in expected_rows]
|
||||
dataset = ray.data.from_items(items)
|
||||
|
||||
dataset.write_snowflake(temp_table, connection_parameters)
|
||||
actual_column_names, actual_rows = execute(
|
||||
f"SELECT * FROM {temp_table}", connection_parameters
|
||||
)
|
||||
|
||||
assert actual_column_names == expected_column_names
|
||||
assert sorted(actual_rows) == sorted(expected_rows)
|
||||
|
||||
|
||||
@pytest.mark.needs_credentials
|
||||
def execute(
|
||||
query: str, connection_parameters: Dict[str, str]
|
||||
) -> Tuple[List[str], List[Tuple[Any]]]:
|
||||
"""Execute a query on Snowflake and return the resulting data.
|
||||
|
||||
Args:
|
||||
query: The SQL query to execute.
|
||||
connection_parameters: Connection params for snowflake.
|
||||
|
||||
Returns:
|
||||
A two-tuple containing the column names and rows.
|
||||
"""
|
||||
with connect(**connection_parameters) as connection, connection.cursor() as cursor:
|
||||
cursor.execute(query)
|
||||
column_names = [column_metadata.name for column_metadata in cursor.description]
|
||||
rows = cursor.fetchall()
|
||||
|
||||
# TODO(mowen): Figure out how to actually handle the Decimal objects, we don't
|
||||
# want a divergenece in behavior here.
|
||||
# The Snowflake Python Connector represents numbers as `Decimal` objects.
|
||||
# rows = [
|
||||
# tuple(float(value) if isinstance(value, Decimal) else value for value in row)
|
||||
# for row in rows
|
||||
# ]
|
||||
|
||||
return column_names, rows
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,176 @@
|
||||
import sqlite3
|
||||
import tempfile
|
||||
from typing import Generator
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.tests.conftest import * # noqa # noqa
|
||||
|
||||
|
||||
@pytest.fixture(name="temp_database")
|
||||
def temp_database_fixture() -> Generator[str, None, None]:
|
||||
with tempfile.NamedTemporaryFile(suffix=".db") as file:
|
||||
yield file.name
|
||||
|
||||
|
||||
def test_read_sql(temp_database: str):
|
||||
connection = sqlite3.connect(temp_database)
|
||||
connection.execute("CREATE TABLE movie(title, year, score)")
|
||||
expected_values = [
|
||||
("Monty Python and the Holy Grail", 1975, 8.2),
|
||||
("And Now for Something Completely Different", 1971, 7.5),
|
||||
]
|
||||
connection.executemany("INSERT INTO movie VALUES (?, ?, ?)", expected_values)
|
||||
connection.commit()
|
||||
connection.close()
|
||||
|
||||
dataset = ray.data.read_sql(
|
||||
"SELECT * FROM movie",
|
||||
lambda: sqlite3.connect(temp_database),
|
||||
)
|
||||
actual_values = [tuple(record.values()) for record in dataset.take_all()]
|
||||
|
||||
assert sorted(actual_values) == sorted(expected_values)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"sql, sql_params",
|
||||
[
|
||||
("SELECT * FROM movie WHERE year >= ?", (1975,)),
|
||||
("SELECT * FROM movie WHERE year >= ?", [1975]),
|
||||
("SELECT * FROM movie WHERE year >= :year", {"year": 1975}),
|
||||
],
|
||||
)
|
||||
def test_read_sql_with_params(temp_database: str, sql: str, sql_params):
|
||||
connection = sqlite3.connect(temp_database)
|
||||
connection.execute("CREATE TABLE movie(title, year, score)")
|
||||
expected_values = [
|
||||
("Monty Python and the Holy Grail", 1975, 8.2),
|
||||
("And Now for Something Completely Different", 1971, 7.5),
|
||||
("Monty Python's Life of Brian", 1979, 8.0),
|
||||
]
|
||||
connection.executemany("INSERT INTO movie VALUES (?, ?, ?)", expected_values)
|
||||
connection.commit()
|
||||
connection.close()
|
||||
|
||||
dataset = ray.data.read_sql(
|
||||
sql,
|
||||
lambda: sqlite3.connect(temp_database),
|
||||
sql_params=sql_params,
|
||||
)
|
||||
actual_values = [tuple(record.values()) for record in dataset.take_all()]
|
||||
|
||||
assert sorted(actual_values) == sorted(
|
||||
[row for row in expected_values if row[1] >= 1975]
|
||||
)
|
||||
|
||||
|
||||
def test_read_sql_with_parallelism_fallback(temp_database: str):
|
||||
connection = sqlite3.connect(temp_database)
|
||||
connection.execute("CREATE TABLE grade(name, id, score)")
|
||||
base_tuple = ("xiaoming", 1, 8.2)
|
||||
# Generate 200 elements
|
||||
expected_values = [
|
||||
(f"{base_tuple[0]}{i}", i, base_tuple[2] + i + 1) for i in range(500)
|
||||
]
|
||||
connection.executemany("INSERT INTO grade VALUES (?, ?, ?)", expected_values)
|
||||
connection.commit()
|
||||
connection.close()
|
||||
|
||||
num_blocks = 2
|
||||
dataset = ray.data.read_sql(
|
||||
"SELECT * FROM grade",
|
||||
lambda: sqlite3.connect(temp_database),
|
||||
override_num_blocks=num_blocks,
|
||||
shard_hash_fn="unicode",
|
||||
shard_keys=["id"],
|
||||
)
|
||||
dataset = dataset.materialize()
|
||||
assert dataset.num_blocks() == num_blocks
|
||||
|
||||
actual_values = [tuple(record.values()) for record in dataset.take_all()]
|
||||
assert sorted(actual_values) == sorted(expected_values)
|
||||
|
||||
|
||||
# for mysql test
|
||||
@pytest.mark.skip(reason="skip this test because mysql env is not ready")
|
||||
def test_read_sql_with_parallelism_mysql(temp_database: str):
|
||||
# connect mysql
|
||||
import pymysql
|
||||
|
||||
connection = pymysql.connect(
|
||||
host="10.10.xx.xx", user="root", password="22222", database="test"
|
||||
)
|
||||
cursor = connection.cursor()
|
||||
|
||||
cursor.execute(
|
||||
"CREATE TABLE IF NOT EXISTS grade (name VARCHAR(255), id INT, score FLOAT)"
|
||||
)
|
||||
|
||||
base_tuple = ("xiaoming", 1, 8.2)
|
||||
expected_values = [
|
||||
(f"{base_tuple[0]}{i}", i, base_tuple[2] + i + 1) for i in range(200)
|
||||
]
|
||||
|
||||
cursor.executemany(
|
||||
"INSERT INTO grade (name, id, score) VALUES (%s, %s, %s)", expected_values
|
||||
)
|
||||
connection.commit()
|
||||
|
||||
cursor.close()
|
||||
connection.close()
|
||||
|
||||
dataset = ray.data.read_sql(
|
||||
"SELECT * FROM grade",
|
||||
lambda: pymysql.connect(host="xxxxx", user="xx", password="xx", database="xx"),
|
||||
parallelism=4,
|
||||
shard_keys=["id"],
|
||||
)
|
||||
actual_values = [tuple(record.values()) for record in dataset.take_all()]
|
||||
|
||||
assert sorted(actual_values) == sorted(expected_values)
|
||||
assert dataset.materialize().num_blocks() == 4
|
||||
|
||||
|
||||
def test_write_sql(temp_database: str):
|
||||
connection = sqlite3.connect(temp_database)
|
||||
connection.cursor().execute("CREATE TABLE test(string, number)")
|
||||
dataset = ray.data.from_items(
|
||||
[{"string": "spam", "number": 0}, {"string": "ham", "number": 1}]
|
||||
)
|
||||
|
||||
dataset.write_sql(
|
||||
"INSERT INTO test VALUES(?, ?)", lambda: sqlite3.connect(temp_database)
|
||||
)
|
||||
|
||||
result = connection.cursor().execute("SELECT * FROM test ORDER BY number")
|
||||
assert result.fetchall() == [("spam", 0), ("ham", 1)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_blocks", (1, 20))
|
||||
def test_write_sql_many_rows(num_blocks: int, temp_database: str):
|
||||
connection = sqlite3.connect(temp_database)
|
||||
connection.cursor().execute("CREATE TABLE test(id)")
|
||||
dataset = ray.data.range(1000).repartition(num_blocks)
|
||||
|
||||
dataset.write_sql(
|
||||
"INSERT INTO test VALUES(?)", lambda: sqlite3.connect(temp_database)
|
||||
)
|
||||
|
||||
result = connection.cursor().execute("SELECT * FROM test ORDER BY id")
|
||||
assert result.fetchall() == [(i,) for i in range(1000)]
|
||||
|
||||
|
||||
def test_write_sql_nonexistant_table(temp_database: str):
|
||||
dataset = ray.data.range(1)
|
||||
with pytest.raises(sqlite3.OperationalError):
|
||||
dataset.write_sql(
|
||||
"INSERT INTO test VALUES(?)", lambda: sqlite3.connect(temp_database)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,38 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.data.tests.test_util import _check_usage_record
|
||||
from ray.data.tests.util import extract_values
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
|
||||
def test_from_tf_e2e(ray_start_regular_shared_2_cpus):
|
||||
import tensorflow as tf
|
||||
import tensorflow_datasets as tfds
|
||||
|
||||
tf_dataset = tfds.load("mnist", split=["train"], as_supervised=True)[0]
|
||||
tf_dataset = tf_dataset.take(8) # Use subset to make test run faster.
|
||||
|
||||
ray_dataset = ray.data.from_tf(tf_dataset)
|
||||
|
||||
actual_data = extract_values("item", ray_dataset.take_all())
|
||||
expected_data = list(tf_dataset)
|
||||
assert len(actual_data) == len(expected_data)
|
||||
for (expected_features, expected_label), (actual_features, actual_label) in zip(
|
||||
expected_data, actual_data
|
||||
):
|
||||
tf.debugging.assert_equal(expected_features, actual_features)
|
||||
tf.debugging.assert_equal(expected_label, actual_label)
|
||||
|
||||
# Check that metadata fetch is included in stats.
|
||||
assert "FromItems" in ray_dataset.stats()
|
||||
# Underlying implementation uses `FromItems` operator
|
||||
assert ray_dataset._logical_plan.dag.name == "FromItems"
|
||||
_check_usage_record(["FromItems"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,100 @@
|
||||
import os
|
||||
|
||||
import pytest
|
||||
from fsspec.implementations.http import HTTPFileSystem
|
||||
|
||||
import ray
|
||||
from ray.data._internal.execution.interfaces.ref_bundle import (
|
||||
_ref_bundles_iterator_to_block_refs_list,
|
||||
)
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.data.tests.mock_http_server import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
|
||||
def _to_lines(rows):
|
||||
return [row["text"] for row in rows]
|
||||
|
||||
|
||||
def test_empty_text_files(ray_start_regular_shared, tmp_path):
|
||||
path = os.path.join(tmp_path, "test_text")
|
||||
os.mkdir(path)
|
||||
# 2 empty files.
|
||||
_ = open(os.path.join(path, "file1.txt"), "w")
|
||||
_ = open(os.path.join(path, "file2.txt"), "w")
|
||||
ds = ray.data.read_text(path)
|
||||
assert ds.count() == 0
|
||||
ds = ray.data.read_text(path, drop_empty_lines=False)
|
||||
assert ds.count() == 0
|
||||
|
||||
|
||||
def test_read_text(ray_start_regular_shared, tmp_path):
|
||||
path = os.path.join(tmp_path, "test_text")
|
||||
os.mkdir(path)
|
||||
with open(os.path.join(path, "file1.txt"), "w") as f:
|
||||
f.write("hello\n")
|
||||
f.write("world")
|
||||
with open(os.path.join(path, "file2.txt"), "w") as f:
|
||||
f.write("goodbye")
|
||||
with open(os.path.join(path, "file3.txt"), "w") as f:
|
||||
f.write("ray\n")
|
||||
ds = ray.data.read_text(path)
|
||||
assert sorted(_to_lines(ds.take())) == ["goodbye", "hello", "ray", "world"]
|
||||
ds = ray.data.read_text(path, drop_empty_lines=False)
|
||||
assert ds.count() == 4
|
||||
|
||||
|
||||
def test_read_text_remote_args(ray_start_cluster, tmp_path):
|
||||
cluster = ray_start_cluster
|
||||
cluster.add_node(
|
||||
resources={"foo": 100},
|
||||
num_cpus=1,
|
||||
_system_config={"max_direct_call_object_size": 0},
|
||||
)
|
||||
cluster.add_node(resources={"bar": 100}, num_cpus=1)
|
||||
|
||||
ray.shutdown()
|
||||
ray.init(cluster.address)
|
||||
|
||||
@ray.remote
|
||||
def get_node_id():
|
||||
return ray.get_runtime_context().get_node_id()
|
||||
|
||||
bar_node_id = ray.get(get_node_id.options(resources={"bar": 1}).remote())
|
||||
|
||||
path = os.path.join(tmp_path, "test_text")
|
||||
os.mkdir(path)
|
||||
with open(os.path.join(path, "file1.txt"), "w") as f:
|
||||
f.write("hello\n")
|
||||
f.write("world")
|
||||
with open(os.path.join(path, "file2.txt"), "w") as f:
|
||||
f.write("goodbye")
|
||||
|
||||
ds = ray.data.read_text(
|
||||
path, override_num_blocks=2, ray_remote_args={"resources": {"bar": 1}}
|
||||
)
|
||||
|
||||
block_refs = _ref_bundles_iterator_to_block_refs_list(
|
||||
ds.iter_internal_ref_bundles()
|
||||
)
|
||||
ray.wait(block_refs, num_returns=len(block_refs), fetch_local=False)
|
||||
location_data = ray.experimental.get_object_locations(block_refs)
|
||||
locations = []
|
||||
for block in block_refs:
|
||||
locations.extend(location_data[block]["node_ids"])
|
||||
assert set(locations) == {bar_node_id}, locations
|
||||
assert sorted(_to_lines(ds.take())) == ["goodbye", "hello", "world"]
|
||||
|
||||
|
||||
def test_fsspec_http_file_system(ray_start_regular_shared, http_server, http_file):
|
||||
ds = ray.data.read_text(http_file, filesystem=HTTPFileSystem())
|
||||
assert ds.count() > 0
|
||||
# Test auto-resolve of HTTP file system when it is not provided.
|
||||
ds = ray.data.read_text(http_file)
|
||||
assert ds.count() > 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,412 @@
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray import train
|
||||
from ray.data.preprocessors import Concatenator
|
||||
from ray.train import ScalingConfig
|
||||
|
||||
if sys.version_info <= (3, 12):
|
||||
# Skip this test for Python 3.12+ due to tensorflow incompatibility
|
||||
import tensorflow as tf
|
||||
|
||||
# if tf version is > 2.16, errors cannot be imported as functions
|
||||
# parse version with packaging
|
||||
from packaging import version
|
||||
|
||||
from ray.train.tensorflow import TensorflowTrainer
|
||||
|
||||
if version.parse(tf.__version__) >= version.parse("2.16"):
|
||||
mse = tf.keras.losses.MeanSquaredError()
|
||||
mae = tf.keras.losses.MeanAbsoluteError()
|
||||
else:
|
||||
mse = tf.keras.losses.mean_squared_error
|
||||
mae = tf.keras.losses.mean_absolute_error
|
||||
|
||||
|
||||
class TestToTF:
|
||||
def test_autosharding_is_disabled(self):
|
||||
ds = ray.data.from_items([{"spam": 0, "ham": 0}])
|
||||
|
||||
dataset = ds.to_tf(feature_columns="spam", label_columns="ham")
|
||||
|
||||
actual_auto_shard_policy = (
|
||||
dataset.options().experimental_distribute.auto_shard_policy
|
||||
)
|
||||
expected_auto_shard_policy = tf.data.experimental.AutoShardPolicy.OFF
|
||||
assert actual_auto_shard_policy is expected_auto_shard_policy
|
||||
|
||||
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
||||
def test_element_spec_type(self, include_additional_columns):
|
||||
ds = ray.data.from_items([{"spam": 0, "ham": 0, "weight": 0}])
|
||||
|
||||
if include_additional_columns:
|
||||
dataset = ds.to_tf(
|
||||
feature_columns="spam", label_columns="ham", additional_columns="weight"
|
||||
)
|
||||
feature_spec, label_spec, additional_spec = dataset.element_spec
|
||||
else:
|
||||
dataset = ds.to_tf(feature_columns="spam", label_columns="ham")
|
||||
feature_spec, label_spec = dataset.element_spec
|
||||
|
||||
assert isinstance(feature_spec, tf.TypeSpec)
|
||||
assert isinstance(label_spec, tf.TypeSpec)
|
||||
if include_additional_columns:
|
||||
assert isinstance(additional_spec, tf.TypeSpec)
|
||||
|
||||
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
||||
def test_element_spec_user_provided(self, include_additional_columns):
|
||||
ds = ray.data.from_items([{"spam": 0, "ham": 0, "eggs": 0, "weight": 0}])
|
||||
|
||||
if include_additional_columns:
|
||||
dataset1 = ds.to_tf(
|
||||
feature_columns=["spam", "ham"],
|
||||
label_columns="eggs",
|
||||
additional_columns="weight",
|
||||
)
|
||||
feature_spec, label_spec, additional_spec = dataset1.element_spec
|
||||
dataset2 = ds.to_tf(
|
||||
feature_columns=["spam", "ham"],
|
||||
label_columns="eggs",
|
||||
additional_columns="weight",
|
||||
feature_type_spec=feature_spec,
|
||||
label_type_spec=label_spec,
|
||||
additional_type_spec=additional_spec,
|
||||
)
|
||||
(
|
||||
feature_output_spec,
|
||||
label_output_spec,
|
||||
additional_output_spec,
|
||||
) = dataset2.element_spec
|
||||
else:
|
||||
dataset1 = ds.to_tf(feature_columns=["spam", "ham"], label_columns="eggs")
|
||||
feature_spec, label_spec = dataset1.element_spec
|
||||
dataset2 = ds.to_tf(
|
||||
feature_columns=["spam", "ham"],
|
||||
label_columns="eggs",
|
||||
feature_type_spec=feature_spec,
|
||||
label_type_spec=label_spec,
|
||||
)
|
||||
feature_output_spec, label_output_spec = dataset2.element_spec
|
||||
|
||||
assert isinstance(label_output_spec, tf.TypeSpec)
|
||||
assert isinstance(feature_output_spec, dict)
|
||||
assert feature_output_spec.keys() == {"spam", "ham"}
|
||||
assert all(
|
||||
isinstance(value, tf.TypeSpec) for value in feature_output_spec.values()
|
||||
)
|
||||
if include_additional_columns:
|
||||
assert isinstance(additional_output_spec, tf.TypeSpec)
|
||||
|
||||
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
||||
def test_element_spec_type_with_multiple_columns(self, include_additional_columns):
|
||||
ds = ray.data.from_items(
|
||||
[{"spam": 0, "ham": 0, "eggs": 0, "weight1": 0, "weight2": 0}]
|
||||
)
|
||||
|
||||
if include_additional_columns:
|
||||
dataset = ds.to_tf(
|
||||
feature_columns=["spam", "ham"],
|
||||
label_columns="eggs",
|
||||
additional_columns=["weight1", "weight2"],
|
||||
)
|
||||
(
|
||||
feature_output_signature,
|
||||
_,
|
||||
additional_output_signature,
|
||||
) = dataset.element_spec
|
||||
else:
|
||||
dataset = ds.to_tf(feature_columns=["spam", "ham"], label_columns="eggs")
|
||||
feature_output_signature, _ = dataset.element_spec
|
||||
|
||||
assert isinstance(feature_output_signature, dict)
|
||||
assert feature_output_signature.keys() == {"spam", "ham"}
|
||||
assert all(
|
||||
isinstance(value, tf.TypeSpec)
|
||||
for value in feature_output_signature.values()
|
||||
)
|
||||
|
||||
if include_additional_columns:
|
||||
assert isinstance(additional_output_signature, dict)
|
||||
assert additional_output_signature.keys() == {"weight1", "weight2"}
|
||||
assert all(
|
||||
isinstance(value, tf.TypeSpec)
|
||||
for value in additional_output_signature.values()
|
||||
)
|
||||
|
||||
df = pd.DataFrame(
|
||||
{
|
||||
"feature1": [0, 1, 2],
|
||||
"feature2": [3, 4, 5],
|
||||
"label": [0, 1, 1],
|
||||
"weight1": [0, 0.1, 0.2],
|
||||
"weight2": [0.3, 0.4, 0.5],
|
||||
}
|
||||
)
|
||||
ds = ray.data.from_pandas(df)
|
||||
|
||||
if include_additional_columns:
|
||||
dataset = ds.to_tf(
|
||||
feature_columns=["feature1", "feature2"],
|
||||
label_columns="label",
|
||||
additional_columns=["weight1", "weight2"],
|
||||
batch_size=3,
|
||||
)
|
||||
(
|
||||
feature_output_signature,
|
||||
_,
|
||||
additional_output_signature,
|
||||
) = dataset.element_spec
|
||||
assert isinstance(additional_output_signature, dict)
|
||||
assert additional_output_signature.keys() == {"weight1", "weight2"}
|
||||
assert all(
|
||||
isinstance(value, tf.TypeSpec)
|
||||
for value in additional_output_signature.values()
|
||||
)
|
||||
else:
|
||||
dataset = ds.to_tf(
|
||||
feature_columns=["feature1", "feature2"],
|
||||
label_columns="label",
|
||||
batch_size=3,
|
||||
)
|
||||
feature_output_signature, _ = dataset.element_spec
|
||||
|
||||
assert isinstance(feature_output_signature, dict)
|
||||
assert feature_output_signature.keys() == {"feature1", "feature2"}
|
||||
assert all(
|
||||
isinstance(value, tf.TypeSpec)
|
||||
for value in feature_output_signature.values()
|
||||
)
|
||||
|
||||
if include_additional_columns:
|
||||
features, labels, additional_metadata = next(iter(dataset))
|
||||
assert (
|
||||
additional_metadata["weight1"].numpy() == df["weight1"].values
|
||||
).all()
|
||||
assert (
|
||||
additional_metadata["weight2"].numpy() == df["weight2"].values
|
||||
).all()
|
||||
else:
|
||||
features, labels = next(iter(dataset))
|
||||
assert (labels.numpy() == df["label"].values).all()
|
||||
assert (features["feature1"].numpy() == df["feature1"].values).all()
|
||||
assert (features["feature2"].numpy() == df["feature2"].values).all()
|
||||
|
||||
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
||||
def test_element_spec_name(self, include_additional_columns):
|
||||
ds = ray.data.from_items([{"spam": 0, "ham": 0, "weight": 0}])
|
||||
|
||||
if include_additional_columns:
|
||||
dataset = ds.to_tf(
|
||||
feature_columns="spam", label_columns="ham", additional_columns="weight"
|
||||
)
|
||||
feature_spec, label_spec, additional_spec = dataset.element_spec
|
||||
else:
|
||||
dataset = ds.to_tf(feature_columns="spam", label_columns="ham")
|
||||
feature_spec, label_spec = dataset.element_spec
|
||||
|
||||
assert feature_spec.name == "spam"
|
||||
assert label_spec.name == "ham"
|
||||
if include_additional_columns:
|
||||
assert additional_spec.name == "weight"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data, expected_dtype",
|
||||
# Skip this test for Python 3.12+ due to tensorflow incompatibility
|
||||
[
|
||||
(0, tf.int64),
|
||||
(0.0, tf.double),
|
||||
(False, tf.bool),
|
||||
("eggs", tf.string),
|
||||
([1.0, 2.0], tf.float64),
|
||||
(np.zeros([2, 2], dtype=np.float32), tf.float32),
|
||||
]
|
||||
if sys.version_info <= (3, 12)
|
||||
else [],
|
||||
)
|
||||
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
||||
def test_element_spec_dtype(self, data, expected_dtype, include_additional_columns):
|
||||
ds = ray.data.from_items([{"spam": data, "ham": data, "weight": data}])
|
||||
|
||||
if include_additional_columns:
|
||||
dataset = ds.to_tf(
|
||||
feature_columns="spam",
|
||||
label_columns="ham",
|
||||
additional_columns="weight",
|
||||
)
|
||||
feature_spec, label_spec, additional_spec = dataset.element_spec
|
||||
else:
|
||||
dataset = ds.to_tf(feature_columns="spam", label_columns="ham")
|
||||
feature_spec, label_spec = dataset.element_spec
|
||||
|
||||
assert feature_spec.dtype == expected_dtype
|
||||
assert label_spec.dtype == expected_dtype
|
||||
if include_additional_columns:
|
||||
assert additional_spec.dtype == expected_dtype
|
||||
|
||||
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
||||
def test_element_spec_shape(self, include_additional_columns):
|
||||
ds = ray.data.from_items(8 * [{"spam": 0, "ham": 0, "weight": 0}])
|
||||
|
||||
if include_additional_columns:
|
||||
dataset = ds.to_tf(
|
||||
feature_columns="spam",
|
||||
label_columns="ham",
|
||||
additional_columns="weight",
|
||||
batch_size=4,
|
||||
)
|
||||
feature_spec, label_spec, additional_spec = dataset.element_spec
|
||||
assert tuple(additional_spec.shape) == (None,)
|
||||
else:
|
||||
dataset = ds.to_tf(
|
||||
feature_columns="spam", label_columns="ham", batch_size=4
|
||||
)
|
||||
feature_spec, label_spec = dataset.element_spec
|
||||
|
||||
assert tuple(feature_spec.shape) == (None,)
|
||||
assert tuple(label_spec.shape) == (None,)
|
||||
|
||||
if include_additional_columns:
|
||||
features, labels, additional_metadata = next(iter(dataset))
|
||||
assert tuple(additional_metadata.shape) == (4,)
|
||||
else:
|
||||
features, labels = next(iter(dataset))
|
||||
assert tuple(features.shape) == (4,)
|
||||
assert tuple(labels.shape) == (4,)
|
||||
|
||||
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
||||
def test_element_spec_shape_with_tensors(self, include_additional_columns):
|
||||
ds = ray.data.from_items(
|
||||
8
|
||||
* [
|
||||
{
|
||||
"spam": np.zeros([3, 32, 32]),
|
||||
"ham": 0,
|
||||
"weight": np.zeros([3, 32, 32]),
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
if include_additional_columns:
|
||||
dataset = ds.to_tf(
|
||||
feature_columns="spam",
|
||||
label_columns="ham",
|
||||
additional_columns="weight",
|
||||
batch_size=4,
|
||||
)
|
||||
feature_spec, _, additional_spec = dataset.element_spec
|
||||
assert tuple(additional_spec.shape) == (None, 3, 32, 32)
|
||||
else:
|
||||
dataset = ds.to_tf(
|
||||
feature_columns="spam", label_columns="ham", batch_size=4
|
||||
)
|
||||
feature_spec, _ = dataset.element_spec
|
||||
|
||||
assert tuple(feature_spec.shape) == (None, 3, 32, 32)
|
||||
|
||||
if include_additional_columns:
|
||||
features, labels, additional_metadata = next(iter(dataset))
|
||||
assert tuple(additional_metadata.shape) == (4, 3, 32, 32)
|
||||
else:
|
||||
features, labels = next(iter(dataset))
|
||||
assert tuple(features.shape) == (4, 3, 32, 32)
|
||||
assert tuple(labels.shape) == (4,)
|
||||
|
||||
@pytest.mark.parametrize("batch_size", [1, 2])
|
||||
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
||||
def test_element_spec_shape_with_ragged_tensors(
|
||||
self, batch_size, include_additional_columns
|
||||
):
|
||||
df = pd.DataFrame(
|
||||
{
|
||||
"spam": [np.zeros([32, 32, 3]), np.zeros([64, 64, 3])],
|
||||
"ham": [0, 0],
|
||||
"weight": [np.zeros([32, 32, 3]), np.zeros([64, 64, 3])],
|
||||
}
|
||||
)
|
||||
ds = ray.data.from_pandas(df)
|
||||
|
||||
if include_additional_columns:
|
||||
dataset = ds.to_tf(
|
||||
feature_columns="spam",
|
||||
label_columns="ham",
|
||||
additional_columns="weight",
|
||||
batch_size=batch_size,
|
||||
)
|
||||
feature_spec, _, additional_spec = dataset.element_spec
|
||||
assert tuple(additional_spec.shape) == (None, None, None, None)
|
||||
else:
|
||||
dataset = ds.to_tf(
|
||||
feature_columns="spam", label_columns="ham", batch_size=batch_size
|
||||
)
|
||||
feature_spec, _ = dataset.element_spec
|
||||
|
||||
assert tuple(feature_spec.shape) == (None, None, None, None)
|
||||
|
||||
if include_additional_columns:
|
||||
features, labels, additional_metadata = next(iter(dataset))
|
||||
assert tuple(additional_metadata.shape) == (batch_size, None, None, None)
|
||||
else:
|
||||
features, labels = next(iter(dataset))
|
||||
assert tuple(features.shape) == (batch_size, None, None, None)
|
||||
assert tuple(labels.shape) == (batch_size,)
|
||||
|
||||
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
||||
def test_training(self, include_additional_columns):
|
||||
def build_model() -> tf.keras.Model:
|
||||
return tf.keras.Sequential([tf.keras.layers.Dense(1)])
|
||||
|
||||
def train_func():
|
||||
strategy = tf.distribute.MultiWorkerMirroredStrategy()
|
||||
with strategy.scope():
|
||||
multi_worker_model = build_model()
|
||||
multi_worker_model.compile(
|
||||
optimizer=tf.keras.optimizers.SGD(),
|
||||
loss=mae,
|
||||
metrics=[mse],
|
||||
)
|
||||
|
||||
if include_additional_columns:
|
||||
dataset = train.get_dataset_shard("train").to_tf(
|
||||
"X", "Y", additional_columns="W", batch_size=4
|
||||
)
|
||||
else:
|
||||
dataset = train.get_dataset_shard("train").to_tf("X", "Y", batch_size=4)
|
||||
multi_worker_model.fit(dataset)
|
||||
|
||||
dataset = ray.data.from_items(8 * [{"X0": 0, "X1": 0, "Y": 0, "W": 0}])
|
||||
concatenator = Concatenator(columns=["X0", "X1"], output_column_name="X")
|
||||
dataset = concatenator.transform(dataset)
|
||||
|
||||
trainer = TensorflowTrainer(
|
||||
train_loop_per_worker=train_func,
|
||||
scaling_config=ScalingConfig(num_workers=2),
|
||||
datasets={"train": dataset},
|
||||
)
|
||||
trainer.fit()
|
||||
|
||||
@pytest.mark.parametrize("include_additional_columns", [False, True])
|
||||
def test_invalid_column_raises_error(self, include_additional_columns):
|
||||
ds = ray.data.from_items([{"spam": 0, "ham": 0, "weight": 0}])
|
||||
with pytest.raises(ValueError):
|
||||
if include_additional_columns:
|
||||
ds.to_tf(
|
||||
feature_columns="spam",
|
||||
label_columns="ham",
|
||||
additional_columns="baz",
|
||||
)
|
||||
else:
|
||||
ds.to_tf(feature_columns="foo", label_columns="bar")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
if sys.version_info >= (3, 12):
|
||||
# Skip this test for Python 3.12+ due to to incompatibility tensorflow
|
||||
sys.exit(0)
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,822 @@
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
from typing import TYPE_CHECKING
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from pandas.api.types import is_float_dtype, is_int64_dtype, is_object_dtype
|
||||
|
||||
import ray
|
||||
from ray.data.dataset import Dataset
|
||||
from ray.tests.conftest import * # noqa: F401,F403
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tensorflow_metadata.proto.v0 import schema_pb2
|
||||
|
||||
if sys.version_info <= (3, 12):
|
||||
# Skip this test for Python 3.12+ due to to incompatibility tensorflow
|
||||
import tensorflow as tf
|
||||
|
||||
|
||||
def _is_object_like(dtype):
|
||||
"""Match the pre-Arrow-dtype semantics of ``is_object_dtype``: pandas used
|
||||
object dtype for lists, bytes, and strings; ArrowBlockAccessor.to_pandas()
|
||||
now preserves these as ``pd.ArrowDtype`` via a ``types_mapper``."""
|
||||
if is_object_dtype(dtype):
|
||||
return True
|
||||
if isinstance(dtype, pd.ArrowDtype):
|
||||
pa_type = dtype.pyarrow_dtype
|
||||
return (
|
||||
pa.types.is_list(pa_type)
|
||||
or pa.types.is_large_list(pa_type)
|
||||
or pa.types.is_binary(pa_type)
|
||||
or pa.types.is_large_binary(pa_type)
|
||||
or pa.types.is_string(pa_type)
|
||||
or pa.types.is_large_string(pa_type)
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
def _is_int64_like(dtype):
|
||||
if is_int64_dtype(dtype):
|
||||
return True
|
||||
if isinstance(dtype, pd.ArrowDtype):
|
||||
return dtype.pyarrow_dtype == pa.int64()
|
||||
return False
|
||||
|
||||
|
||||
def _is_float_like(dtype):
|
||||
if is_float_dtype(dtype):
|
||||
return True
|
||||
if isinstance(dtype, pd.ArrowDtype):
|
||||
return pa.types.is_floating(dtype.pyarrow_dtype)
|
||||
return False
|
||||
|
||||
|
||||
def tf_records_partial():
|
||||
"""Underlying data corresponds to `data_partial` fixture."""
|
||||
import tensorflow as tf
|
||||
|
||||
return [
|
||||
# Record one (corresponding to row one).
|
||||
tf.train.Example(
|
||||
features=tf.train.Features(
|
||||
feature={
|
||||
"int_item": tf.train.Feature(
|
||||
int64_list=tf.train.Int64List(value=[1])
|
||||
),
|
||||
"int_list": tf.train.Feature(
|
||||
int64_list=tf.train.Int64List(value=[2, 2, 3])
|
||||
),
|
||||
"int_partial": tf.train.Feature(
|
||||
int64_list=tf.train.Int64List(value=[])
|
||||
),
|
||||
"float_item": tf.train.Feature(
|
||||
float_list=tf.train.FloatList(value=[1.0])
|
||||
),
|
||||
"float_list": tf.train.Feature(
|
||||
float_list=tf.train.FloatList(value=[2.0, 3.0, 4.0])
|
||||
),
|
||||
"float_partial": tf.train.Feature(
|
||||
float_list=tf.train.FloatList(value=[1.0])
|
||||
),
|
||||
"bytes_item": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[b"abc"])
|
||||
),
|
||||
"bytes_list": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[b"def", b"1234"])
|
||||
),
|
||||
"bytes_partial": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[])
|
||||
),
|
||||
"string_item": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[b"uvw"])
|
||||
),
|
||||
"string_list": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[b"xyz", b"999"])
|
||||
),
|
||||
"string_partial": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[])
|
||||
),
|
||||
}
|
||||
)
|
||||
),
|
||||
# Record two (corresponding to row two).
|
||||
tf.train.Example(
|
||||
features=tf.train.Features(
|
||||
feature={
|
||||
"int_item": tf.train.Feature(
|
||||
int64_list=tf.train.Int64List(value=[2])
|
||||
),
|
||||
"int_list": tf.train.Feature(
|
||||
int64_list=tf.train.Int64List(value=[3, 3, 4])
|
||||
),
|
||||
"int_partial": tf.train.Feature(
|
||||
int64_list=tf.train.Int64List(value=[9, 2])
|
||||
),
|
||||
"float_item": tf.train.Feature(
|
||||
float_list=tf.train.FloatList(value=[2.0])
|
||||
),
|
||||
"float_list": tf.train.Feature(
|
||||
float_list=tf.train.FloatList(value=[5.0, 6.0, 7.0])
|
||||
),
|
||||
"float_partial": tf.train.Feature(
|
||||
float_list=tf.train.FloatList(value=[])
|
||||
),
|
||||
"bytes_item": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[b"ghi"])
|
||||
),
|
||||
"bytes_list": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[b"jkl", b"5678"])
|
||||
),
|
||||
"bytes_partial": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[b"hello"])
|
||||
),
|
||||
"string_item": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[b"mno"])
|
||||
),
|
||||
"string_list": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[b"pqr", b"111"])
|
||||
),
|
||||
"string_partial": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[b"world"])
|
||||
),
|
||||
}
|
||||
)
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def data_partial(with_tf_schema):
|
||||
"""TFRecords generated from this corresponds to `tf_records_partial`."""
|
||||
return [
|
||||
# Row one.
|
||||
{
|
||||
"int_item": [1] if with_tf_schema else 1,
|
||||
"int_list": [2, 2, 3],
|
||||
"int_partial": [],
|
||||
"float_item": [1.0] if with_tf_schema else 1.0,
|
||||
"float_list": [2.0, 3.0, 4.0],
|
||||
"float_partial": [1.0] if with_tf_schema else 1.0,
|
||||
"bytes_item": [b"abc"] if with_tf_schema else b"abc",
|
||||
"bytes_list": [b"def", b"1234"],
|
||||
"bytes_partial": [] if with_tf_schema else None,
|
||||
"string_item": ["uvw"] if with_tf_schema else "uvw",
|
||||
"string_list": ["xyz", "999"],
|
||||
"string_partial": [] if with_tf_schema else None,
|
||||
},
|
||||
# Row two.
|
||||
{
|
||||
"int_item": [2] if with_tf_schema else 2,
|
||||
"int_list": [3, 3, 4],
|
||||
"int_partial": [9, 2],
|
||||
"float_item": [2.0] if with_tf_schema else 2.0,
|
||||
"float_list": [5.0, 6.0, 7.0],
|
||||
"float_partial": [] if with_tf_schema else None,
|
||||
"bytes_item": [b"ghi"] if with_tf_schema else b"ghi",
|
||||
"bytes_list": [b"jkl", b"5678"],
|
||||
"bytes_partial": [b"hello"] if with_tf_schema else b"hello",
|
||||
"string_item": ["mno"] if with_tf_schema else "mno",
|
||||
"string_list": ["pqr", "111"],
|
||||
"string_partial": ["world"] if with_tf_schema else "world",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def tf_records_empty():
|
||||
"""Underlying data corresponds to `data_empty` fixture."""
|
||||
import tensorflow as tf
|
||||
|
||||
return [
|
||||
# Record one (corresponding to row one).
|
||||
tf.train.Example(
|
||||
features=tf.train.Features(
|
||||
feature={
|
||||
"int_item": tf.train.Feature(
|
||||
int64_list=tf.train.Int64List(value=[1])
|
||||
),
|
||||
"int_list": tf.train.Feature(
|
||||
int64_list=tf.train.Int64List(value=[2, 2, 3])
|
||||
),
|
||||
"int_partial": tf.train.Feature(
|
||||
int64_list=tf.train.Int64List(value=[])
|
||||
),
|
||||
"int_empty": tf.train.Feature(
|
||||
int64_list=tf.train.Int64List(value=[])
|
||||
),
|
||||
"float_item": tf.train.Feature(
|
||||
float_list=tf.train.FloatList(value=[1.0])
|
||||
),
|
||||
"float_list": tf.train.Feature(
|
||||
float_list=tf.train.FloatList(value=[2.0, 3.0, 4.0])
|
||||
),
|
||||
"float_partial": tf.train.Feature(
|
||||
float_list=tf.train.FloatList(value=[1.0])
|
||||
),
|
||||
"float_empty": tf.train.Feature(
|
||||
float_list=tf.train.FloatList(value=[])
|
||||
),
|
||||
"bytes_item": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[b"abc"])
|
||||
),
|
||||
"bytes_list": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[b"def", b"1234"])
|
||||
),
|
||||
"bytes_partial": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[])
|
||||
),
|
||||
"bytes_empty": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[])
|
||||
),
|
||||
"string_item": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[b"uvw"])
|
||||
),
|
||||
"string_list": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[b"xyz", b"999"])
|
||||
),
|
||||
"string_partial": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[])
|
||||
),
|
||||
"string_empty": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[])
|
||||
),
|
||||
}
|
||||
)
|
||||
),
|
||||
# Record two (corresponding to row two).
|
||||
tf.train.Example(
|
||||
features=tf.train.Features(
|
||||
feature={
|
||||
"int_item": tf.train.Feature(
|
||||
int64_list=tf.train.Int64List(value=[2])
|
||||
),
|
||||
"int_list": tf.train.Feature(
|
||||
int64_list=tf.train.Int64List(value=[3, 3, 4])
|
||||
),
|
||||
"int_partial": tf.train.Feature(
|
||||
int64_list=tf.train.Int64List(value=[9, 2])
|
||||
),
|
||||
"int_empty": tf.train.Feature(
|
||||
int64_list=tf.train.Int64List(value=[])
|
||||
),
|
||||
"float_item": tf.train.Feature(
|
||||
float_list=tf.train.FloatList(value=[2.0])
|
||||
),
|
||||
"float_list": tf.train.Feature(
|
||||
float_list=tf.train.FloatList(value=[5.0, 6.0, 7.0])
|
||||
),
|
||||
"float_partial": tf.train.Feature(
|
||||
float_list=tf.train.FloatList(value=[])
|
||||
),
|
||||
"float_empty": tf.train.Feature(
|
||||
float_list=tf.train.FloatList(value=[])
|
||||
),
|
||||
"bytes_item": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[b"ghi"])
|
||||
),
|
||||
"bytes_list": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[b"jkl", b"5678"])
|
||||
),
|
||||
"bytes_partial": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[b"hello"])
|
||||
),
|
||||
"bytes_empty": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[])
|
||||
),
|
||||
"string_item": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[b"mno"])
|
||||
),
|
||||
"string_list": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[b"pqr", b"111"])
|
||||
),
|
||||
"string_partial": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[b"world"])
|
||||
),
|
||||
"string_empty": tf.train.Feature(
|
||||
bytes_list=tf.train.BytesList(value=[])
|
||||
),
|
||||
}
|
||||
)
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def data_empty(with_tf_schema):
|
||||
"""TFRecords generated from this corresponds to
|
||||
the `tf_records_empty` fixture."""
|
||||
return [
|
||||
# Row one.
|
||||
{
|
||||
"int_item": [1] if with_tf_schema else 1,
|
||||
"int_list": [2, 2, 3],
|
||||
"int_partial": [],
|
||||
"int_empty": [],
|
||||
"float_item": [1.0] if with_tf_schema else 1.0,
|
||||
"float_list": [2.0, 3.0, 4.0],
|
||||
"float_partial": [1.0] if with_tf_schema else 1.0,
|
||||
"float_empty": [],
|
||||
"bytes_item": [b"abc"] if with_tf_schema else b"abc",
|
||||
"bytes_list": [b"def", b"1234"],
|
||||
"bytes_partial": [],
|
||||
"bytes_empty": [],
|
||||
"string_item": ["uvw"] if with_tf_schema else "uvw",
|
||||
"string_list": ["xyz", "999"],
|
||||
"string_partial": [] if with_tf_schema else None,
|
||||
"string_empty": [],
|
||||
},
|
||||
# Row two.
|
||||
{
|
||||
"int_item": [2] if with_tf_schema else 2,
|
||||
"int_list": [3, 3, 4],
|
||||
"int_partial": [9, 2],
|
||||
"int_empty": [],
|
||||
"float_item": [2.0] if with_tf_schema else 2.0,
|
||||
"float_list": [5.0, 6.0, 7.0],
|
||||
"float_partial": [],
|
||||
"float_empty": [],
|
||||
"bytes_item": [b"ghi"] if with_tf_schema else b"ghi",
|
||||
"bytes_list": [b"jkl", b"5678"],
|
||||
"bytes_partial": [b"hello"] if with_tf_schema else b"hello",
|
||||
"bytes_empty": [],
|
||||
"string_item": ["mno"] if with_tf_schema else "mno",
|
||||
"string_list": ["pqr", "111"],
|
||||
"string_partial": ["world"] if with_tf_schema else "world",
|
||||
"string_empty": [],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def _features_to_schema(features: "tf.train.Features") -> "schema_pb2.Schema":
|
||||
from tensorflow_metadata.proto.v0 import schema_pb2
|
||||
|
||||
tf_schema = schema_pb2.Schema()
|
||||
for feature_name, feature_msg in features.feature.items():
|
||||
schema_feature = tf_schema.feature.add()
|
||||
schema_feature.name = feature_name
|
||||
if feature_msg.HasField("bytes_list"):
|
||||
schema_feature.type = schema_pb2.FeatureType.BYTES
|
||||
elif feature_msg.HasField("float_list"):
|
||||
schema_feature.type = schema_pb2.FeatureType.FLOAT
|
||||
elif feature_msg.HasField("int64_list"):
|
||||
schema_feature.type = schema_pb2.FeatureType.INT
|
||||
return tf_schema
|
||||
|
||||
|
||||
def _ds_eq_streaming(ds_expected, ds_actual) -> bool:
|
||||
# Casting the strings to bytes for comparing string features
|
||||
def _str2bytes(d):
|
||||
for k, v in d.items():
|
||||
if "string" in k:
|
||||
if isinstance(v, list):
|
||||
d[k] = [vv.encode() for vv in v]
|
||||
elif isinstance(v, str):
|
||||
d[k] = v.encode()
|
||||
return d
|
||||
|
||||
ds_expected = ds_expected.map(_str2bytes)
|
||||
assert ds_expected.take() == ds_actual.take()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"with_tf_schema,compression",
|
||||
[
|
||||
(True, None),
|
||||
(False, None),
|
||||
],
|
||||
)
|
||||
def test_read_tfrecords(
|
||||
with_tf_schema,
|
||||
compression,
|
||||
ray_start_regular_shared_2_cpus,
|
||||
tmp_path,
|
||||
):
|
||||
import pandas as pd
|
||||
import tensorflow as tf
|
||||
|
||||
example = tf_records_empty()[0]
|
||||
|
||||
tf_schema = None
|
||||
if with_tf_schema:
|
||||
tf_schema = _features_to_schema(example.features)
|
||||
|
||||
path = os.path.join(tmp_path, "data.tfrecords")
|
||||
with tf.io.TFRecordWriter(
|
||||
path=path, options=tf.io.TFRecordOptions(compression_type=compression)
|
||||
) as writer:
|
||||
writer.write(example.SerializeToString())
|
||||
|
||||
arrow_open_stream_args = None
|
||||
if compression:
|
||||
arrow_open_stream_args = {"compression": compression}
|
||||
|
||||
ds = ray.data.read_tfrecords(
|
||||
path,
|
||||
tf_schema=tf_schema,
|
||||
arrow_open_stream_args=arrow_open_stream_args,
|
||||
)
|
||||
|
||||
df = ds.to_pandas()
|
||||
# Protobuf serializes features in a non-deterministic order.
|
||||
if with_tf_schema:
|
||||
assert _is_object_like(dict(df.dtypes)["int_item"])
|
||||
else:
|
||||
assert _is_int64_like(dict(df.dtypes)["int_item"])
|
||||
assert _is_object_like(dict(df.dtypes)["int_list"])
|
||||
assert _is_object_like(dict(df.dtypes)["int_partial"])
|
||||
assert _is_object_like(dict(df.dtypes)["int_empty"])
|
||||
|
||||
if with_tf_schema:
|
||||
assert _is_object_like(dict(df.dtypes)["float_item"])
|
||||
assert _is_object_like(dict(df.dtypes)["float_partial"])
|
||||
else:
|
||||
assert _is_float_like(dict(df.dtypes)["float_item"])
|
||||
assert _is_float_like(dict(df.dtypes)["float_partial"])
|
||||
assert _is_object_like(dict(df.dtypes)["float_list"])
|
||||
assert _is_object_like(dict(df.dtypes)["float_empty"])
|
||||
|
||||
assert _is_object_like(dict(df.dtypes)["bytes_item"])
|
||||
assert _is_object_like(dict(df.dtypes)["bytes_partial"])
|
||||
assert _is_object_like(dict(df.dtypes)["bytes_list"])
|
||||
assert _is_object_like(dict(df.dtypes)["bytes_empty"])
|
||||
|
||||
assert _is_object_like(dict(df.dtypes)["string_item"])
|
||||
assert _is_object_like(dict(df.dtypes)["string_partial"])
|
||||
assert _is_object_like(dict(df.dtypes)["string_list"])
|
||||
assert _is_object_like(dict(df.dtypes)["string_empty"])
|
||||
|
||||
# If the schema is specified, we should not perform the
|
||||
# automatic unwrapping of single-element lists.
|
||||
if with_tf_schema:
|
||||
assert isinstance(df["int_item"], pd.Series)
|
||||
assert df["int_item"].tolist() == [[1]]
|
||||
else:
|
||||
assert list(df["int_item"]) == [1]
|
||||
assert np.array_equal(df["int_list"][0], np.array([2, 2, 3]))
|
||||
assert np.array_equal(df["int_partial"][0], np.array([], dtype=np.int64))
|
||||
assert np.array_equal(df["int_empty"][0], np.array([], dtype=np.int64))
|
||||
|
||||
if with_tf_schema:
|
||||
assert isinstance(df["float_item"], pd.Series)
|
||||
assert df["float_item"].tolist() == [[1.0]]
|
||||
assert df["float_partial"].tolist() == [[1.0]]
|
||||
else:
|
||||
assert list(df["float_item"]) == [1.0]
|
||||
assert list(df["float_partial"]) == [1.0]
|
||||
assert np.array_equal(df["float_list"][0], np.array([2.0, 3.0, 4.0]))
|
||||
assert np.array_equal(df["float_empty"][0], np.array([], dtype=np.float32))
|
||||
|
||||
if with_tf_schema:
|
||||
assert isinstance(df["bytes_item"], pd.Series)
|
||||
assert df["bytes_item"].tolist() == [[b"abc"]]
|
||||
assert isinstance(df["string_item"], pd.Series)
|
||||
assert df["string_item"].tolist() == [[b"uvw"]] # strings are read as bytes
|
||||
else:
|
||||
assert list(df["bytes_item"]) == [b"abc"]
|
||||
assert list(df["string_item"]) == [b"uvw"]
|
||||
assert np.array_equal(df["bytes_list"][0], np.array([b"def", b"1234"]))
|
||||
assert np.array_equal(df["bytes_partial"][0], np.array([], dtype=np.bytes_))
|
||||
assert np.array_equal(df["bytes_empty"][0], np.array([], dtype=np.bytes_))
|
||||
|
||||
assert np.array_equal(df["string_list"][0], np.array([b"xyz", b"999"]))
|
||||
assert np.array_equal(df["string_partial"][0], np.array([], dtype=np.bytes_))
|
||||
assert np.array_equal(df["string_empty"][0], np.array([], dtype=np.bytes_))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_ray_data_read_tfrecords(mocker):
|
||||
mock_read_tfrecords = mocker.patch("ray.data.read_tfrecords")
|
||||
mock_read_tfrecords.return_value = MagicMock(spec=Dataset)
|
||||
return mock_read_tfrecords
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_cpus", [1, 2, 4])
|
||||
def test_read_tfrecords_ray_remote_args(
|
||||
ray_start_regular_shared_2_cpus,
|
||||
mock_ray_data_read_tfrecords,
|
||||
tmp_path,
|
||||
num_cpus,
|
||||
):
|
||||
import tensorflow as tf
|
||||
|
||||
example = tf_records_empty()[0]
|
||||
path = os.path.join(tmp_path, "data.tfrecords")
|
||||
with tf.io.TFRecordWriter(path=path) as writer:
|
||||
writer.write(example.SerializeToString())
|
||||
ray_remote_args = {"num_cpus": num_cpus}
|
||||
ds = ray.data.read_tfrecords(
|
||||
paths=[path],
|
||||
ray_remote_args=ray_remote_args,
|
||||
)
|
||||
assert isinstance(ds, Dataset)
|
||||
mock_ray_data_read_tfrecords.assert_called_once()
|
||||
args, kwargs = mock_ray_data_read_tfrecords.call_args
|
||||
assert kwargs["paths"] == [path]
|
||||
assert kwargs["ray_remote_args"] == ray_remote_args
|
||||
|
||||
|
||||
@pytest.mark.parametrize("with_tf_schema", (True, False))
|
||||
def test_write_tfrecords(
|
||||
with_tf_schema,
|
||||
ray_start_regular_shared_2_cpus,
|
||||
tmp_path,
|
||||
):
|
||||
"""Test that write_tfrecords writes TFRecords correctly.
|
||||
|
||||
Test this by writing a Dataset to a TFRecord (function under test),
|
||||
reading it back out into a tf.train.Example,
|
||||
and checking that the result is analogous to the original Dataset.
|
||||
"""
|
||||
|
||||
import tensorflow as tf
|
||||
|
||||
# The dataset we will write to a .tfrecords file.
|
||||
ds = ray.data.from_items(
|
||||
data_partial(with_tf_schema),
|
||||
# Here, we specify `override_num_blocks=1` to ensure that all rows end up in
|
||||
# the same block, which is required for type inference involving partially
|
||||
# missing columns.
|
||||
override_num_blocks=1,
|
||||
)
|
||||
|
||||
# The corresponding tf.train.Example that we would expect to read
|
||||
# from this dataset.
|
||||
expected_records = tf_records_partial()
|
||||
|
||||
tf_schema = None
|
||||
if with_tf_schema:
|
||||
features = expected_records[0].features
|
||||
tf_schema = _features_to_schema(features)
|
||||
|
||||
# Perform the test.
|
||||
# Write the dataset to a .tfrecords file.
|
||||
ds.write_tfrecords(tmp_path, tf_schema=tf_schema)
|
||||
|
||||
# Read the Examples back out from the .tfrecords file.
|
||||
# This follows the offical TFRecords tutorial:
|
||||
# https://www.tensorflow.org/tutorials/load_data/tfrecord#reading_a_tfrecord_file_2
|
||||
|
||||
filenames = sorted(os.listdir(tmp_path))
|
||||
filepaths = [os.path.join(tmp_path, filename) for filename in filenames]
|
||||
raw_dataset = tf.data.TFRecordDataset(filepaths)
|
||||
|
||||
tfrecords = []
|
||||
for raw_record in raw_dataset:
|
||||
example = tf.train.Example()
|
||||
example.ParseFromString(raw_record.numpy())
|
||||
tfrecords.append(example)
|
||||
|
||||
assert tfrecords == expected_records
|
||||
|
||||
|
||||
@pytest.mark.parametrize("with_tf_schema", (True, False))
|
||||
def test_write_tfrecords_empty_features(
|
||||
with_tf_schema,
|
||||
ray_start_regular_shared_2_cpus,
|
||||
tmp_path,
|
||||
):
|
||||
"""Test that write_tfrecords writes TFRecords with completely empty features
|
||||
correctly (i.e. the case where type inference from partially filled features
|
||||
is not possible). We expect this to succeed when passing an explicit `tf_schema`
|
||||
param, and otherwise will raise a `ValueError`.
|
||||
|
||||
Test this by writing a Dataset to a TFRecord (function under test),
|
||||
reading it back out into a tf.train.Example,
|
||||
and checking that the result is analogous to the original Dataset.
|
||||
"""
|
||||
|
||||
import tensorflow as tf
|
||||
|
||||
# The dataset we will write to a .tfrecords file.
|
||||
ds = ray.data.from_items(data_empty(with_tf_schema))
|
||||
|
||||
# The corresponding tf.train.Example that we would expect to read
|
||||
# from this dataset.
|
||||
expected_records = tf_records_empty()
|
||||
|
||||
if not with_tf_schema:
|
||||
with pytest.raises(ValueError):
|
||||
# Type inference from fully empty columns should fail if
|
||||
# no schema is specified.
|
||||
ds.write_tfrecords(tmp_path)
|
||||
else:
|
||||
features = expected_records[0].features
|
||||
tf_schema = _features_to_schema(features)
|
||||
|
||||
# Perform the test.
|
||||
# Write the dataset to a .tfrecords file.
|
||||
ds.write_tfrecords(tmp_path, tf_schema=tf_schema)
|
||||
|
||||
# Read the Examples back out from the .tfrecords file.
|
||||
# This follows the offical TFRecords tutorial:
|
||||
# https://www.tensorflow.org/tutorials/load_data/tfrecord#reading_a_tfrecord_file_2
|
||||
|
||||
filenames = sorted(os.listdir(tmp_path))
|
||||
filepaths = [os.path.join(tmp_path, filename) for filename in filenames]
|
||||
raw_dataset = tf.data.TFRecordDataset(filepaths)
|
||||
|
||||
tfrecords = []
|
||||
for raw_record in raw_dataset:
|
||||
example = tf.train.Example()
|
||||
example.ParseFromString(raw_record.numpy())
|
||||
tfrecords.append(example)
|
||||
|
||||
assert tfrecords == expected_records
|
||||
|
||||
|
||||
@pytest.mark.parametrize("with_tf_schema", (True, False))
|
||||
def test_readback_tfrecords(
|
||||
ray_start_regular_shared_2_cpus,
|
||||
tmp_path,
|
||||
with_tf_schema,
|
||||
):
|
||||
"""
|
||||
Test reading back TFRecords written using datasets.
|
||||
The dataset we read back should be the same that we wrote.
|
||||
"""
|
||||
|
||||
# The dataset we will write to a .tfrecords file.
|
||||
# Here and in the read_tfrecords call below, we specify `override_num_blocks=1`
|
||||
# to ensure that all rows end up in the same block, which is required
|
||||
# for type inference involving partially missing columns.
|
||||
ds = ray.data.from_items(data_partial(with_tf_schema), override_num_blocks=1)
|
||||
expected_records = tf_records_partial()
|
||||
|
||||
tf_schema = None
|
||||
if with_tf_schema:
|
||||
features = expected_records[0].features
|
||||
tf_schema = _features_to_schema(features)
|
||||
|
||||
# Write the TFRecords.
|
||||
ds.write_tfrecords(tmp_path, tf_schema=tf_schema)
|
||||
# Read the TFRecords.
|
||||
readback_ds = ray.data.read_tfrecords(
|
||||
tmp_path, tf_schema=tf_schema, override_num_blocks=1
|
||||
)
|
||||
_ds_eq_streaming(ds, readback_ds)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("with_tf_schema", (True, False))
|
||||
def test_readback_tfrecords_empty_features(
|
||||
ray_start_regular_shared_2_cpus,
|
||||
tmp_path,
|
||||
with_tf_schema,
|
||||
):
|
||||
"""
|
||||
Test reading back TFRecords written using datasets.
|
||||
The dataset we read back should be the same that we wrote.
|
||||
"""
|
||||
|
||||
# The dataset we will write to a .tfrecords file.
|
||||
ds = ray.data.from_items(data_empty(with_tf_schema))
|
||||
if not with_tf_schema:
|
||||
with pytest.raises(ValueError):
|
||||
# With no schema specified, this should fail because
|
||||
# type inference on completely empty columns is ambiguous.
|
||||
ds.write_tfrecords(tmp_path)
|
||||
else:
|
||||
ds = ray.data.from_items(data_empty(with_tf_schema), override_num_blocks=1)
|
||||
expected_records = tf_records_empty()
|
||||
|
||||
features = expected_records[0].features
|
||||
tf_schema = _features_to_schema(features)
|
||||
|
||||
# Write the TFRecords.
|
||||
ds.write_tfrecords(tmp_path, tf_schema=tf_schema)
|
||||
|
||||
# Read the TFRecords.
|
||||
readback_ds = ray.data.read_tfrecords(
|
||||
tmp_path,
|
||||
tf_schema=tf_schema,
|
||||
override_num_blocks=1,
|
||||
)
|
||||
_ds_eq_streaming(ds, readback_ds)
|
||||
|
||||
|
||||
def test_write_tfrecords_tensor(
|
||||
ray_start_regular_shared_2_cpus, tmp_path, tensor_format_context
|
||||
):
|
||||
"""Test that write_tfrecords handles tensor data by serializing
|
||||
tensors to bytes via tf.io.serialize_tensor, preserving shape and dtype."""
|
||||
import tensorflow as tf
|
||||
|
||||
ds = ray.data.range_tensor(3, shape=(2, 2))
|
||||
|
||||
ds.write_tfrecords(tmp_path)
|
||||
|
||||
# Read back the raw TFRecord examples and deserialize tensors.
|
||||
filenames = sorted(os.listdir(tmp_path))
|
||||
filepaths = [os.path.join(tmp_path, filename) for filename in filenames]
|
||||
raw_dataset = tf.data.TFRecordDataset(filepaths)
|
||||
|
||||
results = []
|
||||
for raw_record in raw_dataset:
|
||||
example = tf.train.Example()
|
||||
example.ParseFromString(raw_record.numpy())
|
||||
serialized = example.features.feature["data"].bytes_list.value[0]
|
||||
tensor = tf.io.parse_tensor(serialized, out_type=tf.int64)
|
||||
results.append(tensor.numpy())
|
||||
|
||||
assert len(results) == 3
|
||||
for i, result in enumerate(results):
|
||||
assert result.shape == (2, 2)
|
||||
expected = np.full((2, 2), i)
|
||||
np.testing.assert_array_equal(result, expected)
|
||||
|
||||
|
||||
def test_write_invalid_tfrecords(ray_start_regular_shared_2_cpus, tmp_path):
|
||||
"""
|
||||
If we try to write a dataset with invalid TFRecord datatypes,
|
||||
ValueError should be raised.
|
||||
"""
|
||||
|
||||
ds = ray.data.from_items([{"item": None}])
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
ds.write_tfrecords(tmp_path)
|
||||
|
||||
|
||||
def test_read_invalid_tfrecords(ray_start_regular_shared_2_cpus, tmp_path):
|
||||
file_path = os.path.join(tmp_path, "file.json")
|
||||
with open(file_path, "w") as file:
|
||||
json.dump({"number": 0, "string": "foo"}, file)
|
||||
|
||||
# Expect RuntimeError raised when reading JSON as TFRecord file.
|
||||
with pytest.raises(RuntimeError, match="Failed to read TFRecord file"):
|
||||
ray.data.read_tfrecords(file_path).schema()
|
||||
|
||||
|
||||
def test_read_with_invalid_schema(
|
||||
ray_start_regular_shared_2_cpus,
|
||||
tmp_path,
|
||||
):
|
||||
from tensorflow_metadata.proto.v0 import schema_pb2
|
||||
|
||||
# The dataset we will write to a .tfrecords file.
|
||||
ds = ray.data.from_items(data_partial(True), override_num_blocks=1)
|
||||
expected_records = tf_records_partial()
|
||||
|
||||
# Build fake schema proto with missing/incorrect field name
|
||||
tf_schema_wrong_name = schema_pb2.Schema()
|
||||
schema_feature = tf_schema_wrong_name.feature.add()
|
||||
schema_feature.name = "wrong_name"
|
||||
schema_feature.type = schema_pb2.FeatureType.INT
|
||||
|
||||
# Build a fake schema proto with incorrect type
|
||||
tf_schema_wrong_type = _features_to_schema(expected_records[0].features)
|
||||
for schema_feature in tf_schema_wrong_type.feature:
|
||||
if schema_feature.name == "bytes_item":
|
||||
schema_feature.type = schema_pb2.FeatureType.INT
|
||||
break
|
||||
|
||||
# Writing with incorrect schema should raise a `ValueError`
|
||||
with pytest.raises(ValueError) as e:
|
||||
ds.write_tfrecords(tmp_path, tf_schema=tf_schema_wrong_name)
|
||||
assert "Found extra unexpected feature" in str(e.value.args[0])
|
||||
|
||||
with pytest.raises(ValueError) as e:
|
||||
ds.write_tfrecords(tmp_path, tf_schema=tf_schema_wrong_type)
|
||||
assert str(e.value.args[0]) == (
|
||||
"Schema field type mismatch during write: "
|
||||
"specified type is int, but underlying type is bytes"
|
||||
)
|
||||
|
||||
# Complete a valid write, then try reading with incorrect schema,
|
||||
# which should raise a `ValueError`.
|
||||
ds.write_tfrecords(tmp_path)
|
||||
with pytest.raises(ValueError) as e:
|
||||
ray.data.read_tfrecords(tmp_path, tf_schema=tf_schema_wrong_name).materialize()
|
||||
assert "Found extra unexpected feature" in str(e.value.args[0])
|
||||
|
||||
with pytest.raises(ValueError) as e:
|
||||
ray.data.read_tfrecords(tmp_path, tf_schema=tf_schema_wrong_type).materialize()
|
||||
assert str(e.value.args[0]) == (
|
||||
"Schema field type mismatch during read: "
|
||||
"specified type is int, but underlying type is bytes"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("min_rows_per_file", [5, 10, 50])
|
||||
def test_write_min_rows_per_file(
|
||||
tmp_path, ray_start_regular_shared_2_cpus, min_rows_per_file
|
||||
):
|
||||
ray.data.range(100, override_num_blocks=20).write_tfrecords(
|
||||
tmp_path, min_rows_per_file=min_rows_per_file
|
||||
)
|
||||
|
||||
for filename in os.listdir(tmp_path):
|
||||
dataset = tf.data.TFRecordDataset(os.path.join(tmp_path, filename))
|
||||
assert len(list(dataset)) == min_rows_per_file
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
if sys.version_info >= (3, 12):
|
||||
# Skip this test for Python 3.12+ due to to incompatibility tensorflow
|
||||
sys.exit(0)
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,93 @@
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import ray
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.data.tests.util import extract_values
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
|
||||
@pytest.mark.parametrize("local_read", [True, False])
|
||||
def test_from_torch_map_style_dataset(ray_start_10_cpus_shared, local_read):
|
||||
class StubDataset(torch.utils.data.Dataset):
|
||||
def __len__(self):
|
||||
return 1
|
||||
|
||||
def __getitem__(self, index):
|
||||
return index
|
||||
|
||||
torch_dataset = StubDataset()
|
||||
|
||||
ray_dataset = ray.data.from_torch(torch_dataset, local_read=local_read)
|
||||
|
||||
actual_data = ray_dataset.take_all()
|
||||
assert actual_data == [{"item": 0}]
|
||||
|
||||
|
||||
def test_from_torch_iterable_style_dataset(ray_start_10_cpus_shared):
|
||||
class StubIterableDataset(torch.utils.data.IterableDataset):
|
||||
def __len__(self):
|
||||
return 1
|
||||
|
||||
def __iter__(self):
|
||||
return iter([0])
|
||||
|
||||
iter_torch_dataset = StubIterableDataset()
|
||||
|
||||
ray_dataset = ray.data.from_torch(iter_torch_dataset)
|
||||
|
||||
actual_data = ray_dataset.take_all()
|
||||
assert actual_data == [{"item": 0}]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("local_read", [True, False])
|
||||
def test_from_torch_boundary_conditions(ray_start_10_cpus_shared, local_read):
|
||||
"""
|
||||
Tests that from_torch respects __len__ for map-style datasets
|
||||
"""
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
class BoundaryTestMapDataset(Dataset):
|
||||
"""A map-style dataset where __len__ is less than the underlying data size."""
|
||||
|
||||
def __init__(self, data, length):
|
||||
super().__init__()
|
||||
self._data = data
|
||||
self._length = length
|
||||
assert self._length <= len(
|
||||
self._data
|
||||
), "Length must be <= data size to properly test boundary conditions"
|
||||
|
||||
def __len__(self):
|
||||
return self._length
|
||||
|
||||
def __getitem__(self, index):
|
||||
if not (0 <= index < self._length):
|
||||
# Note: don't use IndexError because we want to fail clearly if
|
||||
# Ray Data tries to access beyond __len__ - 1
|
||||
raise RuntimeError(
|
||||
f"Index {index} out of bounds for dataset with length {self._length}"
|
||||
)
|
||||
return self._data[index]
|
||||
|
||||
source_data = list(range(10))
|
||||
dataset_len = 8 # Intentionally less than len(source_data)
|
||||
|
||||
# --- Test MapDataset ---
|
||||
map_ds = BoundaryTestMapDataset(source_data, dataset_len)
|
||||
# Expected data only includes elements up to dataset_len - 1
|
||||
expected_items = source_data[:dataset_len]
|
||||
|
||||
ray_ds_map = ray.data.from_torch(map_ds, local_read=local_read)
|
||||
actual_items_map = extract_values("item", list(ray_ds_map.take_all()))
|
||||
|
||||
# This assertion verifies that ray_ds_map didn't try to access index 8 or 9,
|
||||
# which would have raised an IndexError in BoundaryTestMapDataset.__getitem__
|
||||
assert actual_items_map == expected_items
|
||||
assert len(actual_items_map) == dataset_len
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,703 @@
|
||||
"""Tests for TurbopufferDatasink.
|
||||
|
||||
Organized by critical paths:
|
||||
1. Constructor validation
|
||||
2. Client initialization
|
||||
3. Arrow table preparation
|
||||
4. Single-namespace batching
|
||||
5. Transform to Turbopuffer format
|
||||
6. Retry logic
|
||||
7. End-to-end write orchestration
|
||||
8. Streaming behavior
|
||||
9. Multi-namespace writes
|
||||
10. Serialization
|
||||
"""
|
||||
|
||||
import pickle
|
||||
import sys
|
||||
import time
|
||||
import uuid
|
||||
from typing import List
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from packaging.version import parse as parse_version
|
||||
|
||||
from ray.data._internal.datasource.turbopuffer_datasink import TurbopufferDatasink
|
||||
from ray.data._internal.utils.arrow_utils import get_pyarrow_version
|
||||
|
||||
# Skip all tests if PyArrow version is less than 19.0
|
||||
pytestmark = pytest.mark.skipif(
|
||||
get_pyarrow_version() < parse_version("19.0.0"),
|
||||
reason="TurbopufferDatasink tests require PyArrow >= 19.0",
|
||||
)
|
||||
|
||||
# =============================================================================
|
||||
# Fixtures
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_turbopuffer_module(monkeypatch):
|
||||
"""Provide a fake turbopuffer module so imports in the datasink succeed."""
|
||||
fake_module = MagicMock()
|
||||
fake_module.Turbopuffer = MagicMock()
|
||||
with patch.dict(sys.modules, {"turbopuffer": fake_module}):
|
||||
yield fake_module
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sink():
|
||||
"""Default sink with minimal required arguments."""
|
||||
return TurbopufferDatasink(
|
||||
namespace="default_ns",
|
||||
region="gcp-us-central1",
|
||||
api_key="test-api-key",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_client():
|
||||
"""Mock Turbopuffer client with namespace support."""
|
||||
client = MagicMock()
|
||||
client.namespace.return_value = MagicMock()
|
||||
return client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_table():
|
||||
"""Standard table with id and vector columns."""
|
||||
return pa.table(
|
||||
{
|
||||
"id": [1, 2, 3],
|
||||
"vector": [[0.1], [0.2], [0.3]],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def make_sink(**kwargs) -> TurbopufferDatasink:
|
||||
"""Helper to construct a sink with minimal required arguments."""
|
||||
params = {
|
||||
"namespace": "default_ns",
|
||||
"region": "gcp-us-central1",
|
||||
"api_key": "test-api-key",
|
||||
}
|
||||
params.update(kwargs)
|
||||
return TurbopufferDatasink(**params)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# 1. Constructor validation
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestConstructorValidation:
|
||||
"""Tests for constructor argument validation."""
|
||||
|
||||
def test_requires_namespace_or_namespace_column(self):
|
||||
"""Must provide exactly one of namespace / namespace_column."""
|
||||
with pytest.raises(ValueError, match="Either.*must be provided"):
|
||||
TurbopufferDatasink(
|
||||
region="gcp-us-central1",
|
||||
api_key="k",
|
||||
)
|
||||
|
||||
def test_rejects_both_namespace_and_namespace_column(self):
|
||||
"""Cannot provide both namespace and namespace_column."""
|
||||
with pytest.raises(ValueError, match="exactly one"):
|
||||
TurbopufferDatasink(
|
||||
namespace="ns",
|
||||
namespace_column="ns_col",
|
||||
region="gcp-us-central1",
|
||||
api_key="k",
|
||||
)
|
||||
|
||||
def test_namespace_column_cannot_be_id_or_vector(self):
|
||||
"""namespace_column must not collide with id_column or vector_column."""
|
||||
with pytest.raises(ValueError, match="namespace_column.*must not be the same"):
|
||||
make_sink(namespace=None, namespace_column="id")
|
||||
|
||||
with pytest.raises(ValueError, match="namespace_column.*must not be the same"):
|
||||
make_sink(namespace=None, namespace_column="vector")
|
||||
|
||||
def test_api_key_from_env(self, monkeypatch):
|
||||
"""API key can come from environment variable."""
|
||||
monkeypatch.delenv("TURBOPUFFER_API_KEY", raising=False)
|
||||
|
||||
# No api_key and no env var -> error
|
||||
with pytest.raises(ValueError):
|
||||
TurbopufferDatasink(namespace="ns", region="gcp-us-central1")
|
||||
|
||||
# With env var, init should succeed
|
||||
monkeypatch.setenv("TURBOPUFFER_API_KEY", "env-api-key")
|
||||
sink = TurbopufferDatasink(namespace="ns", region="gcp-us-central1")
|
||||
assert sink.api_key == "env-api-key"
|
||||
|
||||
def test_rejects_same_id_and_vector_column(self):
|
||||
"""id_column and vector_column must be distinct."""
|
||||
with pytest.raises(ValueError, match="id_column and vector_column"):
|
||||
make_sink(id_column="doc_id", vector_column="doc_id")
|
||||
|
||||
def test_accepts_region_only(self):
|
||||
"""Constructor succeeds with region and no base_url."""
|
||||
sink = make_sink(region="gcp-us-central1")
|
||||
assert sink.region == "gcp-us-central1"
|
||||
assert sink.base_url is None
|
||||
|
||||
def test_accepts_base_url_only(self):
|
||||
"""Constructor succeeds with base_url and no region."""
|
||||
sink = make_sink(
|
||||
region=None,
|
||||
base_url="https://gcp-us-central1.turbopuffer.com",
|
||||
)
|
||||
assert sink.base_url == "https://gcp-us-central1.turbopuffer.com"
|
||||
assert sink.region is None
|
||||
|
||||
def test_rejects_both_region_and_base_url(self):
|
||||
"""Cannot provide both region and base_url."""
|
||||
with pytest.raises(ValueError, match="exactly one of 'region' or 'base_url'"):
|
||||
make_sink(
|
||||
region="gcp-us-central1",
|
||||
base_url="https://gcp-us-central1.turbopuffer.com",
|
||||
)
|
||||
|
||||
def test_rejects_neither_region_nor_base_url(self):
|
||||
"""Must provide at least one of region or base_url."""
|
||||
with pytest.raises(ValueError, match="Either 'region' or 'base_url'"):
|
||||
TurbopufferDatasink(
|
||||
namespace="ns",
|
||||
api_key="k",
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# 2. Client initialization
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestClientInitialization:
|
||||
"""Tests for Turbopuffer client lazy initialization."""
|
||||
|
||||
def test_lazy_initialization(self, sink, mock_turbopuffer_module):
|
||||
"""Client is created lazily and cached."""
|
||||
client1 = sink._get_client()
|
||||
client2 = sink._get_client()
|
||||
|
||||
assert client1 is client2
|
||||
mock_turbopuffer_module.Turbopuffer.assert_called_once_with(
|
||||
api_key="test-api-key",
|
||||
region="gcp-us-central1",
|
||||
)
|
||||
|
||||
def test_uses_explicit_region(self, mock_turbopuffer_module):
|
||||
"""Client uses the configured region."""
|
||||
sink = make_sink(region="custom-region")
|
||||
sink._get_client()
|
||||
|
||||
mock_turbopuffer_module.Turbopuffer.assert_called_once_with(
|
||||
api_key="test-api-key",
|
||||
region="custom-region",
|
||||
)
|
||||
|
||||
def test_uses_base_url(self, mock_turbopuffer_module):
|
||||
"""Client uses base_url when region is not provided."""
|
||||
sink = make_sink(
|
||||
region=None,
|
||||
base_url="https://gcp-us-central1.turbopuffer.com",
|
||||
)
|
||||
sink._get_client()
|
||||
|
||||
mock_turbopuffer_module.Turbopuffer.assert_called_once_with(
|
||||
api_key="test-api-key",
|
||||
base_url="https://gcp-us-central1.turbopuffer.com",
|
||||
)
|
||||
|
||||
def test_base_url_does_not_pass_region(self, mock_turbopuffer_module):
|
||||
"""When base_url is used, region is not passed to the client."""
|
||||
sink = make_sink(
|
||||
region=None,
|
||||
base_url="https://custom.turbopuffer.com",
|
||||
)
|
||||
sink._get_client()
|
||||
|
||||
call_kwargs = mock_turbopuffer_module.Turbopuffer.call_args[1]
|
||||
assert "region" not in call_kwargs
|
||||
assert call_kwargs["base_url"] == "https://custom.turbopuffer.com"
|
||||
|
||||
def test_region_does_not_pass_base_url(self, mock_turbopuffer_module):
|
||||
"""When region is used, base_url is not passed to the client."""
|
||||
sink = make_sink(region="gcp-us-central1")
|
||||
sink._get_client()
|
||||
|
||||
call_kwargs = mock_turbopuffer_module.Turbopuffer.call_args[1]
|
||||
assert "base_url" not in call_kwargs
|
||||
assert call_kwargs["region"] == "gcp-us-central1"
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# 3. Arrow table preparation
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestArrowTablePreparation:
|
||||
"""Tests for _prepare_arrow_table."""
|
||||
|
||||
def test_renames_columns_and_filters_null_ids(self):
|
||||
"""Custom columns are renamed and null IDs filtered."""
|
||||
table = pa.table(
|
||||
{
|
||||
"doc_id": [1, 2, None],
|
||||
"emb": [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]],
|
||||
}
|
||||
)
|
||||
sink = make_sink(id_column="doc_id", vector_column="emb")
|
||||
|
||||
prepared = sink._prepare_arrow_table(table)
|
||||
|
||||
# Null ID row filtered, columns renamed to id/vector
|
||||
expected = pa.table(
|
||||
{
|
||||
"id": [1, 2],
|
||||
"vector": [[0.1, 0.2], [0.3, 0.4]],
|
||||
}
|
||||
)
|
||||
assert prepared.equals(expected)
|
||||
|
||||
def test_missing_id_column_raises(self):
|
||||
"""Missing custom ID column raises ValueError."""
|
||||
table = pa.table({"other": [1, 2, 3]})
|
||||
sink = make_sink(id_column="doc_id")
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
sink._prepare_arrow_table(table)
|
||||
|
||||
def test_missing_vector_column_raises(self):
|
||||
"""Missing vector column raises ValueError."""
|
||||
table = pa.table({"id": [1, 2, 3]})
|
||||
sink = make_sink(vector_column="embedding")
|
||||
|
||||
with pytest.raises(ValueError, match="Vector column 'embedding' not found"):
|
||||
sink._prepare_arrow_table(table)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"existing_col,custom_col,expected_match",
|
||||
[
|
||||
("id", "doc_id", "already has.*'id' column"),
|
||||
("vector", "emb", "already has.*'vector' column"),
|
||||
],
|
||||
ids=["id_conflict", "vector_conflict"],
|
||||
)
|
||||
def test_conflicting_column_names_raise(
|
||||
self, existing_col, custom_col, expected_match
|
||||
):
|
||||
"""Raise if table already has target column name."""
|
||||
if existing_col == "id":
|
||||
table = pa.table(
|
||||
{"id": [1, 2], "doc_id": [10, 20], "vector": [[0.1], [0.2]]}
|
||||
)
|
||||
sink = make_sink(id_column="doc_id")
|
||||
else:
|
||||
table = pa.table(
|
||||
{"id": [1, 2], "vector": [[0.1], [0.2]], "emb": [[0.3], [0.4]]}
|
||||
)
|
||||
sink = make_sink(vector_column="emb")
|
||||
|
||||
with pytest.raises(ValueError, match=expected_match):
|
||||
sink._prepare_arrow_table(table)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# 4. Single-namespace batching
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestSingleNamespaceBatching:
|
||||
"""Tests for write batching behavior."""
|
||||
|
||||
def test_batches_by_batch_size(self, mock_client):
|
||||
"""Large tables are split into batches."""
|
||||
num_rows = 25
|
||||
table = pa.table(
|
||||
{
|
||||
"id": list(range(num_rows)),
|
||||
"vector": [[float(i)] for i in range(num_rows)],
|
||||
}
|
||||
)
|
||||
sink = make_sink(batch_size=10)
|
||||
batch_sizes: List[int] = []
|
||||
|
||||
def track_batch(ns, batch, namespace_name=None):
|
||||
# batch is a RecordBatch, get its row count
|
||||
batch_sizes.append(batch.num_rows)
|
||||
|
||||
with patch.object(sink, "_get_client", return_value=mock_client):
|
||||
with patch.object(sink, "_write_batch_with_retry", side_effect=track_batch):
|
||||
sink.write([table], ctx=None)
|
||||
|
||||
assert batch_sizes == [10, 10, 5]
|
||||
|
||||
def test_skips_empty_blocks(self, sink):
|
||||
"""Empty blocks don't trigger namespace writes."""
|
||||
empty_table = pa.table({"id": [], "vector": []})
|
||||
|
||||
with patch.object(sink, "_get_client") as mock_get_client:
|
||||
with patch.object(sink, "_write_batch_with_retry") as mock_write:
|
||||
mock_get_client.return_value = MagicMock()
|
||||
sink.write([empty_table], ctx=None)
|
||||
|
||||
mock_write.assert_not_called()
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# 5. Transform to Turbopuffer format
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestTransformToTurbopufferFormat:
|
||||
"""Tests for _transform_to_turbopuffer_format."""
|
||||
|
||||
def test_requires_id_column(self, sink):
|
||||
"""Table must have 'id' column."""
|
||||
table = pa.table({"col": [1, 2, 3]})
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
sink._transform_to_turbopuffer_format(table)
|
||||
|
||||
def test_converts_uuid_bytes_to_native_uuid(self, sink):
|
||||
"""16-byte binary IDs become native uuid.UUID objects.
|
||||
|
||||
Per Turbopuffer performance docs, native UUIDs (16 bytes) are more
|
||||
efficient than string UUIDs (36 bytes).
|
||||
"""
|
||||
u = uuid.uuid4()
|
||||
|
||||
# ID column must be binary(16) for UUID conversion
|
||||
table = pa.table(
|
||||
{
|
||||
"id": pa.array([u.bytes], type=pa.binary(16)),
|
||||
"vector": [[0.1, 0.2]],
|
||||
}
|
||||
)
|
||||
|
||||
columns = sink._transform_to_turbopuffer_format(table)
|
||||
|
||||
expected = {
|
||||
"id": [u], # Native uuid.UUID, not bytes
|
||||
"vector": [[0.1, 0.2]],
|
||||
}
|
||||
assert columns == expected
|
||||
assert isinstance(columns["id"][0], uuid.UUID)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# 6. Retry logic
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestRetryLogic:
|
||||
"""Tests for _write_batch_with_retry."""
|
||||
|
||||
@pytest.fixture
|
||||
def sample_batch(self):
|
||||
"""A simple batch for retry tests."""
|
||||
return pa.table({"id": [1], "vector": [[0.1]]})
|
||||
|
||||
def test_success_first_try(self, sink, sample_batch):
|
||||
"""Successful write on first attempt."""
|
||||
namespace = MagicMock()
|
||||
|
||||
sink._write_batch_with_retry(namespace, sample_batch)
|
||||
|
||||
namespace.write.assert_called_once_with(
|
||||
upsert_columns={"id": [1], "vector": [[0.1]]},
|
||||
schema=None,
|
||||
distance_metric="cosine_distance",
|
||||
)
|
||||
|
||||
def test_retries_then_succeeds(self, sink, sample_batch, monkeypatch):
|
||||
"""Transient failures are retried."""
|
||||
monkeypatch.setattr(time, "sleep", lambda _: None)
|
||||
namespace = MagicMock()
|
||||
attempts = {"count": 0}
|
||||
|
||||
def flaky_write(*args, **kwargs):
|
||||
attempts["count"] += 1
|
||||
if attempts["count"] < 3:
|
||||
raise RuntimeError("temporary error")
|
||||
|
||||
namespace.write.side_effect = flaky_write
|
||||
|
||||
sink._write_batch_with_retry(namespace, sample_batch)
|
||||
|
||||
assert attempts["count"] == 3
|
||||
|
||||
def test_exhausts_retries_and_raises(self, sink, sample_batch, monkeypatch):
|
||||
"""Persistent failures exhaust retries and raise."""
|
||||
monkeypatch.setattr(time, "sleep", lambda _: None)
|
||||
namespace = MagicMock()
|
||||
namespace.write.side_effect = RuntimeError("persistent error")
|
||||
|
||||
with pytest.raises(RuntimeError, match="persistent error"):
|
||||
sink._write_batch_with_retry(namespace, sample_batch)
|
||||
|
||||
assert namespace.write.call_count == 5 # max_attempts=5
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"schema,distance_metric",
|
||||
[
|
||||
({"field": "value"}, "cosine_distance"),
|
||||
(None, "euclidean_squared"),
|
||||
({"type": "string"}, "euclidean_squared"),
|
||||
],
|
||||
ids=["with_schema", "alt_metric", "both"],
|
||||
)
|
||||
def test_configurable_options(self, schema, distance_metric):
|
||||
"""Schema and distance_metric are passed to write."""
|
||||
sink = make_sink(schema=schema, distance_metric=distance_metric)
|
||||
namespace = MagicMock()
|
||||
batch = pa.table({"id": [1], "vector": [[0.1]]})
|
||||
|
||||
sink._write_batch_with_retry(namespace, batch)
|
||||
|
||||
namespace.write.assert_called_once_with(
|
||||
upsert_columns={"id": [1], "vector": [[0.1]]},
|
||||
schema=schema,
|
||||
distance_metric=distance_metric,
|
||||
)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# 7. End-to-end write orchestration
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestWriteOrchestration:
|
||||
"""Tests for top-level write() method."""
|
||||
|
||||
def test_write_multiple_blocks(self, sink):
|
||||
"""Multiple blocks are processed and written."""
|
||||
blocks = [
|
||||
pa.table({"id": [1, 2], "vector": [[1.0], [2.0]]}),
|
||||
pa.table({"id": [3], "vector": [[3.0]]}),
|
||||
]
|
||||
write_calls = []
|
||||
|
||||
def track_write(ns, batch, namespace_name=None):
|
||||
write_calls.append(batch.num_rows)
|
||||
|
||||
with patch.object(sink, "_get_client") as mock_get_client:
|
||||
mock_client = MagicMock()
|
||||
mock_get_client.return_value = mock_client
|
||||
with patch.object(sink, "_write_batch_with_retry", side_effect=track_write):
|
||||
sink.write(blocks, ctx=None)
|
||||
|
||||
# Two blocks written
|
||||
assert len(write_calls) == 2
|
||||
assert write_calls == [2, 1]
|
||||
|
||||
# Namespace accessed with correct name
|
||||
mock_client.namespace.assert_called_with("default_ns")
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# 8. Streaming behavior (memory efficiency)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestStreamingBehavior:
|
||||
"""Tests for memory-efficient streaming writes."""
|
||||
|
||||
def test_processes_blocks_independently(self, sink):
|
||||
"""Each block is processed and written separately."""
|
||||
blocks = [pa.table({"id": [i], "vector": [[float(i)]]}) for i in range(5)]
|
||||
write_counts = []
|
||||
|
||||
def track_write(ns, batch, namespace_name=None):
|
||||
write_counts.append(batch.num_rows)
|
||||
|
||||
with patch.object(sink, "_get_client", return_value=MagicMock()):
|
||||
with patch.object(sink, "_write_batch_with_retry", side_effect=track_write):
|
||||
sink.write(blocks, ctx=None)
|
||||
|
||||
# 5 blocks → 5 writes of 1 row each
|
||||
assert len(write_counts) == 5
|
||||
assert all(c == 1 for c in write_counts)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# 9. Multi-namespace writes
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestMultiNamespaceWrites:
|
||||
"""Tests for namespace_column-driven multi-namespace writes."""
|
||||
|
||||
def test_routes_rows_to_correct_namespaces(self):
|
||||
"""Rows are grouped by namespace_column and written to the right ns."""
|
||||
sink = make_sink(namespace=None, namespace_column="tenant")
|
||||
table = pa.table(
|
||||
{
|
||||
"tenant": ["ns_a", "ns_b", "ns_a", "ns_b"],
|
||||
"id": [1, 2, 3, 4],
|
||||
"vector": [[0.1], [0.2], [0.3], [0.4]],
|
||||
}
|
||||
)
|
||||
|
||||
writes = {} # namespace_name -> list of row counts
|
||||
|
||||
def track_write(ns, batch, namespace_name=None):
|
||||
writes.setdefault(namespace_name, []).append(batch.num_rows)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.namespace.return_value = MagicMock()
|
||||
with patch.object(sink, "_get_client", return_value=mock_client):
|
||||
with patch.object(sink, "_write_batch_with_retry", side_effect=track_write):
|
||||
sink.write([table], ctx=None)
|
||||
|
||||
assert "ns_a" in writes
|
||||
assert "ns_b" in writes
|
||||
assert sum(writes["ns_a"]) == 2
|
||||
assert sum(writes["ns_b"]) == 2
|
||||
|
||||
def test_drops_namespace_column_before_writing(self):
|
||||
"""The namespace column is not included in the written data."""
|
||||
sink = make_sink(namespace=None, namespace_column="tenant")
|
||||
table = pa.table(
|
||||
{
|
||||
"tenant": ["ns_a"],
|
||||
"id": [1],
|
||||
"vector": [[0.1]],
|
||||
}
|
||||
)
|
||||
|
||||
written_batches = []
|
||||
|
||||
def capture_batch(ns, batch, namespace_name=None):
|
||||
written_batches.append(batch)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.namespace.return_value = MagicMock()
|
||||
with patch.object(sink, "_get_client", return_value=mock_client):
|
||||
with patch.object(
|
||||
sink, "_write_batch_with_retry", side_effect=capture_batch
|
||||
):
|
||||
sink.write([table], ctx=None)
|
||||
|
||||
assert len(written_batches) == 1
|
||||
assert "tenant" not in written_batches[0].column_names
|
||||
assert "id" in written_batches[0].column_names
|
||||
|
||||
def test_missing_namespace_column_raises(self):
|
||||
"""Missing namespace column in data raises ValueError."""
|
||||
sink = make_sink(namespace=None, namespace_column="tenant")
|
||||
table = pa.table(
|
||||
{
|
||||
"id": [1],
|
||||
"vector": [[0.1]],
|
||||
}
|
||||
)
|
||||
|
||||
mock_client = MagicMock()
|
||||
with patch.object(sink, "_get_client", return_value=mock_client):
|
||||
with pytest.raises(ValueError, match="Namespace column.*not found"):
|
||||
sink.write([table], ctx=None)
|
||||
|
||||
def test_null_namespace_values_raise(self):
|
||||
"""Null values in namespace column raise ValueError."""
|
||||
sink = make_sink(namespace=None, namespace_column="tenant")
|
||||
table = pa.table(
|
||||
{
|
||||
"tenant": ["ns_a", None],
|
||||
"id": [1, 2],
|
||||
"vector": [[0.1], [0.2]],
|
||||
}
|
||||
)
|
||||
|
||||
mock_client = MagicMock()
|
||||
with patch.object(sink, "_get_client", return_value=mock_client):
|
||||
with pytest.raises(ValueError, match="contains null values"):
|
||||
sink.write([table], ctx=None)
|
||||
|
||||
def test_skips_empty_blocks_in_multi_namespace(self):
|
||||
"""Empty blocks are skipped in multi-namespace mode."""
|
||||
sink = make_sink(namespace=None, namespace_column="tenant")
|
||||
empty_table = pa.table(
|
||||
{
|
||||
"tenant": pa.array([], type=pa.string()),
|
||||
"id": pa.array([], type=pa.int64()),
|
||||
"vector": pa.array([], type=pa.list_(pa.float64())),
|
||||
}
|
||||
)
|
||||
|
||||
mock_client = MagicMock()
|
||||
with patch.object(sink, "_get_client", return_value=mock_client):
|
||||
with patch.object(sink, "_write_batch_with_retry") as mock_write:
|
||||
sink.write([empty_table], ctx=None)
|
||||
|
||||
mock_write.assert_not_called()
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# 10. Serialization behavior
|
||||
# =============================================================================
|
||||
|
||||
|
||||
class TestSerialization:
|
||||
"""Tests for pickle serialization support."""
|
||||
|
||||
def test_preserves_configuration(self, sink, mock_turbopuffer_module):
|
||||
"""Configuration is preserved after pickle round-trip."""
|
||||
pickled = pickle.dumps(sink)
|
||||
unpickled = pickle.loads(pickled)
|
||||
|
||||
assert unpickled.namespace == sink.namespace
|
||||
assert unpickled.namespace_column == sink.namespace_column
|
||||
assert unpickled.api_key == sink.api_key
|
||||
assert unpickled.region == sink.region
|
||||
assert unpickled.base_url == sink.base_url
|
||||
assert unpickled.batch_size == sink.batch_size
|
||||
assert unpickled._client is None
|
||||
|
||||
# Lazy initialization works after unpickling
|
||||
client = unpickled._get_client()
|
||||
assert client is not None
|
||||
mock_turbopuffer_module.Turbopuffer.assert_called()
|
||||
|
||||
def test_preserves_namespace_column_configuration(self, mock_turbopuffer_module):
|
||||
"""namespace_column configuration survives pickle round-trip."""
|
||||
sink = make_sink(namespace=None, namespace_column="tenant")
|
||||
pickled = pickle.dumps(sink)
|
||||
unpickled = pickle.loads(pickled)
|
||||
|
||||
assert unpickled.namespace is None
|
||||
assert unpickled.namespace_column == "tenant"
|
||||
assert unpickled._client is None
|
||||
|
||||
def test_preserves_base_url_configuration(self, mock_turbopuffer_module):
|
||||
"""base_url configuration survives pickle round-trip."""
|
||||
sink = make_sink(
|
||||
region=None,
|
||||
base_url="https://gcp-us-central1.turbopuffer.com",
|
||||
)
|
||||
pickled = pickle.dumps(sink)
|
||||
unpickled = pickle.loads(pickled)
|
||||
|
||||
assert unpickled.region is None
|
||||
assert unpickled.base_url == "https://gcp-us-central1.turbopuffer.com"
|
||||
assert unpickled._client is None
|
||||
|
||||
# Lazy initialization works and uses base_url
|
||||
unpickled._get_client()
|
||||
mock_turbopuffer_module.Turbopuffer.assert_called_once_with(
|
||||
api_key="test-api-key",
|
||||
base_url="https://gcp-us-central1.turbopuffer.com",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,35 @@
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
|
||||
|
||||
def test_read_videos():
|
||||
uri = "s3://anonymous@ray-example-data/basketball.mp4"
|
||||
ds = ray.data.read_videos(uri, include_timestamps=True)
|
||||
|
||||
assert ds.count() == 333
|
||||
assert ds.schema().names == ["frame", "frame_index", "frame_timestamp"]
|
||||
|
||||
frame_indices = ds.select_columns(["frame_index"]).take_all()
|
||||
assert sorted(frame_indices, key=lambda item: item["frame_index"]) == [
|
||||
{"frame_index": i} for i in range(333)
|
||||
]
|
||||
|
||||
frame_timestamps = ds.select_columns(["frame_timestamp"]).take_all()
|
||||
for t in frame_timestamps:
|
||||
assert isinstance(t["frame_timestamp"], np.ndarray)
|
||||
assert t["frame_timestamp"].shape[0] == 2
|
||||
|
||||
frame_type, frame_index_type, _ = ds.schema().types
|
||||
|
||||
assert frame_type.shape == (720, 1280, 3)
|
||||
assert frame_type.value_type == pa.uint8()
|
||||
assert frame_index_type == pa.int64()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,335 @@
|
||||
# Copyright NVIDIA Corporation 2023
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import glob
|
||||
import io
|
||||
import os
|
||||
import pickle
|
||||
import tarfile
|
||||
|
||||
import pytest
|
||||
import webdataset as wds
|
||||
|
||||
import ray
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
|
||||
class TarWriter:
|
||||
def __init__(self, path):
|
||||
self.path = path
|
||||
self.tar = tarfile.open(path, "w")
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *args):
|
||||
self.tar.close()
|
||||
|
||||
def write(self, name, data):
|
||||
f = self.tar.tarinfo()
|
||||
f.name = name
|
||||
f.size = len(data)
|
||||
self.tar.addfile(f, io.BytesIO(data))
|
||||
|
||||
|
||||
def test_webdataset_read(ray_start_2_cpus, tmp_path):
|
||||
path = os.path.join(tmp_path, "bar_000000.tar")
|
||||
with TarWriter(path) as tf:
|
||||
for i in range(100):
|
||||
tf.write(f"{i}.a", str(i).encode("utf-8"))
|
||||
tf.write(f"{i}.b", str(i**2).encode("utf-8"))
|
||||
assert os.path.exists(path)
|
||||
assert len(glob.glob(f"{tmp_path}/*.tar")) == 1
|
||||
ds = ray.data.read_webdataset(paths=[str(tmp_path)])
|
||||
samples = ds.take(100)
|
||||
assert len(samples) == 100
|
||||
for i, sample in enumerate(samples):
|
||||
assert isinstance(sample, dict), sample
|
||||
assert sample["__key__"] == str(i)
|
||||
assert sample["a"].decode("utf-8") == str(i)
|
||||
assert sample["b"].decode("utf-8") == str(i**2)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def allow_unsafe_deserialization(monkeypatch):
|
||||
monkeypatch.setenv("RAY_DATA_WEBDATASET_ALLOW_UNSAFE_DESERIALIZATION", "1")
|
||||
|
||||
|
||||
def test_webdataset_expand_json(
|
||||
ray_start_2_cpus, tmp_path, allow_unsafe_deserialization
|
||||
):
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
|
||||
gray = np.random.randint(0, 255, (100, 100), dtype=np.uint8)
|
||||
dstruct = dict(a=[1, 2], b=dict(c=2), d="hello")
|
||||
ttensor = torch.tensor([1, 2, 3]).numpy()
|
||||
|
||||
sample = {
|
||||
"__key__": "foo",
|
||||
"jpg": image,
|
||||
"gray.png": gray,
|
||||
"mp": dstruct,
|
||||
"json": dstruct,
|
||||
"pt": ttensor,
|
||||
"und": b"undecoded",
|
||||
"custom": b"nothing",
|
||||
}
|
||||
|
||||
# write the encoded data using the default encoder
|
||||
data = [sample]
|
||||
ds = ray.data.from_items(data).repartition(1)
|
||||
ds.write_webdataset(path=tmp_path, try_create_dir=True)
|
||||
ds = ray.data.read_webdataset(
|
||||
paths=[str(tmp_path)], override_num_blocks=1, expand_json=True
|
||||
)
|
||||
record = ds.take(1)
|
||||
assert [1, 2] == record[0]["a"]
|
||||
|
||||
|
||||
def test_webdataset_suffixes(ray_start_2_cpus, tmp_path):
|
||||
path = os.path.join(tmp_path, "bar_000000.tar")
|
||||
with TarWriter(path) as tf:
|
||||
for i in range(100):
|
||||
tf.write(f"{i}.txt", str(i).encode("utf-8"))
|
||||
tf.write(f"{i}.test.txt", str(i**2).encode("utf-8"))
|
||||
tf.write(f"{i}.cls", str(i**2).encode("utf-8"))
|
||||
tf.write(f"{i}.test.cls2", str(i**2).encode("utf-8"))
|
||||
assert os.path.exists(path)
|
||||
assert len(glob.glob(f"{tmp_path}/*.tar")) == 1
|
||||
|
||||
# test simple suffixes
|
||||
ds = ray.data.read_webdataset(paths=[str(tmp_path)], suffixes=["txt", "cls"])
|
||||
samples = ds.take(100)
|
||||
assert len(samples) == 100
|
||||
for i, sample in enumerate(samples):
|
||||
assert set(sample.keys()) == {"__url__", "__key__", "txt", "cls"}
|
||||
|
||||
# test fnmatch patterns for suffixes
|
||||
ds = ray.data.read_webdataset(paths=[str(tmp_path)], suffixes=["*.txt", "*.cls"])
|
||||
samples = ds.take(100)
|
||||
assert len(samples) == 100
|
||||
for i, sample in enumerate(samples):
|
||||
assert set(sample.keys()) == {"__url__", "__key__", "txt", "cls", "test.txt"}
|
||||
|
||||
# test selection function
|
||||
def select(name):
|
||||
return name.endswith("txt")
|
||||
|
||||
ds = ray.data.read_webdataset(paths=[str(tmp_path)], suffixes=select)
|
||||
samples = ds.take(100)
|
||||
assert len(samples) == 100
|
||||
for i, sample in enumerate(samples):
|
||||
assert set(sample.keys()) == {"__url__", "__key__", "txt", "test.txt"}
|
||||
|
||||
# test filerename
|
||||
def renamer(name):
|
||||
result = name.replace("txt", "text")
|
||||
print("***", name, result)
|
||||
return result
|
||||
|
||||
ds = ray.data.read_webdataset(paths=[str(tmp_path)], filerename=renamer)
|
||||
samples = ds.take(100)
|
||||
assert len(samples) == 100
|
||||
for i, sample in enumerate(samples):
|
||||
assert set(sample.keys()) == {
|
||||
"__url__",
|
||||
"__key__",
|
||||
"text",
|
||||
"cls",
|
||||
"test.text",
|
||||
"test.cls2",
|
||||
}
|
||||
|
||||
|
||||
def test_webdataset_write(ray_start_2_cpus, tmp_path):
|
||||
print(ray.available_resources())
|
||||
data = [dict(__key__=str(i), a=str(i), b=str(i**2)) for i in range(100)]
|
||||
ds = ray.data.from_items(data).repartition(1)
|
||||
ds.write_webdataset(path=tmp_path, try_create_dir=True)
|
||||
paths = glob.glob(f"{tmp_path}/*.tar")
|
||||
assert len(paths) == 1
|
||||
with open(paths[0], "rb") as stream:
|
||||
tf = tarfile.open(fileobj=stream)
|
||||
for i in range(100):
|
||||
assert tf.extractfile(f"{i}.a").read().decode("utf-8") == str(i)
|
||||
assert tf.extractfile(f"{i}.b").read().decode("utf-8") == str(i**2)
|
||||
|
||||
|
||||
def custom_decoder(sample):
|
||||
for key, value in sample.items():
|
||||
if key == "png":
|
||||
# check that images have already been decoded
|
||||
assert not isinstance(value, bytes)
|
||||
elif key.endswith("custom"):
|
||||
sample[key] = "custom-value"
|
||||
return sample
|
||||
|
||||
|
||||
def test_webdataset_coding(ray_start_2_cpus, tmp_path, allow_unsafe_deserialization):
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
import torch
|
||||
|
||||
image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
|
||||
gray = np.random.randint(0, 255, (100, 100), dtype=np.uint8)
|
||||
dstruct = dict(a=[1], b=dict(c=2), d="hello")
|
||||
ttensor = torch.tensor([1, 2, 3]).numpy()
|
||||
|
||||
sample = {
|
||||
"__key__": "foo",
|
||||
"jpg": image,
|
||||
"gray.png": gray,
|
||||
"mp": dstruct,
|
||||
"json": dstruct,
|
||||
"pt": ttensor,
|
||||
"und": b"undecoded",
|
||||
"custom": b"nothing",
|
||||
}
|
||||
|
||||
# write the encoded data using the default encoder
|
||||
data = [sample]
|
||||
ds = ray.data.from_items(data).repartition(1)
|
||||
ds.write_webdataset(path=tmp_path, try_create_dir=True)
|
||||
|
||||
# read the encoded data using the default decoder
|
||||
paths = glob.glob(f"{tmp_path}/*.tar")
|
||||
assert len(paths) == 1
|
||||
path = paths[0]
|
||||
assert os.path.exists(path)
|
||||
ds = ray.data.read_webdataset(paths=[str(tmp_path)])
|
||||
samples = ds.take(1)
|
||||
assert len(samples) == 1
|
||||
for sample in samples:
|
||||
assert isinstance(sample, dict), sample
|
||||
assert sample["__key__"] == "foo"
|
||||
assert isinstance(sample["jpg"], np.ndarray)
|
||||
assert sample["jpg"].shape == (100, 100, 3)
|
||||
assert isinstance(sample["gray.png"], np.ndarray)
|
||||
assert sample["gray.png"].shape == (100, 100)
|
||||
assert isinstance(sample["mp"], dict)
|
||||
assert sample["mp"]["a"] == [1]
|
||||
assert sample["mp"]["b"]["c"] == 2
|
||||
assert isinstance(sample["json"], dict)
|
||||
assert sample["json"]["a"] == [1]
|
||||
assert isinstance(sample["pt"], np.ndarray)
|
||||
assert sample["pt"].tolist() == [1, 2, 3]
|
||||
|
||||
# test the format argument to the default decoder and multiple decoders
|
||||
ds = ray.data.read_webdataset(
|
||||
paths=[str(tmp_path)], decoder=["PIL", custom_decoder]
|
||||
)
|
||||
samples = ds.take(1)
|
||||
assert len(samples) == 1
|
||||
for sample in samples:
|
||||
assert isinstance(sample, dict), sample
|
||||
assert sample["__key__"] == "foo"
|
||||
assert isinstance(sample["jpg"], PIL.Image.Image)
|
||||
assert isinstance(sample["gray.png"], PIL.Image.Image)
|
||||
assert isinstance(sample["und"], bytes)
|
||||
assert sample["und"] == b"undecoded"
|
||||
assert sample["custom"] == "custom-value"
|
||||
|
||||
|
||||
def test_webdataset_decoding(ray_start_2_cpus, tmp_path):
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
image = np.random.randint(0, 255, (100, 100, 3), dtype=np.uint8)
|
||||
gray = np.random.randint(0, 255, (100, 100), dtype=np.uint8)
|
||||
dstruct = dict(a=np.nan, b=dict(c=2), d="hello", e={"img_filename": "for_test.jpg"})
|
||||
ttensor = torch.tensor([1, 2, 3]).numpy()
|
||||
|
||||
sample = {
|
||||
"__key__": "foo",
|
||||
"jpg": image,
|
||||
"gray.png": gray,
|
||||
"mp": dstruct,
|
||||
"json": dstruct,
|
||||
"pt": ttensor,
|
||||
"und": b"undecoded",
|
||||
"custom": b"nothing",
|
||||
}
|
||||
|
||||
# write the encoded data using the default encoder
|
||||
data = [sample]
|
||||
ds = ray.data.from_items(data).repartition(1)
|
||||
ds.write_webdataset(path=tmp_path, try_create_dir=True)
|
||||
|
||||
ds = ray.data.read_webdataset(
|
||||
paths=[str(tmp_path)],
|
||||
override_num_blocks=1,
|
||||
decoder=None,
|
||||
)
|
||||
samples = ds.take(1)
|
||||
import json
|
||||
|
||||
meta_json = json.loads(samples[0]["json"].decode("utf-8"))
|
||||
assert meta_json["e"]["img_filename"] == "for_test.jpg"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("min_rows_per_file", [5, 10, 50])
|
||||
def test_write_min_rows_per_file(tmp_path, ray_start_2_cpus, min_rows_per_file):
|
||||
ray.data.from_items(
|
||||
[{"id": str(i)} for i in range(100)], override_num_blocks=20
|
||||
).write_webdataset(tmp_path, min_rows_per_file=min_rows_per_file)
|
||||
|
||||
for filename in os.listdir(tmp_path):
|
||||
dataset = wds.WebDataset(os.path.join(tmp_path, filename))
|
||||
assert len(list(dataset)) == min_rows_per_file
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"filename",
|
||||
["000000.pkl", "000000.pickle", "000000.pt", "000000.pth"],
|
||||
)
|
||||
def test_default_decoder_rejects_unsafe_extensions(
|
||||
ray_start_2_cpus, tmp_path, filename
|
||||
):
|
||||
path = os.path.join(tmp_path, "unsafe.tar")
|
||||
with TarWriter(path) as tf:
|
||||
tf.write(filename, b"fake-payload")
|
||||
|
||||
ds = ray.data.read_webdataset(paths=[str(tmp_path)])
|
||||
with pytest.raises(Exception, match="Refusing to"):
|
||||
ds.take_all()
|
||||
|
||||
|
||||
def test_default_decoder_allows_unsafe_with_env_var(
|
||||
ray_start_2_cpus, tmp_path, allow_unsafe_deserialization
|
||||
):
|
||||
path = os.path.join(tmp_path, "trusted.tar")
|
||||
with TarWriter(path) as tf:
|
||||
tf.write("000000.pkl", pickle.dumps({"key": "value"}))
|
||||
|
||||
ds = ray.data.read_webdataset(paths=[str(tmp_path)])
|
||||
rows = ds.take_all()
|
||||
|
||||
assert len(rows) == 1
|
||||
assert rows[0]["pkl"] == {"key": "value"}
|
||||
|
||||
|
||||
def test_custom_decoder_bypasses_unsafe_guard(ray_start_2_cpus, tmp_path):
|
||||
path = os.path.join(tmp_path, "custom.tar")
|
||||
with TarWriter(path) as tf:
|
||||
tf.write("000000.pkl", pickle.dumps({"key": "value"}))
|
||||
|
||||
def safe_pkl_decoder(sample):
|
||||
sample = dict(sample)
|
||||
for key, value in sample.items():
|
||||
if key == "pkl":
|
||||
sample[key] = pickle.loads(value)
|
||||
return sample
|
||||
|
||||
ds = ray.data.read_webdataset(paths=[str(tmp_path)], decoder=safe_pkl_decoder)
|
||||
rows = ds.take_all()
|
||||
|
||||
assert len(rows) == 1
|
||||
assert rows[0]["pkl"] == {"key": "value"}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,937 @@
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import fsspec
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pyarrow.fs
|
||||
import pytest
|
||||
import zarr
|
||||
from pytest_lazy_fixtures import lf as lazy_fixture
|
||||
|
||||
import ray
|
||||
from ray.data._internal.datasource import zarrv2_datasource
|
||||
from ray.data.block import BlockAccessor
|
||||
from ray.data.tests.conftest import * # noqa: F401, F403
|
||||
|
||||
|
||||
def _execute_read_tasks(tasks) -> pd.DataFrame:
|
||||
frames = [
|
||||
BlockAccessor.for_block(block).to_pandas() for task in tasks for block in task()
|
||||
]
|
||||
return pd.concat(frames, ignore_index=True)
|
||||
|
||||
|
||||
def _reconstruct_array(df: pd.DataFrame, array_name: str) -> np.ndarray:
|
||||
"""Concatenate all chunks of one array from a long-form result frame."""
|
||||
sub = df[df["array"] == array_name].sort_values(
|
||||
"chunk_index", key=lambda col: col.map(tuple)
|
||||
)
|
||||
return np.concatenate(list(sub["chunk"]), axis=0)
|
||||
|
||||
|
||||
def _write_real_zarr_store(
|
||||
store_path: Path,
|
||||
arrays: dict, # {name: (data, chunks)}
|
||||
) -> Path:
|
||||
"""Write a real Zarr v2 store from numpy arrays and consolidate metadata."""
|
||||
root = zarr.open_group(str(store_path), mode="w")
|
||||
for name, (data, chunks) in arrays.items():
|
||||
root.create_dataset(name, data=data, chunks=chunks, dtype=data.dtype)
|
||||
zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
|
||||
return store_path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def zarrv2_group_store(tmp_path) -> Path:
|
||||
"""Two arrays at the store root, both 2-D and 1-D, axis-0-aligned (shape[0]==5)."""
|
||||
return _write_real_zarr_store(
|
||||
tmp_path / "group.zarr",
|
||||
{
|
||||
"images": (np.arange(20, dtype="<i4").reshape(5, 4), (2, 4)),
|
||||
"nested": (np.arange(5, dtype="|u1"), (2,)),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def zarrv2_root_store(tmp_path) -> Path:
|
||||
"""Single-array store with the array sitting directly at the store root."""
|
||||
store_path = tmp_path / "root.zarr"
|
||||
arr = zarr.open(
|
||||
str(store_path),
|
||||
mode="w",
|
||||
shape=(5, 4),
|
||||
chunks=(2, 4),
|
||||
dtype="<i4",
|
||||
)
|
||||
arr[:] = np.arange(20, dtype="<i4").reshape(5, 4)
|
||||
zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
|
||||
return store_path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def local_fsspec_fs():
|
||||
"""fsspec local filesystem (for parametrized cross-fs read tests)."""
|
||||
return fsspec.filesystem("file")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def heterogeneous_zarrv2_store(tmp_path) -> Path:
|
||||
"""A store mixing different ranks, shape[0]s, dtypes, and native chunk sizes.
|
||||
|
||||
Mirrors the UMI-style real-world layout where ``data/*`` arrays share an
|
||||
axis-0 timestep count but differ in everything else, and ``meta/*``
|
||||
arrays live in a separate axis-0 universe entirely. The chunk-per-row
|
||||
datasource handles all of these in one read; nothing has to align.
|
||||
"""
|
||||
store_path = tmp_path / "heterogeneous.zarr"
|
||||
root = zarr.open_group(str(store_path), mode="w")
|
||||
# 4-D image tensor with tiny axis-0 chunks (1 image per chunk).
|
||||
root.create_dataset(
|
||||
"data/camera0_rgb",
|
||||
data=np.arange(20 * 2 * 2 * 3, dtype="|u1").reshape(20, 2, 2, 3),
|
||||
chunks=(1, 2, 2, 3),
|
||||
)
|
||||
# 2-D pose array, same shape[0]=20, much larger axis-0 chunks (10).
|
||||
root.create_dataset(
|
||||
"data/robot0_eef_pos",
|
||||
data=np.arange(20 * 3, dtype="<f4").reshape(20, 3),
|
||||
chunks=(10, 3),
|
||||
)
|
||||
# Episode-boundary metadata: separate axis-0 universe.
|
||||
root.create_dataset(
|
||||
"meta/episode_ends",
|
||||
data=np.array([5, 12, 20], dtype="<i8"),
|
||||
chunks=(3,),
|
||||
)
|
||||
zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
|
||||
return store_path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def unconsolidated_zarrv2_store(tmp_path) -> Path:
|
||||
"""Two arrays at the store root, no ``.zmetadata``.
|
||||
|
||||
Exercises the no-``.zmetadata`` code paths (per-array ``.zarray``
|
||||
discovery and full-store walk) — the common shape of real-world stores
|
||||
behind plain HTTPS or other listing-less filesystems.
|
||||
"""
|
||||
store_path = tmp_path / "unconsolidated.zarr"
|
||||
root = zarr.open_group(str(store_path), mode="w")
|
||||
root.create_dataset(
|
||||
"images", data=np.arange(20, dtype="<i4").reshape(5, 4), chunks=(2, 4)
|
||||
)
|
||||
root.create_dataset("nested", data=np.arange(5, dtype="|u1"), chunks=(2,))
|
||||
return store_path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def aligned_zarrv2_store(tmp_path) -> Path:
|
||||
"""Three arrays sharing ``shape[0]=8``, different ranks and native chunks.
|
||||
|
||||
Models the UMI-style case where data arrays co-stride on the timestep
|
||||
axis but differ in everything else.
|
||||
"""
|
||||
store_path = tmp_path / "aligned.zarr"
|
||||
root = zarr.open_group(str(store_path), mode="w")
|
||||
root.create_dataset(
|
||||
"img",
|
||||
data=np.arange(8 * 4 * 4 * 3, dtype="|u1").reshape(8, 4, 4, 3),
|
||||
chunks=(2, 4, 4, 3),
|
||||
)
|
||||
root.create_dataset(
|
||||
"state",
|
||||
data=np.arange(8 * 3, dtype="<f4").reshape(8, 3),
|
||||
chunks=(4, 3), # different native axis-0 chunks than img
|
||||
)
|
||||
root.create_dataset(
|
||||
"label",
|
||||
data=np.arange(8, dtype="<i8"),
|
||||
chunks=(8,),
|
||||
)
|
||||
zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
|
||||
return store_path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def zarr_zip_store(tmp_path) -> Path:
|
||||
"""A small Zarr store packed into a ``.zip`` for URL-detection tests."""
|
||||
src = tmp_path / "src.zarr"
|
||||
_write_real_zarr_store(
|
||||
src,
|
||||
{
|
||||
"data": (np.arange(12, dtype="<i4").reshape(6, 2), (3, 2)),
|
||||
},
|
||||
)
|
||||
zip_path = tmp_path / "store.zarr.zip"
|
||||
import shutil
|
||||
|
||||
shutil.make_archive(
|
||||
base_name=str(tmp_path / "store.zarr"),
|
||||
format="zip",
|
||||
root_dir=str(src),
|
||||
)
|
||||
assert zip_path.exists()
|
||||
return zip_path
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Metadata discovery
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_normalizes_requested_root_array_path(zarrv2_root_store):
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(zarrv2_root_store),
|
||||
array_paths=[""],
|
||||
)
|
||||
assert list(datasource._metadata_by_path) == [""]
|
||||
|
||||
|
||||
def test_normalizes_requested_array_paths(zarrv2_group_store):
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(zarrv2_group_store),
|
||||
array_paths=["images/", "nested"],
|
||||
)
|
||||
assert list(datasource._metadata_by_path) == ["images", "nested"]
|
||||
|
||||
|
||||
def test_rejects_missing_array_paths(zarrv2_group_store):
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=r"Array\(s\) not found: 'missing'\. Available: 'images', 'nested'",
|
||||
):
|
||||
zarrv2_datasource.ZarrV2Datasource(
|
||||
str(zarrv2_group_store),
|
||||
array_paths=["missing"],
|
||||
)
|
||||
|
||||
|
||||
def test_loads_per_array_zarray_without_zmetadata(unconsolidated_zarrv2_store):
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(unconsolidated_zarrv2_store),
|
||||
array_paths=["images", "nested"],
|
||||
)
|
||||
assert set(datasource._metadata_by_path) == {"images", "nested"}
|
||||
|
||||
|
||||
def test_full_scan_discovers_arrays_without_zmetadata(unconsolidated_zarrv2_store):
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(unconsolidated_zarrv2_store),
|
||||
allow_full_metadata_scan=True,
|
||||
)
|
||||
assert set(datasource._metadata_by_path) == {"images", "nested"}
|
||||
|
||||
|
||||
def test_requires_array_paths_or_full_scan_when_unconsolidated(
|
||||
unconsolidated_zarrv2_store,
|
||||
):
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
r"No array_paths were provided and this Zarr store does not "
|
||||
r"contain \.zmetadata"
|
||||
),
|
||||
):
|
||||
zarrv2_datasource.ZarrV2Datasource(str(unconsolidated_zarrv2_store))
|
||||
|
||||
|
||||
def test_array_paths_missing_zarray_file_raises_value_error(
|
||||
unconsolidated_zarrv2_store,
|
||||
):
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=r"Array path 'missing' not found",
|
||||
):
|
||||
zarrv2_datasource.ZarrV2Datasource(
|
||||
str(unconsolidated_zarrv2_store),
|
||||
array_paths=["missing"],
|
||||
)
|
||||
|
||||
|
||||
def test_local_scheme_pins_reads_to_driver_node(zarrv2_group_store):
|
||||
"""``local://`` stores can't be read distributed; plain/cloud paths can."""
|
||||
local = zarrv2_datasource.ZarrV2Datasource("local://" + str(zarrv2_group_store))
|
||||
assert local.supports_distributed_reads is False
|
||||
|
||||
plain = zarrv2_datasource.ZarrV2Datasource(str(zarrv2_group_store))
|
||||
assert plain.supports_distributed_reads is True
|
||||
|
||||
|
||||
def test_consolidation_detected_via_open_consolidated(
|
||||
zarrv2_group_store, unconsolidated_zarrv2_store
|
||||
):
|
||||
"""``_consolidated`` reflects whether ``.zmetadata`` actually opened."""
|
||||
consolidated = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(zarrv2_group_store), array_paths=["images"]
|
||||
)
|
||||
assert consolidated._consolidated is True
|
||||
|
||||
unconsolidated = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(unconsolidated_zarrv2_store), array_paths=["images"]
|
||||
)
|
||||
assert unconsolidated._consolidated is False
|
||||
|
||||
|
||||
def test_array_paths_rejects_group_path(tmp_path):
|
||||
"""Requesting a group path (not an array) on an unconsolidated store errors."""
|
||||
store_path = tmp_path / "withgroup.zarr"
|
||||
root = zarr.open_group(str(store_path), mode="w")
|
||||
grp = root.create_group("grp")
|
||||
grp.create_dataset("inner", data=np.arange(4, dtype="<i4"), chunks=(2,))
|
||||
# Not consolidated -> the per-array ``.zarray`` lookup path.
|
||||
with pytest.raises(ValueError, match="is a group, not an array"):
|
||||
zarrv2_datasource.ZarrV2Datasource(str(store_path), array_paths=["grp"])
|
||||
|
||||
|
||||
def test_root_array_rejects_non_root_array_paths(zarrv2_root_store):
|
||||
"""A single root-level array rejects array_paths that aren't the root ''."""
|
||||
with pytest.raises(ValueError, match="single root-level array"):
|
||||
zarrv2_datasource.ZarrV2Datasource(
|
||||
str(zarrv2_root_store), array_paths=["missing"]
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# chunk_shapes validation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"chunk_shapes, match",
|
||||
[
|
||||
("invalid", "positive integers"),
|
||||
({"images": 1}, "positive integers"),
|
||||
({"does_not_exist": [2]}, "Unknown array path"),
|
||||
],
|
||||
)
|
||||
def test_rejects_invalid_chunk_shapes(zarrv2_group_store, chunk_shapes, match):
|
||||
with pytest.raises(ValueError, match=match):
|
||||
zarrv2_datasource.ZarrV2Datasource(
|
||||
str(zarrv2_group_store), chunk_shapes=chunk_shapes
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"chunk_shapes,array_paths,expected",
|
||||
[
|
||||
# No chunk_shapes: every array reads at its native chunk size.
|
||||
# 4-D image with tiny chunks coexists with 2-D pose with big chunks —
|
||||
# nothing is forced into a shared min/max.
|
||||
(
|
||||
None,
|
||||
None,
|
||||
{
|
||||
"data/camera0_rgb": (1, 2, 2, 3),
|
||||
"data/robot0_eef_pos": (10, 3),
|
||||
"meta/episode_ends": (3,),
|
||||
},
|
||||
),
|
||||
# ``[5]`` prefix overrides axis 0 across arrays of all ranks at once.
|
||||
(
|
||||
[5],
|
||||
None,
|
||||
{
|
||||
"data/camera0_rgb": (5, 2, 2, 3),
|
||||
"data/robot0_eef_pos": (5, 3),
|
||||
"meta/episode_ends": (5,),
|
||||
},
|
||||
),
|
||||
# Length-2 prefix overrides axes 0+1; needs every selected array to
|
||||
# have rank >= 2, so we filter out ``meta/episode_ends`` (rank 1).
|
||||
(
|
||||
[5, 1],
|
||||
["data/camera0_rgb", "data/robot0_eef_pos"],
|
||||
{
|
||||
"data/camera0_rgb": (5, 1, 2, 3),
|
||||
"data/robot0_eef_pos": (5, 1),
|
||||
},
|
||||
),
|
||||
# Per-array overrides may retile only some arrays while others keep
|
||||
# their native chunks.
|
||||
(
|
||||
{
|
||||
"data/camera0_rgb": [5],
|
||||
"data/robot0_eef_pos": [7],
|
||||
},
|
||||
None,
|
||||
{
|
||||
"data/camera0_rgb": (5, 2, 2, 3),
|
||||
"data/robot0_eef_pos": (7, 3),
|
||||
"meta/episode_ends": (3,),
|
||||
},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_chunk_shapes_resolution_across_mixed_rank(
|
||||
heterogeneous_zarrv2_store, chunk_shapes, array_paths, expected
|
||||
):
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(heterogeneous_zarrv2_store),
|
||||
chunk_shapes=chunk_shapes,
|
||||
array_paths=array_paths,
|
||||
)
|
||||
assert datasource._array_chunks == expected
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# align_axis_0 (wide-form mode)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_align_axis_0_emits_wide_rows(ray_start_regular_shared, aligned_zarrv2_store):
|
||||
"""Wide-row schema: ``t_start``, ``t_stop``, one column per selected array."""
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(aligned_zarrv2_store),
|
||||
align_axis_0=True,
|
||||
chunk_shapes=[4],
|
||||
)
|
||||
df = _execute_read_tasks(datasource.get_read_tasks(parallelism=4))
|
||||
assert set(df.columns) == {"t_start", "t_stop", "img", "state", "label"}
|
||||
# shape[0]=8, chunk_shapes=[4] -> 2 rows.
|
||||
assert len(df) == 2
|
||||
# Reconstruct each array by concatenating slices in order.
|
||||
img_recon = np.concatenate(list(df["img"]), axis=0)
|
||||
assert img_recon.shape == (8, 4, 4, 3)
|
||||
state_recon = np.concatenate(list(df["state"]), axis=0)
|
||||
assert state_recon.shape == (8, 3)
|
||||
label_recon = np.concatenate(list(df["label"]), axis=0)
|
||||
assert label_recon.shape == (8,)
|
||||
# t_start/t_stop are correct.
|
||||
starts = sorted(df["t_start"].tolist())
|
||||
stops = sorted(df["t_stop"].tolist())
|
||||
assert starts == [0, 4]
|
||||
assert stops == [4, 8]
|
||||
|
||||
|
||||
def test_align_axis_0_column_set(ray_start_regular_shared, aligned_zarrv2_store):
|
||||
"""array_paths selects which arrays are read; aligned mode emits one column
|
||||
per selected array (plus t_start/t_stop)."""
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(aligned_zarrv2_store),
|
||||
array_paths=["img", "state"],
|
||||
align_axis_0=True,
|
||||
chunk_shapes=[4],
|
||||
)
|
||||
df = _execute_read_tasks(datasource.get_read_tasks(parallelism=4))
|
||||
assert set(df.columns) == {"t_start", "t_stop", "img", "state"}
|
||||
|
||||
|
||||
def test_align_axis_0_rejects_misaligned_shape0(heterogeneous_zarrv2_store):
|
||||
"""Misalignment raises with the per-array shape[0] breakdown."""
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=r"All selected arrays must share shape\[0\]",
|
||||
):
|
||||
zarrv2_datasource.ZarrV2Datasource(
|
||||
str(heterogeneous_zarrv2_store),
|
||||
align_axis_0=True,
|
||||
chunk_shapes=[5],
|
||||
)
|
||||
|
||||
|
||||
def test_align_axis_0_rejects_divergent_axis_0_chunks(aligned_zarrv2_store):
|
||||
"""If aligned arrays end up with different axis-0 chunks, error clearly.
|
||||
|
||||
The native chunks differ (img=2, state=4, label=8) — without a
|
||||
``chunk_shapes`` re-tile they all stay at native, and the validator
|
||||
catches the mismatch.
|
||||
"""
|
||||
with pytest.raises(
|
||||
ValueError, match="Aligned arrays must share the same axis-0 chunk size"
|
||||
):
|
||||
zarrv2_datasource.ZarrV2Datasource(
|
||||
str(aligned_zarrv2_store),
|
||||
align_axis_0=True,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# overlap (aligned-mode lookahead)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_overlap_extends_chunk_data(ray_start_regular_shared, aligned_zarrv2_store):
|
||||
"""``overlap=N`` makes each row's per-array slice cover ``N`` extra timesteps.
|
||||
|
||||
Aligned store has shape[0]=8, ``chunk_shapes=[4]`` -> rows own [0,4) and [4,8).
|
||||
With ``overlap=2``, row 0's data covers [0,6) and row 1's data covers [4,8)
|
||||
(clipped at the store end since 4+4+2 > 8).
|
||||
"""
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(aligned_zarrv2_store),
|
||||
align_axis_0=True,
|
||||
chunk_shapes=[4],
|
||||
overlap=2,
|
||||
)
|
||||
df = _execute_read_tasks(datasource.get_read_tasks(parallelism=4))
|
||||
# Ownership unchanged: 2 rows of width 4 each.
|
||||
assert sorted(zip(df["t_start"], df["t_stop"])) == [(0, 4), (4, 8)]
|
||||
# Data extents: row 0 has 6 timesteps, row 1 has 4 (clipped at shape[0]=8).
|
||||
rows = sorted(df.to_dict("records"), key=lambda r: r["t_start"])
|
||||
assert rows[0]["img"].shape[0] == 6 # 4 owned + 2 overlap
|
||||
assert rows[0]["state"].shape[0] == 6
|
||||
assert rows[1]["img"].shape[0] == 4 # 4 owned + 0 overlap (clipped)
|
||||
assert rows[1]["state"].shape[0] == 4
|
||||
|
||||
|
||||
def test_overlap_requires_align_axis_0(aligned_zarrv2_store):
|
||||
"""``overlap`` in long-form (no ``align_axis_0``) is a clear error."""
|
||||
with pytest.raises(ValueError, match="overlap requires align_axis_0=True"):
|
||||
zarrv2_datasource.ZarrV2Datasource(
|
||||
str(aligned_zarrv2_store),
|
||||
overlap=2,
|
||||
)
|
||||
|
||||
|
||||
def test_overlap_rejects_negative_and_non_int(aligned_zarrv2_store):
|
||||
bad_values: list[Any] = [-1, 1.5, "two"]
|
||||
|
||||
for bad in bad_values:
|
||||
with pytest.raises(ValueError, match="overlap must be a non-negative integer"):
|
||||
zarrv2_datasource.ZarrV2Datasource(
|
||||
str(aligned_zarrv2_store),
|
||||
align_axis_0=True,
|
||||
chunk_shapes=[4],
|
||||
overlap=bad,
|
||||
)
|
||||
|
||||
|
||||
def test_chunk_shapes_rejected_when_longer_than_smallest_array(
|
||||
heterogeneous_zarrv2_store,
|
||||
):
|
||||
"""A shared ``chunk_shapes`` override longer than a target rank is an error."""
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=r"chunk_shapes override for array .* has 2 axes but array of shape .* has rank 1",
|
||||
):
|
||||
zarrv2_datasource.ZarrV2Datasource(
|
||||
str(heterogeneous_zarrv2_store),
|
||||
chunk_shapes=[2, 2], # OK for 2-D and 4-D, fails for 1-D episode_ends
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Filesystem handling
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_accepts_pyarrow_fs_filesystem(zarrv2_group_store):
|
||||
"""A pyarrow.fs.FileSystem passed in is wrapped into fsspec internally."""
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(zarrv2_group_store),
|
||||
filesystem=pyarrow.fs.LocalFileSystem(),
|
||||
)
|
||||
from fsspec.spec import AbstractFileSystem
|
||||
|
||||
assert isinstance(datasource._fs, AbstractFileSystem)
|
||||
assert set(datasource._metadata_by_path) == {"images", "nested"}
|
||||
|
||||
|
||||
def test_rejects_unsupported_filesystem_type():
|
||||
"""Filesystem that's neither pyarrow.fs nor fsspec raises ``TypeError``."""
|
||||
with pytest.raises(
|
||||
TypeError,
|
||||
match=r"filesystem must be pyarrow\.fs\.FileSystem or",
|
||||
):
|
||||
zarrv2_datasource.ZarrV2Datasource(
|
||||
"/tmp/some.zarr",
|
||||
filesystem="not-a-filesystem",
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# .zarr.zip URL support
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_reads_zarr_zip_local_path(ray_start_regular_shared, zarr_zip_store):
|
||||
"""A local ``.zarr.zip`` path auto-wires fsspec's ZipFileSystem."""
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(str(zarr_zip_store))
|
||||
# The store has one array "data" of shape (6, 2) chunks (3, 2) -> 2 chunks.
|
||||
df = _execute_read_tasks(datasource.get_read_tasks(parallelism=2))
|
||||
assert len(df) == 2
|
||||
assert set(df["array"]) == {"data"}
|
||||
recon = _reconstruct_array(df, "data")
|
||||
np.testing.assert_array_equal(recon, np.arange(12, dtype="<i4").reshape(6, 2))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Read task generation and execution (end-to-end)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_get_read_tasks_batches_chunks_by_parallelism(tmp_path):
|
||||
"""5 chunks split across parallelism=3 produces batches [2, 2, 1]."""
|
||||
store_path = tmp_path / "store.zarr"
|
||||
_write_real_zarr_store(
|
||||
store_path,
|
||||
{"images": (np.arange(5 * 4, dtype="<i4").reshape(5, 4), (1, 4))},
|
||||
)
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path))
|
||||
|
||||
read_tasks = datasource.get_read_tasks(parallelism=3)
|
||||
|
||||
assert len(read_tasks) == 3
|
||||
assert [task.metadata.num_rows for task in read_tasks] == [2, 2, 1]
|
||||
assert all(task.metadata.input_files == (str(store_path),) for task in read_tasks)
|
||||
|
||||
|
||||
def test_long_form_chunk_index_order_matches_grid(ray_start_regular_shared, tmp_path):
|
||||
"""Lazy grid-range tasks emit chunk_index in the same row-major order as a
|
||||
full grid enumeration (regression guard for the lazy-unravel refactor)."""
|
||||
from itertools import product
|
||||
|
||||
store_path = tmp_path / "order.zarr"
|
||||
# shape (6, 4), chunks (2, 2) -> grid (3, 2) = 6 chunks.
|
||||
_write_real_zarr_store(
|
||||
store_path, {"a": (np.arange(6 * 4, dtype="<i4").reshape(6, 4), (2, 2))}
|
||||
)
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path))
|
||||
# parallelism=2 -> two flat-index ranges; concatenated they must be in order.
|
||||
df = _execute_read_tasks(datasource.get_read_tasks(parallelism=2))
|
||||
got = [tuple(int(x) for x in ci) for ci in df["chunk_index"]]
|
||||
assert got == list(product(range(3), range(2)))
|
||||
|
||||
|
||||
def test_per_task_row_limit_caps_chunks_read(
|
||||
ray_start_regular_shared, tmp_path, monkeypatch
|
||||
):
|
||||
"""per_task_row_limit bounds how many chunks a task actually reads, so a
|
||||
downstream ``limit`` doesn't pull the whole batch's I/O."""
|
||||
store_path = tmp_path / "limit.zarr"
|
||||
_write_real_zarr_store(store_path, {"data": (np.arange(10, dtype="<i4"), (1,))})
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path))
|
||||
|
||||
reads = []
|
||||
real_read_chunk = zarrv2_datasource._read_chunk
|
||||
|
||||
def _spy(*args, **kwargs):
|
||||
reads.append(1)
|
||||
return real_read_chunk(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(zarrv2_datasource, "_read_chunk", _spy)
|
||||
|
||||
# parallelism=1 -> one task batching all 10 chunks; cap it at 3.
|
||||
tasks = datasource.get_read_tasks(parallelism=1, per_task_row_limit=3)
|
||||
blocks = [block for task in tasks for block in task()]
|
||||
|
||||
total_rows = sum(BlockAccessor.for_block(b).num_rows() for b in blocks)
|
||||
assert total_rows == 3
|
||||
# The fix: only 3 chunks were actually read (not all 10, then truncated).
|
||||
assert len(reads) == 3
|
||||
|
||||
|
||||
def test_read_chunk_retries_transient_io(monkeypatch):
|
||||
"""_read_chunk retries reads whose error matches retry_match (Ray Data's
|
||||
DataContext.retried_io_errors), then succeeds."""
|
||||
monkeypatch.setattr("time.sleep", lambda *_: None) # no backoff in the test
|
||||
|
||||
class _FlakyArray:
|
||||
attempts = 0
|
||||
|
||||
def __getitem__(self, _idx):
|
||||
type(self).attempts += 1
|
||||
if self.attempts < 3:
|
||||
raise OSError("Connection reset by peer")
|
||||
return np.arange(4, dtype="<i4")
|
||||
|
||||
class _Root:
|
||||
def __getitem__(self, _name):
|
||||
return _FlakyArray()
|
||||
|
||||
out = zarrv2_datasource._read_chunk(
|
||||
_Root(), # pyrefly: ignore[bad-argument-type]
|
||||
"x",
|
||||
((0, 4),),
|
||||
retry_match=["Connection reset"],
|
||||
)
|
||||
np.testing.assert_array_equal(out, np.arange(4, dtype="<i4"))
|
||||
assert _FlakyArray.attempts == 3 # failed twice, then succeeded
|
||||
|
||||
|
||||
def test_long_form_schema_and_materialization(ray_start_regular_shared, tmp_path):
|
||||
"""End-to-end: long-form rows are emitted with the expected columns and data."""
|
||||
store_path = tmp_path / "aligned.zarr"
|
||||
images_src = np.arange(20, dtype="<i4").reshape(5, 4)
|
||||
labels_src = np.arange(5, dtype="|u1")
|
||||
_write_real_zarr_store(
|
||||
store_path,
|
||||
{
|
||||
"images": (images_src, (2, 4)),
|
||||
"labels": (labels_src, (2,)),
|
||||
},
|
||||
)
|
||||
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path))
|
||||
df = _execute_read_tasks(datasource.get_read_tasks(parallelism=16))
|
||||
|
||||
# Schema is the long-form quad.
|
||||
assert list(df.columns) == ["array", "chunk_index", "chunk_slices", "chunk"]
|
||||
# 3 chunks for images (5/2), 3 chunks for labels (5/2) = 6 rows total.
|
||||
assert len(df) == 6
|
||||
assert set(df["array"]) == {"images", "labels"}
|
||||
|
||||
np.testing.assert_array_equal(_reconstruct_array(df, "images"), images_src)
|
||||
np.testing.assert_array_equal(_reconstruct_array(df, "labels"), labels_src)
|
||||
|
||||
# ``chunk_slices`` matches the actual chunk shape and indexes back to
|
||||
# the source array: arr[start:stop, ...] equals the chunk.
|
||||
for _, row in df.iterrows():
|
||||
slices = row["chunk_slices"]
|
||||
chunk = row["chunk"]
|
||||
assert len(slices) == chunk.ndim
|
||||
for axis, (start, stop) in enumerate(slices):
|
||||
assert stop - start == chunk.shape[axis]
|
||||
if row["array"] == "images":
|
||||
np.testing.assert_array_equal(
|
||||
chunk,
|
||||
images_src[slices[0][0] : slices[0][1], slices[1][0] : slices[1][1]],
|
||||
)
|
||||
|
||||
|
||||
def test_chunk_shapes_override_changes_grid(ray_start_regular_shared, tmp_path):
|
||||
"""User-supplied chunk_shapes controls the chunk grid and row count."""
|
||||
store_path = tmp_path / "tile.zarr"
|
||||
src = np.arange(10, dtype="<i4")
|
||||
_write_real_zarr_store(store_path, {"data": (src, (2,))}) # native: 5 chunks
|
||||
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path), chunk_shapes=[5])
|
||||
df = _execute_read_tasks(datasource.get_read_tasks(parallelism=16))
|
||||
assert sorted(chunk.shape[0] for chunk in df["chunk"]) == [5, 5]
|
||||
|
||||
|
||||
def test_heterogeneous_store_emits_one_row_per_chunk(
|
||||
ray_start_regular_shared, heterogeneous_zarrv2_store
|
||||
):
|
||||
"""Mixed-rank/shape/dtype arrays each contribute their chunk count to the output."""
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(str(heterogeneous_zarrv2_store))
|
||||
df = _execute_read_tasks(datasource.get_read_tasks(parallelism=16))
|
||||
|
||||
# Expected chunk counts:
|
||||
# data/camera0_rgb shape=(20,2,2,3) chunks=(1,2,2,3) → 20 chunks
|
||||
# data/robot0_eef_pos shape=(20,3) chunks=(10,3) → 2 chunks
|
||||
# meta/episode_ends shape=(3,) chunks=(3,) → 1 chunk
|
||||
counts = df.groupby("array").size().to_dict()
|
||||
assert counts == {
|
||||
"data/camera0_rgb": 20,
|
||||
"data/robot0_eef_pos": 2,
|
||||
"meta/episode_ends": 1,
|
||||
}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Estimator
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_estimate_inmemory_data_size(tmp_path):
|
||||
"""Estimate = sum over arrays of numel * dtype.itemsize."""
|
||||
store_path = tmp_path / "est.zarr"
|
||||
_write_real_zarr_store(
|
||||
store_path,
|
||||
{
|
||||
"a": (np.zeros((5, 4), dtype="<i4"), (2, 4)),
|
||||
"b": (np.zeros(5, dtype="|u1"), (2,)),
|
||||
},
|
||||
)
|
||||
datasource = zarrv2_datasource.ZarrV2Datasource(str(store_path))
|
||||
# 5*4*4 (a) + 5*1 (b) = 80 + 5 = 85
|
||||
assert datasource.estimate_inmemory_data_size() == 85
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Cross-filesystem end-to-end (Ray Data convention)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"fs",
|
||||
[
|
||||
None,
|
||||
lazy_fixture("local_fs"), # pyarrow.fs (gets wrapped to fsspec internally)
|
||||
lazy_fixture("local_fsspec_fs"), # native fsspec
|
||||
],
|
||||
)
|
||||
def test_read_zarr_basic_across_filesystems(ray_start_regular_shared, fs, local_path):
|
||||
"""Round-trip a real Zarr store through read_zarr for each filesystem flavor.
|
||||
|
||||
Mirrors the parametrized read-path coverage other Ray Data datasources use
|
||||
(lance, parquet, json, hudi, …) — exercises None / pyarrow.fs / fsspec
|
||||
input shapes against the same store written to a local path.
|
||||
"""
|
||||
store_path = os.path.join(local_path, "data.zarr")
|
||||
images_src = np.arange(20, dtype="<i4").reshape(5, 4)
|
||||
labels_src = np.arange(5, dtype="|u1")
|
||||
_write_real_zarr_store(
|
||||
Path(store_path),
|
||||
{
|
||||
"images": (images_src, (2, 4)),
|
||||
"labels": (labels_src, (2,)),
|
||||
},
|
||||
)
|
||||
|
||||
ds = ray.data.read_zarr(store_path, filesystem=fs)
|
||||
|
||||
# 3 chunks each for images and labels (5/2 → ceil = 3) → 6 rows total.
|
||||
assert ds.count() == 6
|
||||
df = pd.DataFrame(ds.take_all())
|
||||
np.testing.assert_array_equal(_reconstruct_array(df, "images"), images_src)
|
||||
np.testing.assert_array_equal(_reconstruct_array(df, "labels"), labels_src)
|
||||
|
||||
|
||||
def test_rejects_zarr_v3(tmp_path, monkeypatch):
|
||||
"""read_zarr targets zarr-python 2.x; an incompatible v3 install must raise a
|
||||
clear, actionable error at construction, not a cryptic ImportError mid-read."""
|
||||
monkeypatch.setattr(zarr, "__version__", "3.0.1")
|
||||
with pytest.raises(ImportError, match=r"zarr-python 2\.x"):
|
||||
zarrv2_datasource.ZarrV2Datasource(str(tmp_path))
|
||||
|
||||
|
||||
def test_explicit_filesystem_strips_uri_scheme(ray_start_regular_shared, tmp_path):
|
||||
"""An explicit ``filesystem=`` plus a scheme-prefixed path must strip the
|
||||
scheme so the store path is backend-relative. Regression: pyarrow
|
||||
filesystems can't resolve a ``file://`` / ``gs://`` prefix in the path."""
|
||||
store_path = tmp_path / "scheme.zarr"
|
||||
_write_real_zarr_store(store_path, {"data": (np.arange(6, dtype="<i4"), (2,))})
|
||||
|
||||
ds = zarrv2_datasource.ZarrV2Datasource(
|
||||
f"file://{store_path}", filesystem=pyarrow.fs.LocalFileSystem()
|
||||
)
|
||||
assert ds._store_path == str(store_path)
|
||||
df = _execute_read_tasks(ds.get_read_tasks(parallelism=2))
|
||||
assert len(df) == 3
|
||||
|
||||
|
||||
def test_get_read_tasks_parallelism_zero(tmp_path):
|
||||
"""parallelism=0 must not divide by zero; fall back to a single task."""
|
||||
store_path = tmp_path / "p0.zarr"
|
||||
_write_real_zarr_store(store_path, {"data": (np.arange(10, dtype="<i4"), (2,))})
|
||||
ds = zarrv2_datasource.ZarrV2Datasource(str(store_path))
|
||||
tasks = ds.get_read_tasks(parallelism=0)
|
||||
assert len(tasks) >= 1
|
||||
|
||||
|
||||
def test_align_axis_0_rejects_scalar_array(tmp_path):
|
||||
"""align_axis_0=True with a 0-D (scalar) array must raise a clear error
|
||||
rather than an IndexError when reading the (empty) axis-0 chunk size."""
|
||||
store_path = tmp_path / "scalar.zarr"
|
||||
root = zarr.open_group(str(store_path), mode="w")
|
||||
root.create_dataset("vec", data=np.arange(8, dtype="<i4"), chunks=(4,))
|
||||
root.create_dataset("scalar", data=np.array(42, dtype="<i4")) # 0-D
|
||||
zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
|
||||
|
||||
with pytest.raises(ValueError, match=r"0-D \(scalar\)"):
|
||||
zarrv2_datasource.ZarrV2Datasource(str(store_path), align_axis_0=True)
|
||||
|
||||
|
||||
def test_reads_zarr_zip_with_explicit_zip_filesystem(
|
||||
ray_start_regular_shared, zarr_zip_store
|
||||
):
|
||||
"""A .zip path read through an explicitly-passed fsspec ZipFileSystem must
|
||||
resolve the store at the archive root (store path ``""``), not treat the
|
||||
``.zip`` name as an entry inside the archive."""
|
||||
zip_fs = fsspec.filesystem("zip", fo=str(zarr_zip_store))
|
||||
ds = zarrv2_datasource.ZarrV2Datasource(str(zarr_zip_store), filesystem=zip_fs)
|
||||
assert ds._store_path == ""
|
||||
df = _execute_read_tasks(ds.get_read_tasks(parallelism=2))
|
||||
assert len(df) == 2
|
||||
|
||||
|
||||
def test_align_axis_0_columns_unify_across_blocks(
|
||||
ray_start_regular_shared, aligned_zarrv2_store
|
||||
):
|
||||
"""Wide-form gives each array its own column, so blocks combine cleanly
|
||||
across the dataset even with trailing edge chunks of differing shape -- the
|
||||
batch-safe schema for row-aligned arrays."""
|
||||
from ray.data._internal.arrow_ops.transform_pyarrow import unify_schemas
|
||||
from ray.data.block import BlockAccessor
|
||||
|
||||
ds = zarrv2_datasource.ZarrV2Datasource(
|
||||
str(aligned_zarrv2_store), align_axis_0=True, chunk_shapes=[3]
|
||||
)
|
||||
blocks = [block for task in ds.get_read_tasks(parallelism=64) for block in task()]
|
||||
assert len(blocks) > 1 # actually exercise cross-block unification
|
||||
schemas = [BlockAccessor.for_block(b).to_arrow().schema for b in blocks]
|
||||
unified = unify_schemas(schemas) # must not raise
|
||||
assert {"t_start", "t_stop", "img", "state", "label"}.issubset(set(unified.names))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Custom codec registration in Ray workers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fresh_ray():
|
||||
"""A clean Ray for a test that needs its own ``ray.init`` (e.g. a custom
|
||||
``runtime_env``). Unlike ``shutdown_only`` (teardown only), it also shuts
|
||||
down any pre-existing cluster, so isolation doesn't depend on test order.
|
||||
"""
|
||||
if ray.is_initialized():
|
||||
ray.shutdown()
|
||||
yield
|
||||
if ray.is_initialized():
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
def test_custom_codec_succeeds_with_worker_setup_hook(fresh_ray, tmp_path):
|
||||
"""Test that we successfully register a custom codec.
|
||||
|
||||
numcodecs' registry is process-local.
|
||||
"""
|
||||
import numcodecs
|
||||
|
||||
def _register_codec():
|
||||
import numcodecs
|
||||
import numpy as np
|
||||
|
||||
class _RayZarrTestCodec(numcodecs.abc.Codec):
|
||||
codec_id = "ray_zarr_test_codec"
|
||||
|
||||
def encode(self, buf):
|
||||
return bytes(buf)
|
||||
|
||||
def decode(self, buf, out=None):
|
||||
arr = np.frombuffer(buf, dtype=np.uint8)
|
||||
if out is not None:
|
||||
out[:] = arr.view(out.dtype)
|
||||
return out
|
||||
return arr.copy()
|
||||
|
||||
numcodecs.register_codec(_RayZarrTestCodec)
|
||||
|
||||
# Register driver-side so we can write the store.
|
||||
_register_codec()
|
||||
|
||||
store_path = tmp_path / "codec_test.zarr"
|
||||
arr = zarr.open(
|
||||
str(store_path),
|
||||
mode="w",
|
||||
shape=(8,),
|
||||
chunks=(4,),
|
||||
dtype="u1",
|
||||
compressor=numcodecs.get_codec({"id": "ray_zarr_test_codec"}),
|
||||
)
|
||||
arr[:] = np.arange(8, dtype="u1")
|
||||
zarr.consolidate_metadata(zarr.DirectoryStore(str(store_path)))
|
||||
|
||||
ray.init(
|
||||
num_cpus=1,
|
||||
logging_level=logging.ERROR,
|
||||
log_to_driver=False,
|
||||
runtime_env={"worker_process_setup_hook": _register_codec},
|
||||
)
|
||||
ds = ray.data.read_zarr(str(store_path))
|
||||
rows = sorted(ds.take_all(), key=lambda r: tuple(r["chunk_index"]))
|
||||
recon = np.concatenate([r["chunk"] for r in rows])
|
||||
np.testing.assert_array_equal(recon, np.arange(8, dtype="u1"))
|
||||
@@ -0,0 +1,25 @@
|
||||
"""This file is injected for Ray Data doctest targets."""
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True, scope="module")
|
||||
def shutdown_ray():
|
||||
ray.shutdown()
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def preserve_block_order():
|
||||
ray.data.context.DataContext.get_current().execution_options.preserve_order = True
|
||||
yield
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def disable_start_message():
|
||||
context = ray.data.context.DataContext.get_current()
|
||||
original_value = context.print_on_execution_start
|
||||
context.print_on_execution_start = False
|
||||
yield
|
||||
context.print_on_execution_start = original_value
|
||||
@@ -0,0 +1,209 @@
|
||||
"""Integration tests for arithmetic expression operations.
|
||||
|
||||
These tests require Ray and test end-to-end arithmetic expression evaluation.
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import pandas as pd
|
||||
import pytest
|
||||
from packaging.version import parse as parse_version
|
||||
|
||||
import ray
|
||||
from ray.data._internal.util import rows_same
|
||||
from ray.data._internal.utils.arrow_utils import get_pyarrow_version
|
||||
from ray.data.expressions import col, lit
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
get_pyarrow_version() < parse_version("20.0.0"),
|
||||
reason="Expression integration tests require PyArrow >= 20.0.0",
|
||||
)
|
||||
|
||||
|
||||
class TestArithmeticIntegration:
|
||||
"""Integration tests for arithmetic expressions with Ray Dataset."""
|
||||
|
||||
def test_arithmetic_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test arithmetic expressions work correctly with Ray Dataset."""
|
||||
ds = ray.data.from_items(
|
||||
[
|
||||
{"price": 10.0, "quantity": 2},
|
||||
{"price": 20.0, "quantity": 3},
|
||||
{"price": 15.0, "quantity": 4},
|
||||
]
|
||||
)
|
||||
|
||||
result = ds.with_column("total", col("price") * col("quantity")).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"price": [10.0, 20.0, 15.0],
|
||||
"quantity": [2, 3, 4],
|
||||
"total": [20.0, 60.0, 60.0],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_chained_arithmetic_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test chained arithmetic expressions with Ray Dataset."""
|
||||
ds = ray.data.from_items(
|
||||
[
|
||||
{"a": 10, "b": 5},
|
||||
{"a": 20, "b": 3},
|
||||
]
|
||||
)
|
||||
|
||||
result = (
|
||||
ds.with_column("sum", col("a") + col("b"))
|
||||
.with_column("diff", col("a") - col("b"))
|
||||
.with_column("product", col("a") * col("b"))
|
||||
.to_pandas()
|
||||
)
|
||||
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"a": [10, 20],
|
||||
"b": [5, 3],
|
||||
"sum": [15, 23],
|
||||
"diff": [5, 17],
|
||||
"product": [50, 60],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_floor_division_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test floor division operations with Ray Dataset."""
|
||||
ds = ray.data.range(5)
|
||||
result = ds.with_column("result", col("id") // 2).to_pandas()
|
||||
expected = pd.DataFrame({"id": [0, 1, 2, 3, 4], "result": [0, 0, 1, 1, 2]})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_literal_floor_division_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test literal floor division by expression with Ray Dataset."""
|
||||
ds = ray.data.range(5)
|
||||
result = ds.with_column("result", lit(10) // (col("id") + 2)).to_pandas()
|
||||
expected = pd.DataFrame({"id": [0, 1, 2, 3, 4], "result": [5, 3, 2, 2, 1]})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"expr_factory,expected_values",
|
||||
[
|
||||
pytest.param(lambda: col("value").ceil(), [-1, 0, 0, 1, 2], id="ceil"),
|
||||
pytest.param(lambda: col("value").floor(), [-2, -1, 0, 0, 1], id="floor"),
|
||||
pytest.param(lambda: col("value").round(), [-2, 0, 0, 0, 2], id="round"),
|
||||
pytest.param(lambda: col("value").trunc(), [-1, 0, 0, 0, 1], id="trunc"),
|
||||
],
|
||||
)
|
||||
def test_rounding_with_dataset(
|
||||
self, ray_start_regular_shared, expr_factory, expected_values
|
||||
):
|
||||
"""Test rounding operations with Ray Dataset."""
|
||||
values = [-1.75, -0.25, 0.0, 0.25, 1.75]
|
||||
ds = ray.data.from_items([{"value": v} for v in values])
|
||||
result = ds.with_column("result", expr_factory()).to_pandas()
|
||||
expected = pd.DataFrame({"value": values, "result": expected_values})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"expr_factory,expected_fn",
|
||||
[
|
||||
pytest.param(lambda: col("value").ln(), math.log, id="ln"),
|
||||
pytest.param(lambda: col("value").log10(), math.log10, id="log10"),
|
||||
pytest.param(lambda: col("value").log2(), math.log2, id="log2"),
|
||||
pytest.param(lambda: col("value").exp(), math.exp, id="exp"),
|
||||
],
|
||||
)
|
||||
def test_logarithmic_with_dataset(
|
||||
self, ray_start_regular_shared, expr_factory, expected_fn
|
||||
):
|
||||
"""Test logarithmic operations with Ray Dataset."""
|
||||
values = [1.0, math.e, 10.0, 4.0]
|
||||
ds = ray.data.from_items([{"value": v} for v in values])
|
||||
expected_values = [expected_fn(v) for v in values]
|
||||
result = ds.with_column("result", expr_factory()).to_pandas()
|
||||
expected = pd.DataFrame({"value": values, "result": expected_values})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"expr_factory,expected_fn",
|
||||
[
|
||||
pytest.param(lambda: col("value").sin(), math.sin, id="sin"),
|
||||
pytest.param(lambda: col("value").cos(), math.cos, id="cos"),
|
||||
pytest.param(lambda: col("value").tan(), math.tan, id="tan"),
|
||||
pytest.param(lambda: col("value").atan(), math.atan, id="atan"),
|
||||
],
|
||||
)
|
||||
def test_trigonometric_with_dataset(
|
||||
self, ray_start_regular_shared, expr_factory, expected_fn
|
||||
):
|
||||
"""Test trigonometric operations with Ray Dataset."""
|
||||
values = [0.0, math.pi / 6, math.pi / 4, math.pi / 3]
|
||||
ds = ray.data.from_items([{"value": v} for v in values])
|
||||
expected_values = [expected_fn(v) for v in values]
|
||||
result = ds.with_column("result", expr_factory()).to_pandas()
|
||||
expected = pd.DataFrame({"value": values, "result": expected_values})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"test_data,expr_factory,expected_results",
|
||||
[
|
||||
pytest.param(
|
||||
[{"x": 5}, {"x": -3}, {"x": 0}],
|
||||
lambda: col("x").negate(),
|
||||
[-5, 3, 0],
|
||||
id="negate",
|
||||
),
|
||||
pytest.param(
|
||||
[{"x": 5}, {"x": -3}, {"x": 0}],
|
||||
lambda: col("x").sign(),
|
||||
[1, -1, 0],
|
||||
id="sign",
|
||||
),
|
||||
pytest.param(
|
||||
[{"x": 5}, {"x": -3}, {"x": 0}],
|
||||
lambda: col("x").abs(),
|
||||
[5, 3, 0],
|
||||
id="abs",
|
||||
),
|
||||
pytest.param(
|
||||
[{"x": 2}, {"x": 3}, {"x": 4}],
|
||||
lambda: col("x").power(2),
|
||||
[4, 9, 16],
|
||||
id="power_int",
|
||||
),
|
||||
pytest.param(
|
||||
[{"x": 4}, {"x": 9}, {"x": 16}],
|
||||
lambda: col("x").power(0.5),
|
||||
[2.0, 3.0, 4.0],
|
||||
id="power_sqrt",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_arithmetic_helpers_with_dataset(
|
||||
self, ray_start_regular_shared, test_data, expr_factory, expected_results
|
||||
):
|
||||
"""Test arithmetic helper operations with Ray Dataset."""
|
||||
ds = ray.data.from_items(test_data)
|
||||
result = ds.with_column("result", expr_factory()).to_pandas()
|
||||
expected = pd.DataFrame(test_data)
|
||||
expected["result"] = expected_results
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_age_group_calculation_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test floor division for grouping values (e.g., age into decades)."""
|
||||
test_data = [
|
||||
{"age": 25},
|
||||
{"age": 17},
|
||||
{"age": 30},
|
||||
]
|
||||
ds = ray.data.from_items(test_data)
|
||||
result = ds.with_column("age_group", col("age") // 10 * 10).to_pandas()
|
||||
expected = pd.DataFrame({"age": [25, 17, 30], "age_group": [20, 10, 30]})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,147 @@
|
||||
"""Integration tests for boolean/logical expression operations.
|
||||
|
||||
These tests require Ray and test end-to-end boolean expression evaluation.
|
||||
"""
|
||||
|
||||
import pandas as pd
|
||||
import pytest
|
||||
from packaging.version import parse as parse_version
|
||||
|
||||
import ray
|
||||
from ray.data._internal.utils.arrow_utils import get_pyarrow_version
|
||||
from ray.data.expressions import col, lit
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
get_pyarrow_version() < parse_version("20.0.0"),
|
||||
reason="Expression integration tests require PyArrow >= 20.0.0",
|
||||
)
|
||||
|
||||
|
||||
class TestBooleanIntegration:
|
||||
"""Integration tests for boolean expressions with Ray Dataset."""
|
||||
|
||||
def test_boolean_filter_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test boolean expressions used for filtering with Ray Dataset."""
|
||||
ds = ray.data.from_items(
|
||||
[
|
||||
{"age": 17, "is_student": True, "name": "Alice"},
|
||||
{"age": 21, "is_student": True, "name": "Bob"},
|
||||
{"age": 25, "is_student": False, "name": "Charlie"},
|
||||
{"age": 30, "is_student": False, "name": "Diana"},
|
||||
]
|
||||
)
|
||||
|
||||
# Add boolean columns using expressions
|
||||
result = (
|
||||
ds.with_column("is_adult", col("age") >= 18)
|
||||
.with_column("adult_student", (col("age") >= 18) & col("is_student"))
|
||||
.with_column("minor_or_student", (col("age") < 18) | col("is_student"))
|
||||
.to_pandas()
|
||||
)
|
||||
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"age": [17, 21, 25, 30],
|
||||
"is_student": [True, True, False, False],
|
||||
"name": ["Alice", "Bob", "Charlie", "Diana"],
|
||||
"is_adult": [False, True, True, True],
|
||||
"adult_student": [False, True, False, False],
|
||||
"minor_or_student": [True, True, False, False],
|
||||
}
|
||||
)
|
||||
pd.testing.assert_frame_equal(result, expected, check_dtype=False)
|
||||
|
||||
def test_complex_boolean_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test complex boolean expressions with Ray Dataset."""
|
||||
ds = ray.data.from_items(
|
||||
[
|
||||
{"score": 85, "passed": True, "bonus": False},
|
||||
{"score": 70, "passed": True, "bonus": True},
|
||||
{"score": 45, "passed": False, "bonus": False},
|
||||
]
|
||||
)
|
||||
|
||||
# Complex: (score > 80) OR (passed AND bonus)
|
||||
result = ds.with_column(
|
||||
"eligible", (col("score") > 80) | (col("passed") & col("bonus"))
|
||||
).to_pandas()
|
||||
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"score": [85, 70, 45],
|
||||
"passed": [True, True, False],
|
||||
"bonus": [False, True, False],
|
||||
"eligible": [True, True, False],
|
||||
}
|
||||
)
|
||||
pd.testing.assert_frame_equal(result, expected, check_dtype=False)
|
||||
|
||||
def test_logical_not_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test logical NOT operation with Ray Dataset."""
|
||||
ds = ray.data.range(5)
|
||||
result = ds.with_column("result", ~(col("id") == 2)).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{"id": [0, 1, 2, 3, 4], "result": [True, True, False, True, True]}
|
||||
)
|
||||
pd.testing.assert_frame_equal(result, expected, check_dtype=False)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"expression_factory,expected_results,test_id",
|
||||
[
|
||||
pytest.param(
|
||||
lambda: (col("age") > 18) & (col("country") == "USA"),
|
||||
[True, False, False],
|
||||
"complex_and",
|
||||
),
|
||||
pytest.param(
|
||||
lambda: (col("age") < 18) | (col("country") == "USA"),
|
||||
[True, True, False],
|
||||
"complex_or",
|
||||
),
|
||||
pytest.param(
|
||||
lambda: ~((col("age") < 25) & (col("country") != "USA")),
|
||||
[True, False, True],
|
||||
"complex_not",
|
||||
),
|
||||
pytest.param(
|
||||
lambda: (col("age") >= 21)
|
||||
& (col("score") >= 10)
|
||||
& col("active").is_not_null()
|
||||
& (col("active") == lit(True)),
|
||||
[True, False, False],
|
||||
"eligibility_flag",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_complex_boolean_expressions_with_dataset(
|
||||
self, ray_start_regular_shared, expression_factory, expected_results, test_id
|
||||
):
|
||||
"""Test complex boolean expressions with Ray Dataset."""
|
||||
test_data = [
|
||||
{"age": 25, "country": "USA", "active": True, "score": 20},
|
||||
{"age": 17, "country": "Canada", "active": False, "score": 10},
|
||||
{"age": 30, "country": "UK", "active": None, "score": 20},
|
||||
]
|
||||
|
||||
ds = ray.data.from_items(test_data)
|
||||
expression = expression_factory()
|
||||
result = ds.with_column("result", expression).to_pandas()
|
||||
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"age": [25, 17, 30],
|
||||
"country": ["USA", "Canada", "UK"],
|
||||
"active": [True, False, None],
|
||||
"score": [20, 10, 20],
|
||||
"result": expected_results,
|
||||
}
|
||||
)
|
||||
pd.testing.assert_frame_equal(result, expected, check_dtype=False)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,192 @@
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from packaging.version import parse as parse_version
|
||||
|
||||
import ray
|
||||
from ray.data._internal.util import rows_same
|
||||
from ray.data._internal.utils.arrow_utils import get_pyarrow_version
|
||||
from ray.data.datatype import DataType
|
||||
from ray.data.exceptions import UserCodeException
|
||||
from ray.data.expressions import col
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
get_pyarrow_version() < parse_version("20.0.0"),
|
||||
reason="with_column requires PyArrow >= 20.0.0",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"expr, target_type, expected_rows",
|
||||
[
|
||||
# Basic type conversions using Ray Data's DataType
|
||||
(col("id"), DataType.int64(), [{"id": i, "result": i} for i in range(5)]),
|
||||
(
|
||||
col("id"),
|
||||
DataType.float64(),
|
||||
[{"id": i, "result": float(i)} for i in range(5)],
|
||||
),
|
||||
(
|
||||
col("id"),
|
||||
DataType.string(),
|
||||
[{"id": i, "result": str(i)} for i in range(5)],
|
||||
),
|
||||
(
|
||||
col("id") / 2,
|
||||
DataType.int64(),
|
||||
[{"id": i, "result": i // 2} for i in range(5)],
|
||||
),
|
||||
# col("id")/2 uses integer division in expression layer, then cast to float64
|
||||
(
|
||||
col("id") / 2,
|
||||
DataType.float64(),
|
||||
[{"id": i, "result": float(i // 2)} for i in range(5)],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_cast_expression_basic(
|
||||
ray_start_regular_shared,
|
||||
expr,
|
||||
target_type,
|
||||
expected_rows,
|
||||
target_max_block_size_infinite_or_default,
|
||||
):
|
||||
"""Test basic type casting with cast() method."""
|
||||
ds = ray.data.range(5).with_column("result", expr.cast(target_type))
|
||||
actual = ds.take_all()
|
||||
assert rows_same(pd.DataFrame(actual), pd.DataFrame(expected_rows))
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
get_pyarrow_version() < parse_version("20.0.0"),
|
||||
reason="with_column requires PyArrow >= 20.0.0",
|
||||
)
|
||||
def test_cast_expression_usecase(
|
||||
ray_start_regular_shared, target_max_block_size_infinite_or_default
|
||||
):
|
||||
"""Test the user use case: converting float result from modulo to int64."""
|
||||
ds = ray.data.range(10)
|
||||
# The modulo operation returns float, cast it to int64
|
||||
ds = ds.with_column("part", (col("id") % 2).cast(DataType.int64()))
|
||||
actual = ds.take_all()
|
||||
expected_rows = [{"id": i, "part": i % 2} for i in range(10)]
|
||||
assert rows_same(pd.DataFrame(actual), pd.DataFrame(expected_rows))
|
||||
|
||||
# Verify the schema shows int64 type
|
||||
schema = ds.schema()
|
||||
assert "part" in schema.names
|
||||
part_type = schema.types[schema.names.index("part")]
|
||||
assert part_type == pa.int64()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
get_pyarrow_version() < parse_version("20.0.0"),
|
||||
reason="with_column requires PyArrow >= 20.0.0",
|
||||
)
|
||||
def test_cast_expression_chained(
|
||||
ray_start_regular_shared, target_max_block_size_infinite_or_default
|
||||
):
|
||||
"""Test that cast() can be chained with other expressions."""
|
||||
ds = ray.data.range(5)
|
||||
# Cast to float64 then multiply
|
||||
ds = ds.with_column("result", col("id").cast(DataType.float64()) * 2.5)
|
||||
actual = ds.take_all()
|
||||
expected_rows = [{"id": i, "result": i * 2.5} for i in range(5)]
|
||||
assert rows_same(pd.DataFrame(actual), pd.DataFrame(expected_rows))
|
||||
|
||||
# Cast result of arithmetic operation
|
||||
ds = ray.data.range(5)
|
||||
ds = ds.with_column("result", (col("id") + 1).cast(DataType.string()))
|
||||
actual = ds.take_all()
|
||||
expected_rows = [{"id": i, "result": str(i + 1)} for i in range(5)]
|
||||
assert rows_same(pd.DataFrame(actual), pd.DataFrame(expected_rows))
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
get_pyarrow_version() < parse_version("20.0.0"),
|
||||
reason="with_column requires PyArrow >= 20.0.0",
|
||||
)
|
||||
def test_cast_expression_safe_mode(
|
||||
ray_start_regular_shared, target_max_block_size_infinite_or_default
|
||||
):
|
||||
"""Test that safe=True (default) raises errors on invalid conversions."""
|
||||
ds = ray.data.from_items([{"value": "not_a_number"}])
|
||||
|
||||
# Attempting to cast non-numeric string to int should raise an error
|
||||
with pytest.raises((UserCodeException, ValueError, pa.ArrowInvalid)):
|
||||
ds.with_column("result", col("value").cast(DataType.int64())).materialize()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
get_pyarrow_version() < parse_version("20.0.0"),
|
||||
reason="with_column requires PyArrow >= 20.0.0",
|
||||
)
|
||||
def test_cast_expression_invalid_type(
|
||||
ray_start_regular_shared, target_max_block_size_infinite_or_default
|
||||
):
|
||||
"""Test that invalid type targets raise appropriate errors."""
|
||||
ds = ray.data.range(5)
|
||||
|
||||
# Passing a non-DataType target should raise TypeError
|
||||
with pytest.raises(
|
||||
TypeError, match="target_type must be a ray.data.datatype.DataType"
|
||||
):
|
||||
ds.with_column("result", col("id").cast("invalid_type")).materialize()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
get_pyarrow_version() < parse_version("20.0.0"),
|
||||
reason="with_column requires PyArrow >= 20.0.0",
|
||||
)
|
||||
def test_cast_expression_multiple_types(
|
||||
ray_start_regular_shared, target_max_block_size_infinite_or_default
|
||||
):
|
||||
"""Test casting with multiple different target types."""
|
||||
ds = ray.data.from_items([{"id": 42, "score": 3.14}])
|
||||
|
||||
# Cast id to different types
|
||||
ds = ds.with_column("id_int", col("id").cast(DataType.int64()))
|
||||
ds = ds.with_column("id_float", col("id").cast(DataType.float64()))
|
||||
ds = ds.with_column("id_str", col("id").cast(DataType.string()))
|
||||
|
||||
# Cast score to int (use safe=False to allow float truncation to int)
|
||||
ds = ds.with_column("score_int", col("score").cast(DataType.int64(), safe=False))
|
||||
|
||||
# Use rows_same to compare the full row content (expects DataFrames).
|
||||
results = ds.take_all()
|
||||
expected = [
|
||||
{
|
||||
"id": 42,
|
||||
"score": 3.14,
|
||||
"id_int": 42,
|
||||
"id_float": 42.0,
|
||||
"id_str": "42",
|
||||
"score_int": 3,
|
||||
}
|
||||
]
|
||||
assert rows_same(pd.DataFrame(results), pd.DataFrame(expected))
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
get_pyarrow_version() < parse_version("20.0.0"),
|
||||
reason="with_column requires PyArrow >= 20.0.0",
|
||||
)
|
||||
def test_cast_expression_python_type_datatype_error(
|
||||
ray_start_regular_shared, target_max_block_size_infinite_or_default
|
||||
):
|
||||
"""Test that using Python-type-backed DataType in cast() raises a clear error."""
|
||||
# Error is raised at expression build time when cast() is called (not at materialize).
|
||||
error_match = "Python-type-backed DataType.*requires.*values"
|
||||
with pytest.raises(TypeError, match=error_match):
|
||||
col("id").cast(DataType(int))
|
||||
with pytest.raises(TypeError, match=error_match):
|
||||
col("id").cast(DataType(str))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,78 @@
|
||||
"""Integration tests for comparison expression operations.
|
||||
|
||||
These tests require Ray and test end-to-end comparison expression evaluation.
|
||||
"""
|
||||
|
||||
import pandas as pd
|
||||
import pytest
|
||||
from packaging.version import parse as parse_version
|
||||
|
||||
import ray
|
||||
from ray.data._internal.utils.arrow_utils import get_pyarrow_version
|
||||
from ray.data.expressions import col
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
get_pyarrow_version() < parse_version("20.0.0"),
|
||||
reason="Expression integration tests require PyArrow >= 20.0.0",
|
||||
)
|
||||
|
||||
|
||||
class TestComparisonIntegration:
|
||||
"""Integration tests for comparison expressions with Ray Dataset."""
|
||||
|
||||
def test_comparison_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test comparison expressions work correctly with Ray Dataset."""
|
||||
ds = ray.data.from_items(
|
||||
[
|
||||
{"age": 17, "name": "Alice"},
|
||||
{"age": 21, "name": "Bob"},
|
||||
{"age": 25, "name": "Charlie"},
|
||||
{"age": 18, "name": "Diana"},
|
||||
]
|
||||
)
|
||||
|
||||
result = ds.with_column("is_adult", col("age") >= 18).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"age": [17, 21, 25, 18],
|
||||
"name": ["Alice", "Bob", "Charlie", "Diana"],
|
||||
"is_adult": [False, True, True, True],
|
||||
}
|
||||
)
|
||||
pd.testing.assert_frame_equal(result, expected, check_dtype=False)
|
||||
|
||||
def test_multiple_comparisons_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test multiple comparison expressions with Ray Dataset."""
|
||||
ds = ray.data.from_items(
|
||||
[
|
||||
{"score": 45, "passing": 50},
|
||||
{"score": 75, "passing": 50},
|
||||
{"score": 50, "passing": 50},
|
||||
]
|
||||
)
|
||||
|
||||
result = (
|
||||
ds.with_column("passed", col("score") >= col("passing"))
|
||||
.with_column("failed", col("score") < col("passing"))
|
||||
.with_column("borderline", col("score") == col("passing"))
|
||||
.to_pandas()
|
||||
)
|
||||
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"score": [45, 75, 50],
|
||||
"passing": [50, 50, 50],
|
||||
"passed": [False, True, True],
|
||||
"failed": [True, False, False],
|
||||
"borderline": [False, False, True],
|
||||
}
|
||||
)
|
||||
pd.testing.assert_frame_equal(result, expected, check_dtype=False)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,63 @@
|
||||
"""Integration tests for array namespace expressions.
|
||||
|
||||
These tests require Ray and test end-to-end array namespace expression evaluation.
|
||||
"""
|
||||
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from packaging import version
|
||||
|
||||
import ray
|
||||
from ray.data._internal.util import rows_same
|
||||
from ray.data.expressions import col
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
version.parse(pa.__version__) < version.parse("19.0.0"),
|
||||
reason="Namespace expressions tests require PyArrow >= 19.0",
|
||||
)
|
||||
|
||||
|
||||
def _make_fixed_size_list_table() -> pa.Table:
|
||||
values = pa.array([1, 2, 3, 4, 5, 6], type=pa.int64())
|
||||
fixed = pa.FixedSizeListArray.from_arrays(values, list_size=2)
|
||||
return pa.Table.from_arrays([fixed], names=["features"])
|
||||
|
||||
|
||||
def test_arr_to_list_fixed_size(ray_start_regular_shared):
|
||||
table = _make_fixed_size_list_table()
|
||||
ds = ray.data.from_arrow(table)
|
||||
|
||||
result = (
|
||||
ds.with_column("features", col("features").arr.to_list())
|
||||
.select_columns(["features"])
|
||||
.to_pandas()
|
||||
)
|
||||
expected = pd.DataFrame(
|
||||
[
|
||||
{"features": [1, 2]},
|
||||
{"features": [3, 4]},
|
||||
{"features": [5, 6]},
|
||||
]
|
||||
)
|
||||
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
def test_arr_to_list_invalid_dtype_raises(ray_start_regular_shared):
|
||||
ds = ray.data.from_items([{"value": 1}, {"value": 2}])
|
||||
|
||||
with pytest.raises(
|
||||
(ray.exceptions.RayTaskError, ray.exceptions.UserCodeException)
|
||||
) as exc_info:
|
||||
ds.with_column("value_list", col("value").arr.to_list()).to_pandas()
|
||||
|
||||
assert "to_list() can only be called on list-like columns" in str(exc_info.value)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,81 @@
|
||||
"""Integration tests for datetime namespace expressions.
|
||||
|
||||
These tests require Ray and test end-to-end datetime namespace expression evaluation.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from packaging import version
|
||||
|
||||
import ray
|
||||
from ray.data._internal.util import rows_same
|
||||
from ray.data.expressions import col
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
version.parse(pa.__version__) < version.parse("19.0.0"),
|
||||
reason="Namespace expressions tests require PyArrow >= 19.0",
|
||||
)
|
||||
|
||||
|
||||
class TestDatetimeNamespace:
|
||||
"""Tests for datetime namespace operations."""
|
||||
|
||||
def test_datetime_namespace_all_operations(self, ray_start_regular_shared):
|
||||
"""Test all datetime namespace operations on a datetime column."""
|
||||
ts = datetime.datetime(2024, 1, 2, 10, 30, 0)
|
||||
|
||||
ds = ray.data.from_items([{"ts": ts}])
|
||||
|
||||
result_ds = (
|
||||
ds.with_column("year", col("ts").dt.year())
|
||||
.with_column("month", col("ts").dt.month())
|
||||
.with_column("day", col("ts").dt.day())
|
||||
.with_column("hour", col("ts").dt.hour())
|
||||
.with_column("minute", col("ts").dt.minute())
|
||||
.with_column("second", col("ts").dt.second())
|
||||
.with_column("date_str", col("ts").dt.strftime("%Y-%m-%d"))
|
||||
.with_column("ts_floor", col("ts").dt.floor("day"))
|
||||
.with_column("ts_ceil", col("ts").dt.ceil("day"))
|
||||
.with_column("ts_round", col("ts").dt.round("day"))
|
||||
.drop_columns(["ts"])
|
||||
)
|
||||
|
||||
actual = result_ds.to_pandas()
|
||||
|
||||
expected = pd.DataFrame(
|
||||
[
|
||||
{
|
||||
"year": 2024,
|
||||
"month": 1,
|
||||
"day": 2,
|
||||
"hour": 10,
|
||||
"minute": 30,
|
||||
"second": 0,
|
||||
"date_str": "2024-01-02",
|
||||
"ts_floor": pd.Timestamp("2024-01-02"),
|
||||
"ts_ceil": pd.Timestamp("2024-01-03"),
|
||||
# round("day") rounds to nearest day; 10:30 < 12:00 so rounds down
|
||||
"ts_round": pd.Timestamp("2024-01-02"),
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
assert rows_same(actual, expected)
|
||||
|
||||
def test_dt_namespace_invalid_dtype_raises(self, ray_start_regular_shared):
|
||||
"""Test that dt namespace on non-datetime column raises an error."""
|
||||
ds = ray.data.from_items([{"value": 1}])
|
||||
|
||||
with pytest.raises(Exception):
|
||||
ds.with_column("year", col("value").dt.year()).to_pandas()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,183 @@
|
||||
"""Integration tests for list namespace expressions.
|
||||
|
||||
These tests require Ray and test end-to-end list namespace expression evaluation.
|
||||
"""
|
||||
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from packaging import version
|
||||
|
||||
import ray
|
||||
from ray.data._internal.util import rows_same
|
||||
from ray.data.expressions import col
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.exceptions import RayTaskError
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
version.parse(pa.__version__) < version.parse("19.0.0"),
|
||||
reason="Namespace expressions tests require PyArrow >= 19.0",
|
||||
)
|
||||
|
||||
|
||||
def _create_dataset(items_data, dataset_format, arrow_table=None):
|
||||
if dataset_format == "arrow":
|
||||
if arrow_table is not None:
|
||||
ds = ray.data.from_arrow(arrow_table)
|
||||
else:
|
||||
table = pa.Table.from_pylist(items_data)
|
||||
ds = ray.data.from_arrow(table)
|
||||
elif dataset_format == "pandas":
|
||||
if arrow_table is not None:
|
||||
df = arrow_table.to_pandas()
|
||||
else:
|
||||
df = pd.DataFrame(items_data)
|
||||
ds = ray.data.from_blocks([df])
|
||||
return ds
|
||||
|
||||
|
||||
DATASET_FORMATS = ["pandas", "arrow"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dataset_format", DATASET_FORMATS)
|
||||
class TestListNamespace:
|
||||
"""Tests for list namespace operations."""
|
||||
|
||||
def test_list_len(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test list.len() returns length of each list."""
|
||||
data = [
|
||||
{"items": [1, 2, 3]},
|
||||
{"items": [4, 5]},
|
||||
{"items": []},
|
||||
]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column("len", col("items").list.len()).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"items": [[1, 2, 3], [4, 5], []],
|
||||
"len": [3, 2, 0],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_list_get(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test list.get() extracts element at index."""
|
||||
data = [
|
||||
{"items": [10, 20, 30]},
|
||||
{"items": [40, 50, 60]},
|
||||
]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column("first", col("items").list.get(0)).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"items": [[10, 20, 30], [40, 50, 60]],
|
||||
"first": [10, 40],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_list_bracket_index(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test list[i] bracket notation for element access."""
|
||||
data = [{"items": [10, 20, 30]}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column("elem", col("items").list[1]).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"items": [[10, 20, 30]],
|
||||
"elem": [20],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_list_with_arithmetic(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test list operations combined with arithmetic."""
|
||||
data = [{"items": [1, 2, 3]}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column("len_plus_one", col("items").list.len() + 1).to_pandas()
|
||||
expected = pd.DataFrame({"items": [[1, 2, 3]], "len_plus_one": [4]})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_list_sort(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test list.sort() sorts each list with custom options."""
|
||||
data = [
|
||||
{"items": [3, 1, 2]},
|
||||
{"items": [None, 4, 2]},
|
||||
]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
method = col("items").list.sort(order="descending", null_placement="at_start")
|
||||
result = ds.with_column("sorted", method).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"items": [[3, 1, 2], [None, 4, 2]],
|
||||
"sorted": [[3, 2, 1], [None, 4, 2]],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_list_flatten(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test list.flatten() removes one nesting level."""
|
||||
data = [
|
||||
{"items": [[1, 2], [3]]},
|
||||
{"items": [[], [4, 5]]},
|
||||
]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column("flattened", col("items").list.flatten()).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"items": [[[1, 2], [3]], [[], [4, 5]]],
|
||||
"flattened": [[1, 2, 3], [4, 5]],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_list_flatten_requires_nested_lists(
|
||||
self, ray_start_regular_shared, dataset_format
|
||||
):
|
||||
"""list.flatten() should raise if elements aren't lists."""
|
||||
data = [{"items": [1, 2]}, {"items": [3, 4]}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
with pytest.raises(RayTaskError):
|
||||
ds.with_column("flattened", col("items").list.flatten()).materialize()
|
||||
|
||||
def test_list_flatten_large_list_type(
|
||||
self, ray_start_regular_shared, dataset_format
|
||||
):
|
||||
"""Flatten should preserve LargeList type when present."""
|
||||
if dataset_format != "arrow":
|
||||
pytest.skip("LargeList type only available via Arrow tables.")
|
||||
|
||||
arrow_type = pa.large_list(pa.list_(pa.int64()))
|
||||
table = pa.Table.from_arrays(
|
||||
[
|
||||
pa.array(
|
||||
[
|
||||
[[1, 2], [3]],
|
||||
[[], [4, 5]],
|
||||
],
|
||||
type=arrow_type,
|
||||
)
|
||||
],
|
||||
names=["items"],
|
||||
)
|
||||
ds = _create_dataset(None, dataset_format, arrow_table=table)
|
||||
result = ds.with_column("flattened", col("items").list.flatten())
|
||||
arrow_refs = result.to_arrow_refs()
|
||||
tables = ray.get(arrow_refs)
|
||||
result_table = pa.concat_tables(tables) if len(tables) > 1 else tables[0]
|
||||
|
||||
flattened_type = result_table.schema.field("flattened").type
|
||||
assert flattened_type == pa.large_list(pa.int64())
|
||||
expected = pa.Table.from_arrays(
|
||||
[
|
||||
pa.array([[1, 2, 3], [4, 5]], type=pa.large_list(pa.int64())),
|
||||
],
|
||||
names=["flattened"],
|
||||
)
|
||||
assert result_table.select(["flattened"]).combine_chunks().equals(expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,150 @@
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from packaging import version
|
||||
|
||||
import ray
|
||||
from ray.data._internal.util import rows_same
|
||||
from ray.data.expressions import col
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
version.parse(pa.__version__) < version.parse("19.0.0"),
|
||||
reason="Namespace expressions tests require PyArrow >= 19.0",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def map_dataset():
|
||||
"""Fixture that creates a dataset backed by an Arrow MapArray column."""
|
||||
map_items = [
|
||||
{"attrs": {"color": "red", "size": "M"}},
|
||||
{"attrs": {"brand": "Ray"}},
|
||||
]
|
||||
map_type = pa.map_(pa.string(), pa.string())
|
||||
arrow_table = pa.table(
|
||||
{"attrs": pa.array([row["attrs"] for row in map_items], type=map_type)}
|
||||
)
|
||||
return ray.data.from_arrow(arrow_table)
|
||||
|
||||
|
||||
def _assert_result(result_df: pd.DataFrame, expected_df: pd.DataFrame, drop_cols: list):
|
||||
"""Helper to drop columns and assert equality."""
|
||||
result_df = result_df.drop(columns=drop_cols)
|
||||
assert rows_same(result_df, expected_df)
|
||||
|
||||
|
||||
class TestMapNamespace:
|
||||
"""Tests for map namespace operations using the shared map_dataset fixture."""
|
||||
|
||||
def test_map_keys(self, map_dataset):
|
||||
result = map_dataset.with_column("keys", col("attrs").map.keys()).to_pandas()
|
||||
expected = pd.DataFrame({"keys": [["color", "size"], ["brand"]]})
|
||||
_assert_result(result, expected, drop_cols=["attrs"])
|
||||
|
||||
def test_map_values(self, map_dataset):
|
||||
result = map_dataset.with_column(
|
||||
"values", col("attrs").map.values()
|
||||
).to_pandas()
|
||||
expected = pd.DataFrame({"values": [["red", "M"], ["Ray"]]})
|
||||
_assert_result(result, expected, drop_cols=["attrs"])
|
||||
|
||||
def test_map_chaining(self, map_dataset):
|
||||
# map.keys() returns a list, so .list.len() should apply
|
||||
result = map_dataset.with_column(
|
||||
"num_keys", col("attrs").map.keys().list.len()
|
||||
).to_pandas()
|
||||
expected = pd.DataFrame({"num_keys": [2, 1]})
|
||||
_assert_result(result, expected, drop_cols=["attrs"])
|
||||
|
||||
|
||||
def test_physical_map_extraction():
|
||||
"""Test extraction works on List<Struct> (Physical Maps)."""
|
||||
# Construct List<Struct<k, v>>
|
||||
struct_type = pa.struct([pa.field("k", pa.string()), pa.field("v", pa.int64())])
|
||||
list_type = pa.list_(struct_type)
|
||||
|
||||
data_py = [[{"k": "a", "v": 1}], [{"k": "b", "v": 2}]]
|
||||
arrow_table = pa.Table.from_arrays(
|
||||
[pa.array(data_py, type=list_type)], names=["data"]
|
||||
)
|
||||
ds = ray.data.from_arrow(arrow_table)
|
||||
|
||||
result = (
|
||||
ds.with_column("keys", col("data").map.keys())
|
||||
.with_column("values", col("data").map.values())
|
||||
.to_pandas()
|
||||
)
|
||||
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"data": data_py,
|
||||
"keys": [["a"], ["b"]],
|
||||
"values": [[1], [2]],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
def test_map_sliced_offsets():
|
||||
"""Test extraction works correctly on sliced Arrow arrays (offset > 0)."""
|
||||
items = [{"m": {"id": i}} for i in range(10)]
|
||||
map_type = pa.map_(pa.string(), pa.int64())
|
||||
arrays = pa.array([row["m"] for row in items], type=map_type)
|
||||
table = pa.Table.from_arrays([arrays], names=["m"])
|
||||
|
||||
# Force offsets by slicing the table before ingestion
|
||||
sliced_table = table.slice(offset=7, length=3)
|
||||
ds = ray.data.from_arrow(sliced_table)
|
||||
|
||||
result = ds.with_column("vals", col("m").map.values()).to_pandas()
|
||||
expected = pd.DataFrame({"vals": [[7], [8], [9]]})
|
||||
_assert_result(result, expected, drop_cols=["m"])
|
||||
|
||||
|
||||
def test_map_nulls_and_empty():
|
||||
"""Test handling of null maps and empty maps."""
|
||||
items_data = [{"m": {"a": 1}}, {"m": {}}, {"m": None}]
|
||||
|
||||
map_type = pa.map_(pa.string(), pa.int64())
|
||||
arrays = pa.array([row["m"] for row in items_data], type=map_type)
|
||||
arrow_table = pa.Table.from_arrays([arrays], names=["m"])
|
||||
ds = ray.data.from_arrow(arrow_table)
|
||||
|
||||
rows = (
|
||||
ds.with_column("keys", col("m").map.keys())
|
||||
.with_column("values", col("m").map.values())
|
||||
.take_all()
|
||||
)
|
||||
|
||||
assert list(rows[0]["keys"]) == ["a"] and list(rows[0]["values"]) == [1]
|
||||
assert len(rows[1]["keys"]) == 0 and len(rows[1]["values"]) == 0
|
||||
assert rows[2]["keys"] is None and rows[2]["values"] is None
|
||||
|
||||
|
||||
def test_empty_chunked_array():
|
||||
"""Test extraction works on empty ChunkedArray (zero chunks)."""
|
||||
from ray.data.namespace_expressions.map_namespace import (
|
||||
MapComponent,
|
||||
_extract_map_component,
|
||||
)
|
||||
|
||||
# Create empty ChunkedArray with map type
|
||||
map_type = pa.map_(pa.string(), pa.int64())
|
||||
empty_chunked = pa.chunked_array([], type=map_type)
|
||||
assert empty_chunked.num_chunks == 0
|
||||
|
||||
# Extract keys - should return empty ChunkedArray with list<string> type
|
||||
keys_result = _extract_map_component(empty_chunked, MapComponent.KEYS)
|
||||
assert len(keys_result) == 0
|
||||
assert keys_result.type == pa.list_(pa.string())
|
||||
|
||||
# Extract values - should return empty ChunkedArray with list<int64> type
|
||||
values_result = _extract_map_component(empty_chunked, MapComponent.VALUES)
|
||||
assert len(values_result) == 0
|
||||
assert values_result.type == pa.list_(pa.int64())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,328 @@
|
||||
"""Integration tests for string namespace expressions.
|
||||
|
||||
These tests require Ray and test end-to-end string namespace expression evaluation.
|
||||
"""
|
||||
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from packaging import version
|
||||
|
||||
import ray
|
||||
from ray.data._internal.util import rows_same
|
||||
from ray.data.expressions import col
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
version.parse(pa.__version__) < version.parse("19.0.0"),
|
||||
reason="Namespace expressions tests require PyArrow >= 19.0",
|
||||
)
|
||||
|
||||
|
||||
def _create_dataset(items_data, dataset_format, arrow_table=None):
|
||||
if dataset_format == "arrow":
|
||||
if arrow_table is not None:
|
||||
ds = ray.data.from_arrow(arrow_table)
|
||||
else:
|
||||
table = pa.Table.from_pylist(items_data)
|
||||
ds = ray.data.from_arrow(table)
|
||||
elif dataset_format == "pandas":
|
||||
if arrow_table is not None:
|
||||
df = arrow_table.to_pandas()
|
||||
else:
|
||||
df = pd.DataFrame(items_data)
|
||||
ds = ray.data.from_blocks([df])
|
||||
return ds
|
||||
|
||||
|
||||
DATASET_FORMATS = ["pandas", "arrow"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dataset_format", DATASET_FORMATS)
|
||||
@pytest.mark.parametrize(
|
||||
"method_name,input_values,expected_results",
|
||||
[
|
||||
("len", ["Alice", "Bob"], [5, 3]),
|
||||
("byte_len", ["ABC"], [3]),
|
||||
],
|
||||
)
|
||||
class TestStringLength:
|
||||
"""Tests for string length operations."""
|
||||
|
||||
def test_string_length(
|
||||
self,
|
||||
ray_start_regular_shared,
|
||||
dataset_format,
|
||||
method_name,
|
||||
input_values,
|
||||
expected_results,
|
||||
):
|
||||
"""Test string length methods."""
|
||||
data = [{"name": v} for v in input_values]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
|
||||
method = getattr(col("name").str, method_name)
|
||||
result = ds.with_column("result", method()).to_pandas()
|
||||
|
||||
expected = pd.DataFrame({"name": input_values, "result": expected_results})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dataset_format", DATASET_FORMATS)
|
||||
@pytest.mark.parametrize(
|
||||
"method_name,input_values,expected_values",
|
||||
[
|
||||
("upper", ["alice", "bob"], ["ALICE", "BOB"]),
|
||||
("lower", ["ALICE", "BOB"], ["alice", "bob"]),
|
||||
("capitalize", ["alice", "bob"], ["Alice", "Bob"]),
|
||||
("title", ["alice smith", "bob jones"], ["Alice Smith", "Bob Jones"]),
|
||||
("swapcase", ["AlIcE"], ["aLiCe"]),
|
||||
],
|
||||
)
|
||||
class TestStringCase:
|
||||
"""Tests for string case conversion."""
|
||||
|
||||
def test_string_case(
|
||||
self,
|
||||
ray_start_regular_shared,
|
||||
dataset_format,
|
||||
method_name,
|
||||
input_values,
|
||||
expected_values,
|
||||
):
|
||||
"""Test string case conversion methods."""
|
||||
data = [{"name": v} for v in input_values]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
|
||||
method = getattr(col("name").str, method_name)
|
||||
result = ds.with_column("result", method()).to_pandas()
|
||||
|
||||
expected = pd.DataFrame({"name": input_values, "result": expected_values})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dataset_format", DATASET_FORMATS)
|
||||
@pytest.mark.parametrize(
|
||||
"method_name,input_values,expected_results",
|
||||
[
|
||||
("is_alpha", ["abc", "abc123", "123"], [True, False, False]),
|
||||
("is_alnum", ["abc123", "abc-123"], [True, False]),
|
||||
("is_digit", ["123", "12a"], [True, False]),
|
||||
("is_space", [" ", " a "], [True, False]),
|
||||
("is_lower", ["abc", "Abc"], [True, False]),
|
||||
("is_upper", ["ABC", "Abc"], [True, False]),
|
||||
("is_ascii", ["hello", "hello😊"], [True, False]),
|
||||
],
|
||||
)
|
||||
class TestStringPredicates:
|
||||
"""Tests for string predicate methods (is_*)."""
|
||||
|
||||
def test_string_predicate(
|
||||
self,
|
||||
ray_start_regular_shared,
|
||||
dataset_format,
|
||||
method_name,
|
||||
input_values,
|
||||
expected_results,
|
||||
):
|
||||
"""Test string predicate methods."""
|
||||
data = [{"val": v} for v in input_values]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
|
||||
method = getattr(col("val").str, method_name)
|
||||
result = ds.with_column("result", method()).to_pandas()
|
||||
|
||||
expected = pd.DataFrame({"val": input_values, "result": expected_results})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dataset_format", DATASET_FORMATS)
|
||||
@pytest.mark.parametrize(
|
||||
"method_name,method_args,input_values,expected_values",
|
||||
[
|
||||
("strip", (), [" hello ", " world "], ["hello", "world"]),
|
||||
("strip", ("x",), ["xxxhelloxxx"], ["hello"]),
|
||||
("lstrip", (), [" hello "], ["hello "]),
|
||||
("rstrip", (), [" hello "], [" hello"]),
|
||||
],
|
||||
)
|
||||
class TestStringTrimming:
|
||||
"""Tests for string trimming operations."""
|
||||
|
||||
def test_string_trimming(
|
||||
self,
|
||||
ray_start_regular_shared,
|
||||
dataset_format,
|
||||
method_name,
|
||||
method_args,
|
||||
input_values,
|
||||
expected_values,
|
||||
):
|
||||
"""Test string trimming methods."""
|
||||
data = [{"val": v} for v in input_values]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
|
||||
method = getattr(col("val").str, method_name)
|
||||
result = ds.with_column("result", method(*method_args)).to_pandas()
|
||||
|
||||
expected = pd.DataFrame({"val": input_values, "result": expected_values})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dataset_format", DATASET_FORMATS)
|
||||
@pytest.mark.parametrize(
|
||||
"method_name,method_kwargs,expected_value",
|
||||
[
|
||||
("pad", {"width": 5, "fillchar": "*", "side": "right"}, "hi***"),
|
||||
("pad", {"width": 5, "fillchar": "*", "side": "left"}, "***hi"),
|
||||
("pad", {"width": 6, "fillchar": "*", "side": "both"}, "**hi**"),
|
||||
("lpad", {"width": 5, "padding": "*"}, "***hi"),
|
||||
("rpad", {"width": 5, "padding": "*"}, "hi***"),
|
||||
("center", {"width": 6, "padding": "*"}, "**hi**"),
|
||||
],
|
||||
)
|
||||
class TestStringPadding:
|
||||
"""Tests for string padding operations."""
|
||||
|
||||
def test_string_padding(
|
||||
self,
|
||||
ray_start_regular_shared,
|
||||
dataset_format,
|
||||
method_name,
|
||||
method_kwargs,
|
||||
expected_value,
|
||||
):
|
||||
"""Test string padding methods."""
|
||||
data = [{"val": "hi"}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
|
||||
method = getattr(col("val").str, method_name)
|
||||
result = ds.with_column("result", method(**method_kwargs)).to_pandas()
|
||||
|
||||
expected = pd.DataFrame({"val": ["hi"], "result": [expected_value]})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dataset_format", DATASET_FORMATS)
|
||||
@pytest.mark.parametrize(
|
||||
"method_name,method_args,method_kwargs,input_values,expected_results",
|
||||
[
|
||||
("starts_with", ("A",), {}, ["Alice", "Bob", "Alex"], [True, False, True]),
|
||||
("starts_with", ("A",), {"ignore_case": True}, ["alice", "bob"], [True, False]),
|
||||
("ends_with", ("e",), {}, ["Alice", "Bob"], [True, False]),
|
||||
("contains", ("li",), {}, ["Alice", "Bob", "Charlie"], [True, False, True]),
|
||||
("find", ("i",), {}, ["Alice", "Bob"], [2, -1]),
|
||||
("count", ("a",), {}, ["banana", "apple"], [3, 1]),
|
||||
("match", ("Al%",), {}, ["Alice", "Bob", "Alex"], [True, False, True]),
|
||||
],
|
||||
)
|
||||
class TestStringSearch:
|
||||
"""Tests for string searching operations."""
|
||||
|
||||
def test_string_search(
|
||||
self,
|
||||
ray_start_regular_shared,
|
||||
dataset_format,
|
||||
method_name,
|
||||
method_args,
|
||||
method_kwargs,
|
||||
input_values,
|
||||
expected_results,
|
||||
):
|
||||
"""Test string searching methods."""
|
||||
data = [{"val": v} for v in input_values]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
|
||||
method = getattr(col("val").str, method_name)
|
||||
result = ds.with_column(
|
||||
"result", method(*method_args, **method_kwargs)
|
||||
).to_pandas()
|
||||
|
||||
expected = pd.DataFrame({"val": input_values, "result": expected_results})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dataset_format", DATASET_FORMATS)
|
||||
class TestStringTransform:
|
||||
"""Tests for string transformation operations."""
|
||||
|
||||
def test_reverse(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test str.reverse() reverses strings."""
|
||||
data = [{"val": "hello"}, {"val": "world"}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column("rev", col("val").str.reverse()).to_pandas()
|
||||
expected = pd.DataFrame({"val": ["hello", "world"], "rev": ["olleh", "dlrow"]})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_slice(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test str.slice() extracts substring."""
|
||||
data = [{"val": "hello"}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column("sliced", col("val").str.slice(1, 4)).to_pandas()
|
||||
expected = pd.DataFrame({"val": ["hello"], "sliced": ["ell"]})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_replace(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test str.replace() replaces substring."""
|
||||
data = [{"val": "hello world"}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column(
|
||||
"replaced", col("val").str.replace("world", "universe")
|
||||
).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{"val": ["hello world"], "replaced": ["hello universe"]}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_replace_with_max(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test str.replace() with max_replacements."""
|
||||
data = [{"val": "aaa"}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column(
|
||||
"replaced", col("val").str.replace("a", "X", max_replacements=2)
|
||||
).to_pandas()
|
||||
expected = pd.DataFrame({"val": ["aaa"], "replaced": ["XXa"]})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_repeat(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test str.repeat() repeats strings."""
|
||||
data = [{"val": "A"}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column("repeated", col("val").str.repeat(3)).to_pandas()
|
||||
expected = pd.DataFrame({"val": ["A"], "repeated": ["AAA"]})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_string_with_comparison(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test string operations combined with comparison."""
|
||||
data = [{"name": "Alice"}, {"name": "Bo"}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = ds.with_column("long_name", col("name").str.len() > 3).to_pandas()
|
||||
expected = pd.DataFrame({"name": ["Alice", "Bo"], "long_name": [True, False]})
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_multiple_string_operations(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test multiple namespace operations in single pipeline."""
|
||||
data = [{"name": "alice"}]
|
||||
ds = _create_dataset(data, dataset_format)
|
||||
result = (
|
||||
ds.with_column("upper", col("name").str.upper())
|
||||
.with_column("len", col("name").str.len())
|
||||
.with_column("starts_a", col("name").str.starts_with("a"))
|
||||
.to_pandas()
|
||||
)
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"name": ["alice"],
|
||||
"upper": ["ALICE"],
|
||||
"len": [5],
|
||||
"starts_a": [True],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,313 @@
|
||||
"""Integration tests for struct namespace expressions.
|
||||
|
||||
These tests require Ray and test end-to-end struct namespace expression evaluation.
|
||||
"""
|
||||
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from packaging import version
|
||||
|
||||
import ray
|
||||
from ray.data._internal.util import rows_same
|
||||
from ray.data.expressions import col
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
version.parse(pa.__version__) < version.parse("19.0.0"),
|
||||
reason="Namespace expressions tests require PyArrow >= 19.0",
|
||||
)
|
||||
|
||||
|
||||
def _create_dataset(items_data, dataset_format, arrow_table=None):
|
||||
if dataset_format == "arrow":
|
||||
if arrow_table is not None:
|
||||
ds = ray.data.from_arrow(arrow_table)
|
||||
else:
|
||||
table = pa.Table.from_pylist(items_data)
|
||||
ds = ray.data.from_arrow(table)
|
||||
elif dataset_format == "pandas":
|
||||
if arrow_table is not None:
|
||||
df = arrow_table.to_pandas()
|
||||
else:
|
||||
df = pd.DataFrame(items_data)
|
||||
ds = ray.data.from_blocks([df])
|
||||
return ds
|
||||
|
||||
|
||||
DATASET_FORMATS = ["pandas", "arrow"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dataset_format", DATASET_FORMATS)
|
||||
class TestStructNamespace:
|
||||
"""Tests for struct namespace operations."""
|
||||
|
||||
def test_struct_bracket_bool_index_raises(self, dataset_format):
|
||||
"""Test struct[bool] raises TypeError instead of being treated as int."""
|
||||
del dataset_format # Unused, required by class-level parametrization.
|
||||
|
||||
with pytest.raises(
|
||||
TypeError, match="Struct indices must be strings or integers"
|
||||
):
|
||||
col("user").struct[True]
|
||||
|
||||
@pytest.mark.parametrize("bad_index", ["1", True])
|
||||
def test_struct_field_by_index_non_integer_raises(self, dataset_format, bad_index):
|
||||
"""Test struct.field_by_index() rejects non-integer indices."""
|
||||
del dataset_format # Unused, required by class-level parametrization.
|
||||
|
||||
with pytest.raises(TypeError, match="Struct field index must be an integer"):
|
||||
col("user").struct.field_by_index(bad_index)
|
||||
|
||||
def test_struct_field_by_index_negative_raises(self, dataset_format):
|
||||
"""Test struct.field_by_index() rejects negative indices."""
|
||||
del dataset_format # Unused, required by class-level parametrization.
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="Struct field index must be non-negative, got -1"
|
||||
):
|
||||
col("user").struct.field_by_index(-1)
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="Struct field index must be non-negative, got -1"
|
||||
):
|
||||
col("user").struct[-1]
|
||||
|
||||
def test_struct_field(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test struct.field() extracts field."""
|
||||
arrow_table = pa.table(
|
||||
{
|
||||
"user": pa.array(
|
||||
[
|
||||
{"name": "Alice", "age": 30},
|
||||
{"name": "Bob", "age": 25},
|
||||
],
|
||||
type=pa.struct(
|
||||
[
|
||||
pa.field("name", pa.string()),
|
||||
pa.field("age", pa.int32()),
|
||||
]
|
||||
),
|
||||
)
|
||||
}
|
||||
)
|
||||
items_data = [
|
||||
{"user": {"name": "Alice", "age": 30}},
|
||||
{"user": {"name": "Bob", "age": 25}},
|
||||
]
|
||||
ds = _create_dataset(items_data, dataset_format, arrow_table)
|
||||
|
||||
result = ds.with_column("age", col("user").struct.field("age")).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"user": [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}],
|
||||
"age": [30, 25],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_struct_bracket(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test struct['field'] bracket notation."""
|
||||
arrow_table = pa.table(
|
||||
{
|
||||
"user": pa.array(
|
||||
[
|
||||
{"name": "Alice", "age": 30},
|
||||
{"name": "Bob", "age": 25},
|
||||
],
|
||||
type=pa.struct(
|
||||
[
|
||||
pa.field("name", pa.string()),
|
||||
pa.field("age", pa.int32()),
|
||||
]
|
||||
),
|
||||
)
|
||||
}
|
||||
)
|
||||
items_data = [
|
||||
{"user": {"name": "Alice", "age": 30}},
|
||||
{"user": {"name": "Bob", "age": 25}},
|
||||
]
|
||||
ds = _create_dataset(items_data, dataset_format, arrow_table)
|
||||
|
||||
result = ds.with_column("name", col("user").struct["name"]).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"user": [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}],
|
||||
"name": ["Alice", "Bob"],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_struct_field_by_index(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test struct.field_by_index() extracts field by position."""
|
||||
if dataset_format == "pandas":
|
||||
pytest.skip(
|
||||
"Index-based struct access requires stable Arrow struct field ordering."
|
||||
)
|
||||
arrow_table = pa.table(
|
||||
{
|
||||
"user": pa.array(
|
||||
[
|
||||
{"name": "Alice", "age": 30},
|
||||
{"name": "Bob", "age": 25},
|
||||
],
|
||||
type=pa.struct(
|
||||
[
|
||||
pa.field("name", pa.string()),
|
||||
pa.field("age", pa.int32()),
|
||||
]
|
||||
),
|
||||
)
|
||||
}
|
||||
)
|
||||
items_data = [
|
||||
{"user": {"name": "Alice", "age": 30}},
|
||||
{"user": {"name": "Bob", "age": 25}},
|
||||
]
|
||||
ds = _create_dataset(items_data, dataset_format, arrow_table)
|
||||
|
||||
result = ds.with_column("age", col("user").struct.field_by_index(1)).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"user": [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}],
|
||||
"age": [30, 25],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_struct_bracket_with_index(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test struct[index] bracket notation."""
|
||||
if dataset_format == "pandas":
|
||||
pytest.skip(
|
||||
"Index-based struct access requires stable Arrow struct field ordering."
|
||||
)
|
||||
arrow_table = pa.table(
|
||||
{
|
||||
"user": pa.array(
|
||||
[
|
||||
{"name": "Alice", "age": 30},
|
||||
{"name": "Bob", "age": 25},
|
||||
],
|
||||
type=pa.struct(
|
||||
[
|
||||
pa.field("name", pa.string()),
|
||||
pa.field("age", pa.int32()),
|
||||
]
|
||||
),
|
||||
)
|
||||
}
|
||||
)
|
||||
items_data = [
|
||||
{"user": {"name": "Alice", "age": 30}},
|
||||
{"user": {"name": "Bob", "age": 25}},
|
||||
]
|
||||
ds = _create_dataset(items_data, dataset_format, arrow_table)
|
||||
|
||||
result = ds.with_column("name", col("user").struct[0]).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"user": [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}],
|
||||
"name": ["Alice", "Bob"],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_struct_nested_field(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test nested struct field access with .field()."""
|
||||
arrow_table = pa.table(
|
||||
{
|
||||
"user": pa.array(
|
||||
[
|
||||
{"name": "Alice", "address": {"city": "NYC", "zip": "10001"}},
|
||||
{"name": "Bob", "address": {"city": "LA", "zip": "90001"}},
|
||||
],
|
||||
type=pa.struct(
|
||||
[
|
||||
pa.field("name", pa.string()),
|
||||
pa.field(
|
||||
"address",
|
||||
pa.struct(
|
||||
[
|
||||
pa.field("city", pa.string()),
|
||||
pa.field("zip", pa.string()),
|
||||
]
|
||||
),
|
||||
),
|
||||
]
|
||||
),
|
||||
)
|
||||
}
|
||||
)
|
||||
items_data = [
|
||||
{"user": {"name": "Alice", "address": {"city": "NYC", "zip": "10001"}}},
|
||||
{"user": {"name": "Bob", "address": {"city": "LA", "zip": "90001"}}},
|
||||
]
|
||||
ds = _create_dataset(items_data, dataset_format, arrow_table)
|
||||
|
||||
result = ds.with_column(
|
||||
"city", col("user").struct.field("address").struct.field("city")
|
||||
).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"user": [
|
||||
{"name": "Alice", "address": {"city": "NYC", "zip": "10001"}},
|
||||
{"name": "Bob", "address": {"city": "LA", "zip": "90001"}},
|
||||
],
|
||||
"city": ["NYC", "LA"],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_struct_nested_bracket(self, ray_start_regular_shared, dataset_format):
|
||||
"""Test nested struct field access with brackets."""
|
||||
arrow_table = pa.table(
|
||||
{
|
||||
"user": pa.array(
|
||||
[
|
||||
{"name": "Alice", "address": {"city": "NYC", "zip": "10001"}},
|
||||
{"name": "Bob", "address": {"city": "LA", "zip": "90001"}},
|
||||
],
|
||||
type=pa.struct(
|
||||
[
|
||||
pa.field("name", pa.string()),
|
||||
pa.field(
|
||||
"address",
|
||||
pa.struct(
|
||||
[
|
||||
pa.field("city", pa.string()),
|
||||
pa.field("zip", pa.string()),
|
||||
]
|
||||
),
|
||||
),
|
||||
]
|
||||
),
|
||||
)
|
||||
}
|
||||
)
|
||||
items_data = [
|
||||
{"user": {"name": "Alice", "address": {"city": "NYC", "zip": "10001"}}},
|
||||
{"user": {"name": "Bob", "address": {"city": "LA", "zip": "90001"}}},
|
||||
]
|
||||
ds = _create_dataset(items_data, dataset_format, arrow_table)
|
||||
|
||||
result = ds.with_column(
|
||||
"zip", col("user").struct["address"].struct["zip"]
|
||||
).to_pandas()
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"user": [
|
||||
{"name": "Alice", "address": {"city": "NYC", "zip": "10001"}},
|
||||
{"name": "Bob", "address": {"city": "LA", "zip": "90001"}},
|
||||
],
|
||||
"zip": ["10001", "90001"],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,287 @@
|
||||
"""Integration tests for predicate expression operations.
|
||||
|
||||
These tests require Ray and test end-to-end predicate expression evaluation.
|
||||
"""
|
||||
|
||||
import pandas as pd
|
||||
import pytest
|
||||
from packaging.version import parse as parse_version
|
||||
|
||||
import ray
|
||||
from ray.data._internal.util import rows_same
|
||||
from ray.data._internal.utils.arrow_utils import get_pyarrow_version
|
||||
from ray.data.expressions import col
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
get_pyarrow_version() < parse_version("20.0.0"),
|
||||
reason="Expression integration tests require PyArrow >= 20.0.0",
|
||||
)
|
||||
|
||||
|
||||
class TestPredicateIntegration:
|
||||
"""Integration tests for predicate expressions with Ray Dataset."""
|
||||
|
||||
def test_null_predicates_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test null predicate expressions with Ray Dataset."""
|
||||
ds = ray.data.from_items(
|
||||
[
|
||||
{"value": 10, "name": "Alice"},
|
||||
{"value": None, "name": "Bob"},
|
||||
{"value": 30, "name": None},
|
||||
{"value": None, "name": None},
|
||||
]
|
||||
)
|
||||
|
||||
result = (
|
||||
ds.with_column("value_is_null", col("value").is_null())
|
||||
.with_column("name_not_null", col("name").is_not_null())
|
||||
.with_column(
|
||||
"both_present", col("value").is_not_null() & col("name").is_not_null()
|
||||
)
|
||||
.to_pandas()
|
||||
)
|
||||
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"value": [10, None, 30, None],
|
||||
"name": ["Alice", "Bob", None, None],
|
||||
"value_is_null": [False, True, False, True],
|
||||
"name_not_null": [True, True, False, False],
|
||||
"both_present": [True, False, False, False],
|
||||
}
|
||||
)
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_membership_predicates_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test membership predicate expressions with Ray Dataset."""
|
||||
ds = ray.data.from_items(
|
||||
[
|
||||
{"status": "active", "category": "A"},
|
||||
{"status": "inactive", "category": "B"},
|
||||
{"status": "pending", "category": "A"},
|
||||
{"status": "deleted", "category": "C"},
|
||||
]
|
||||
)
|
||||
|
||||
result = (
|
||||
ds.with_column(
|
||||
"is_valid_status", col("status").is_in(["active", "pending"])
|
||||
)
|
||||
.with_column("not_deleted", col("status").not_in(["deleted"]))
|
||||
.with_column("category_a", col("category").is_in(["A"]))
|
||||
.to_pandas()
|
||||
)
|
||||
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"status": ["active", "inactive", "pending", "deleted"],
|
||||
"category": ["A", "B", "A", "C"],
|
||||
"is_valid_status": [True, False, True, False],
|
||||
"not_deleted": [True, True, True, False],
|
||||
"category_a": [True, False, True, False],
|
||||
}
|
||||
)
|
||||
pd.testing.assert_frame_equal(result, expected, check_dtype=False)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"test_data,expression,expected_results,test_id",
|
||||
[
|
||||
pytest.param(
|
||||
[{"value": 1}, {"value": None}, {"value": 3}],
|
||||
col("value").is_null(),
|
||||
[False, True, False],
|
||||
"is_null_with_actual_nulls",
|
||||
),
|
||||
pytest.param(
|
||||
[{"value": 1}, {"value": None}, {"value": 3}],
|
||||
col("value").is_not_null(),
|
||||
[True, False, True],
|
||||
"is_not_null_with_actual_nulls",
|
||||
),
|
||||
pytest.param(
|
||||
[{"value": 1}, {"value": 2}, {"value": 3}],
|
||||
col("value").is_in([1, 3]),
|
||||
[True, False, True],
|
||||
"isin_operation",
|
||||
),
|
||||
pytest.param(
|
||||
[{"value": 1}, {"value": 2}, {"value": 3}],
|
||||
col("value").not_in([1, 3]),
|
||||
[False, True, False],
|
||||
"not_in_operation",
|
||||
),
|
||||
pytest.param(
|
||||
[{"name": "Alice"}, {"name": "Bob"}, {"name": "Charlie"}],
|
||||
col("name") == "Bob",
|
||||
[False, True, False],
|
||||
"string_equality",
|
||||
),
|
||||
pytest.param(
|
||||
[{"name": "Alice"}, {"name": "Bob"}, {"name": "Charlie"}],
|
||||
col("name") != "Bob",
|
||||
[True, False, True],
|
||||
"string_not_equal",
|
||||
),
|
||||
pytest.param(
|
||||
[{"name": "included"}, {"name": "excluded"}, {"name": None}],
|
||||
col("name").is_not_null() & (col("name") != "excluded"),
|
||||
[True, False, False],
|
||||
"string_filter",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_null_and_membership_with_dataset(
|
||||
self, ray_start_regular_shared, test_data, expression, expected_results, test_id
|
||||
):
|
||||
"""Test null checking and membership operations with Ray Dataset."""
|
||||
ds = ray.data.from_items(test_data)
|
||||
result = ds.with_column("result", expression).to_pandas()
|
||||
|
||||
expected_data = {}
|
||||
for key in test_data[0].keys():
|
||||
expected_data[key] = [row[key] for row in test_data]
|
||||
expected_data["result"] = expected_results
|
||||
expected = pd.DataFrame(expected_data)
|
||||
|
||||
assert rows_same(result, expected)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"filter_expr,test_data,expected_flags,test_id",
|
||||
[
|
||||
pytest.param(
|
||||
col("age") >= 21,
|
||||
[
|
||||
{"age": 20, "name": "Alice"},
|
||||
{"age": 21, "name": "Bob"},
|
||||
{"age": 25, "name": "Charlie"},
|
||||
],
|
||||
[False, True, True],
|
||||
"age_filter",
|
||||
),
|
||||
pytest.param(
|
||||
col("score") > 50,
|
||||
[
|
||||
{"score": 30, "status": "fail"},
|
||||
{"score": 50, "status": "pass"},
|
||||
{"score": 70, "status": "pass"},
|
||||
],
|
||||
[False, False, True],
|
||||
"score_filter",
|
||||
),
|
||||
pytest.param(
|
||||
(col("age") >= 18) & col("active"),
|
||||
[
|
||||
{"age": 17, "active": True},
|
||||
{"age": 18, "active": False},
|
||||
{"age": 25, "active": True},
|
||||
],
|
||||
[False, False, True],
|
||||
"complex_and_filter",
|
||||
),
|
||||
pytest.param(
|
||||
(col("status") == "approved") | (col("priority") == "high"),
|
||||
[
|
||||
{"status": "pending", "priority": "low"},
|
||||
{"status": "approved", "priority": "low"},
|
||||
{"status": "pending", "priority": "high"},
|
||||
],
|
||||
[False, True, True],
|
||||
"complex_or_filter",
|
||||
),
|
||||
pytest.param(
|
||||
col("value").is_not_null() & (col("value") > 0),
|
||||
[{"value": None}, {"value": -5}, {"value": 10}],
|
||||
[False, False, True],
|
||||
"null_aware_filter",
|
||||
),
|
||||
pytest.param(
|
||||
col("name").is_not_null() & (col("name") != "excluded"),
|
||||
[{"name": "included"}, {"name": "excluded"}, {"name": None}],
|
||||
[True, False, False],
|
||||
"string_filter",
|
||||
),
|
||||
pytest.param(
|
||||
col("category").is_in(["A", "B"]),
|
||||
[
|
||||
{"category": "A"},
|
||||
{"category": "B"},
|
||||
{"category": "C"},
|
||||
{"category": "D"},
|
||||
],
|
||||
[True, True, False, False],
|
||||
"membership_filter",
|
||||
),
|
||||
pytest.param(
|
||||
(col("score") >= 50) & (col("grade") != "F"),
|
||||
[
|
||||
{"score": 45, "grade": "F"},
|
||||
{"score": 55, "grade": "D"},
|
||||
{"score": 75, "grade": "B"},
|
||||
{"score": 30, "grade": "F"},
|
||||
],
|
||||
[False, True, True, False],
|
||||
"nested_filters",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_filter_expressions_with_dataset(
|
||||
self, ray_start_regular_shared, filter_expr, test_data, expected_flags, test_id
|
||||
):
|
||||
"""Test filter expressions with Ray Dataset."""
|
||||
ds = ray.data.from_items(test_data)
|
||||
result = ds.with_column("is_filtered", filter_expr).to_pandas()
|
||||
|
||||
expected = pd.DataFrame(test_data)
|
||||
expected["is_filtered"] = expected_flags
|
||||
|
||||
assert rows_same(result, expected)
|
||||
|
||||
def test_filter_in_pipeline_with_dataset(self, ray_start_regular_shared):
|
||||
"""Test filter expressions in a data processing pipeline."""
|
||||
test_data = [
|
||||
{"product": "A", "quantity": 10, "price": 100, "region": "North"},
|
||||
{"product": "B", "quantity": 5, "price": 200, "region": "South"},
|
||||
{"product": "C", "quantity": 20, "price": 50, "region": "North"},
|
||||
{"product": "D", "quantity": 15, "price": 75, "region": "East"},
|
||||
{"product": "E", "quantity": 3, "price": 300, "region": "West"},
|
||||
]
|
||||
|
||||
ds = ray.data.from_items(test_data)
|
||||
|
||||
result = (
|
||||
ds.with_column("revenue", col("quantity") * col("price"))
|
||||
.with_column("is_high_value", col("revenue") >= 1000)
|
||||
.with_column("is_bulk_order", col("quantity") >= 10)
|
||||
.with_column("is_premium", col("price") >= 100)
|
||||
.with_column(
|
||||
"needs_special_handling",
|
||||
(col("is_high_value")) | (col("is_bulk_order") & col("is_premium")),
|
||||
)
|
||||
.with_column("is_north_region", col("region") == "North")
|
||||
.to_pandas()
|
||||
)
|
||||
|
||||
expected = pd.DataFrame(
|
||||
{
|
||||
"product": ["A", "B", "C", "D", "E"],
|
||||
"quantity": [10, 5, 20, 15, 3],
|
||||
"price": [100, 200, 50, 75, 300],
|
||||
"region": ["North", "South", "North", "East", "West"],
|
||||
"revenue": [1000, 1000, 1000, 1125, 900],
|
||||
"is_high_value": [True, True, True, True, False],
|
||||
"is_bulk_order": [True, False, True, True, False],
|
||||
"is_premium": [True, True, False, False, True],
|
||||
"needs_special_handling": [True, True, True, True, False],
|
||||
"is_north_region": [True, False, True, False, False],
|
||||
}
|
||||
)
|
||||
|
||||
assert rows_same(result, expected)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,84 @@
|
||||
# extracted from fsspec
|
||||
# https://github.com/fsspec/filesystem_spec/blob/a8cfd9c52a20c930c67ff296b60dbcda89d64db9/fsspec/tests/conftest.py
|
||||
import contextlib
|
||||
import threading
|
||||
from http.server import BaseHTTPRequestHandler, HTTPServer
|
||||
|
||||
import pytest
|
||||
|
||||
port = 9898
|
||||
data = b"\n".join([b"some test data"] * 1000)
|
||||
data_file = "http://localhost:%i/index/data_file" % port
|
||||
index = b'<a href="%s">Link</a>' % data_file.encode()
|
||||
|
||||
|
||||
class HTTPTestHandler(BaseHTTPRequestHandler):
|
||||
files = {
|
||||
"/index/data_file": data,
|
||||
"/index": index,
|
||||
}
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def _respond(self, code=200, headers=None, data=b""):
|
||||
headers = headers or {}
|
||||
headers.update({"User-Agent": "test"})
|
||||
self.send_response(code)
|
||||
for k, v in headers.items():
|
||||
self.send_header(k, str(v))
|
||||
self.end_headers()
|
||||
if data:
|
||||
self.wfile.write(data)
|
||||
|
||||
def do_GET(self):
|
||||
file_path = self.path.rstrip("/")
|
||||
file_data = self.files.get(file_path)
|
||||
if file_data is None:
|
||||
return self._respond(404)
|
||||
if "Range" in self.headers:
|
||||
ran = self.headers["Range"]
|
||||
b, ran = ran.split("=")
|
||||
start, end = ran.split("-")
|
||||
if start:
|
||||
file_data = file_data[int(start) : (int(end) + 1) if end else None]
|
||||
else:
|
||||
# suffix only
|
||||
file_data = file_data[-int(end) :]
|
||||
if "give_length" in self.headers:
|
||||
response_headers = {"Content-Length": len(file_data)}
|
||||
self._respond(200, response_headers, file_data)
|
||||
elif "give_range" in self.headers:
|
||||
self._respond(
|
||||
200,
|
||||
{"Content-Range": "0-%i/%i" % (len(file_data) - 1, len(file_data))},
|
||||
file_data,
|
||||
)
|
||||
else:
|
||||
self._respond(200, data=file_data)
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def serve():
|
||||
server_address = ("", port)
|
||||
httpd = HTTPServer(server_address, HTTPTestHandler)
|
||||
th = threading.Thread(target=httpd.serve_forever)
|
||||
th.daemon = True
|
||||
th.start()
|
||||
try:
|
||||
yield "http://localhost:%i" % port
|
||||
finally:
|
||||
httpd.socket.close()
|
||||
httpd.shutdown()
|
||||
th.join()
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def http_server():
|
||||
with serve() as s:
|
||||
yield s
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def http_file():
|
||||
return data_file
|
||||
@@ -0,0 +1,126 @@
|
||||
import shutil
|
||||
import signal
|
||||
import socket
|
||||
import subprocess as sp
|
||||
import time
|
||||
|
||||
# extracted from aioboto3
|
||||
# https://github.com/terrycain/aioboto3/blob/16a1a1085191ebe6d40ee45d9588b2173738af0c/tests/mock_server.py
|
||||
import pytest
|
||||
import requests
|
||||
|
||||
from ray._common.network_utils import build_address
|
||||
|
||||
_proxy_bypass = {
|
||||
"http": None,
|
||||
"https": None,
|
||||
}
|
||||
|
||||
|
||||
def _is_port_available(host, port):
|
||||
"""Check if a port is available for use."""
|
||||
try:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind((host, port))
|
||||
return True
|
||||
except OSError:
|
||||
return False
|
||||
|
||||
|
||||
def _find_available_port(host, preferred_port, max_attempts=10):
|
||||
"""Find an available port starting from preferred_port."""
|
||||
|
||||
# Try the preferred port first
|
||||
if _is_port_available(host, preferred_port):
|
||||
return preferred_port
|
||||
|
||||
# Try a wider range if preferred port is busy
|
||||
for i in range(1, max_attempts):
|
||||
port = preferred_port + i
|
||||
if _is_port_available(host, port):
|
||||
return port
|
||||
|
||||
# If all else fails, let the OS pick a port
|
||||
try:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
||||
s.bind((host, 0)) # Let OS pick port
|
||||
_, port = s.getsockname()
|
||||
return port
|
||||
except OSError as e:
|
||||
raise RuntimeError(
|
||||
f"Could not find any available port starting from " f"{preferred_port}: {e}"
|
||||
) from e
|
||||
|
||||
|
||||
def start_service(service_name, host, port):
|
||||
moto_svr_path = shutil.which("moto_server")
|
||||
if not moto_svr_path:
|
||||
pytest.skip("moto not installed")
|
||||
|
||||
# Always use port conflict resolution to be safe
|
||||
port = _find_available_port(host, port)
|
||||
|
||||
# moto 5.x no longer accepts a service name argument - all services
|
||||
# are served on a single endpoint
|
||||
args = [moto_svr_path, "-H", host, "-p", str(port)]
|
||||
# For debugging
|
||||
# args = '{0} {1} -H {2} -p {3} 2>&1 | \
|
||||
# tee -a /tmp/moto.log'.format(moto_svr_path, service_name, host, port)
|
||||
process = sp.Popen(
|
||||
args, stdin=sp.PIPE, stdout=sp.PIPE, stderr=sp.PIPE
|
||||
) # shell=True
|
||||
url = f"http://{build_address(host, port)}"
|
||||
|
||||
for i in range(0, 30):
|
||||
output = process.poll()
|
||||
if output is not None:
|
||||
print("moto_server exited status {0}".format(output))
|
||||
stdout, stderr = process.communicate()
|
||||
print("moto_server stdout: {0}".format(stdout))
|
||||
print("moto_server stderr: {0}".format(stderr))
|
||||
pytest.fail("Can not start service: {}".format(service_name))
|
||||
|
||||
try:
|
||||
# we need to bypass the proxies due to monkeypatches
|
||||
requests.get(url, timeout=5, proxies=_proxy_bypass)
|
||||
break
|
||||
except requests.exceptions.ConnectionError:
|
||||
time.sleep(0.5)
|
||||
else:
|
||||
stop_process(process) # pytest.fail doesn't call stop_process
|
||||
pytest.fail("Can not start service: {}".format(service_name))
|
||||
|
||||
return process, url
|
||||
|
||||
|
||||
def stop_process(process):
|
||||
"""Stop process with shorter timeout to prevent test hangs."""
|
||||
if process is None or process.poll() is not None:
|
||||
return # Already stopped
|
||||
|
||||
try:
|
||||
process.send_signal(signal.SIGTERM)
|
||||
process.communicate(timeout=20)
|
||||
except sp.TimeoutExpired:
|
||||
process.kill()
|
||||
try:
|
||||
process.communicate(timeout=5) # Short timeout for kill
|
||||
except sp.TimeoutExpired:
|
||||
print("Warning: Process cleanup timed out")
|
||||
except Exception as e:
|
||||
print(f"Warning: Error during process cleanup: {e}")
|
||||
|
||||
|
||||
# TODO(Clark): We should be able to use "session" scope here, but we've found
|
||||
# that the s3_fs fixture ends up hanging with S3 ops timing out (or the server
|
||||
# being unreachable). This appears to only be an issue when using the tmp_dir
|
||||
# fixture as the S3 dir path. We should fix this since "session" scope should
|
||||
# reduce a lot of the per-test overhead (2x faster execution for IO methods in
|
||||
# test_dataset.py).
|
||||
@pytest.fixture(scope="function")
|
||||
def s3_server():
|
||||
host = "localhost"
|
||||
port = 5002
|
||||
process, url = start_service("s3", host, port)
|
||||
yield url
|
||||
stop_process(process)
|
||||
@@ -0,0 +1,764 @@
|
||||
"""
|
||||
Backwards compatibility tests for Preprocessor private field renaming.
|
||||
|
||||
These tests verify that preprocessors pickled with old public field names
|
||||
(e.g., 'columns') can be deserialized correctly after fields were renamed
|
||||
to private (e.g., '_columns').
|
||||
|
||||
The __setstate__ method in each preprocessor handles migration automatically.
|
||||
"""
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data.preprocessors import (
|
||||
Categorizer,
|
||||
Chain,
|
||||
Concatenator,
|
||||
CountVectorizer,
|
||||
CustomKBinsDiscretizer,
|
||||
FeatureHasher,
|
||||
HashingVectorizer,
|
||||
LabelEncoder,
|
||||
MaxAbsScaler,
|
||||
MinMaxScaler,
|
||||
MultiHotEncoder,
|
||||
Normalizer,
|
||||
OneHotEncoder,
|
||||
OrdinalEncoder,
|
||||
PowerTransformer,
|
||||
RobustScaler,
|
||||
SimpleImputer,
|
||||
StandardScaler,
|
||||
Tokenizer,
|
||||
TorchVisionPreprocessor,
|
||||
UniformKBinsDiscretizer,
|
||||
)
|
||||
|
||||
# =============================================================================
|
||||
# Field Migration Tests
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"preprocessor_class,old_state,expected_attrs",
|
||||
[
|
||||
(
|
||||
Concatenator,
|
||||
{
|
||||
"columns": ["A", "B"],
|
||||
"output_column_name": "concat",
|
||||
"dtype": np.float32,
|
||||
"raise_if_missing": True,
|
||||
"flatten": True,
|
||||
},
|
||||
{
|
||||
"columns": ["A", "B"],
|
||||
"output_column_name": "concat",
|
||||
"dtype": np.float32,
|
||||
"raise_if_missing": True,
|
||||
"flatten": True,
|
||||
},
|
||||
),
|
||||
(
|
||||
Normalizer,
|
||||
{
|
||||
"columns": ["A", "B"],
|
||||
"norm": "l1",
|
||||
"output_columns": ["A_norm", "B_norm"],
|
||||
},
|
||||
{
|
||||
"columns": ["A", "B"],
|
||||
"norm": "l1",
|
||||
"output_columns": ["A_norm", "B_norm"],
|
||||
},
|
||||
),
|
||||
(
|
||||
Tokenizer,
|
||||
{
|
||||
"columns": ["text"],
|
||||
"tokenization_fn": lambda s: s.split(),
|
||||
"output_columns": ["text_tokens"],
|
||||
},
|
||||
{
|
||||
"columns": ["text"],
|
||||
"tokenization_fn": "callable", # Special marker
|
||||
"output_columns": ["text_tokens"],
|
||||
},
|
||||
),
|
||||
(
|
||||
PowerTransformer,
|
||||
{
|
||||
"columns": ["A", "B"],
|
||||
"power": 3,
|
||||
"method": "box-cox",
|
||||
"output_columns": ["A_pow", "B_pow"],
|
||||
},
|
||||
{
|
||||
"columns": ["A", "B"],
|
||||
"power": 3,
|
||||
"method": "box-cox",
|
||||
"output_columns": ["A_pow", "B_pow"],
|
||||
},
|
||||
),
|
||||
(
|
||||
HashingVectorizer,
|
||||
{
|
||||
"columns": ["text"],
|
||||
"num_features": 200,
|
||||
"tokenization_fn": lambda s: s.split(),
|
||||
"output_columns": ["text_vec"],
|
||||
},
|
||||
{
|
||||
"columns": ["text"],
|
||||
"num_features": 200,
|
||||
"tokenization_fn": "callable",
|
||||
"output_columns": ["text_vec"],
|
||||
},
|
||||
),
|
||||
(
|
||||
CountVectorizer,
|
||||
{
|
||||
"columns": ["text"],
|
||||
"tokenization_fn": lambda s: s.split(),
|
||||
"max_features": 100,
|
||||
"output_columns": ["text_count"],
|
||||
},
|
||||
{
|
||||
"columns": ["text"],
|
||||
"tokenization_fn": "callable",
|
||||
"max_features": 100,
|
||||
"output_columns": ["text_count"],
|
||||
},
|
||||
),
|
||||
(
|
||||
FeatureHasher,
|
||||
{"columns": ["A", "B"], "num_features": 20, "output_column": "features"},
|
||||
{"columns": ["A", "B"], "num_features": 20, "output_column": "features"},
|
||||
),
|
||||
(
|
||||
OrdinalEncoder,
|
||||
{
|
||||
"columns": ["color"],
|
||||
"output_columns": ["color_encoded"],
|
||||
"encode_lists": False,
|
||||
},
|
||||
{
|
||||
"columns": ["color"],
|
||||
"output_columns": ["color_encoded"],
|
||||
"encode_lists": False,
|
||||
},
|
||||
),
|
||||
(
|
||||
OneHotEncoder,
|
||||
{
|
||||
"columns": ["color"],
|
||||
"output_columns": ["color_encoded"],
|
||||
"max_categories": {"color": 5},
|
||||
},
|
||||
{
|
||||
"columns": ["color"],
|
||||
"output_columns": ["color_encoded"],
|
||||
"max_categories": {"color": 5},
|
||||
},
|
||||
),
|
||||
(
|
||||
MultiHotEncoder,
|
||||
{
|
||||
"columns": ["tags"],
|
||||
"output_columns": ["tags_encoded"],
|
||||
"max_categories": {},
|
||||
},
|
||||
{"columns": ["tags"], "output_columns": ["tags_encoded"]},
|
||||
),
|
||||
(
|
||||
LabelEncoder,
|
||||
{"label_column": "label", "output_column": "label_id"},
|
||||
{"label_column": "label", "output_column": "label_id"},
|
||||
),
|
||||
(
|
||||
Categorizer,
|
||||
{"columns": ["sex"], "output_columns": ["sex_cat"], "dtypes": {}},
|
||||
{"columns": ["sex"], "output_columns": ["sex_cat"]},
|
||||
),
|
||||
(
|
||||
StandardScaler,
|
||||
{"columns": ["A", "B"], "output_columns": ["A_scaled", "B_scaled"]},
|
||||
{"columns": ["A", "B"], "output_columns": ["A_scaled", "B_scaled"]},
|
||||
),
|
||||
(
|
||||
MinMaxScaler,
|
||||
{"columns": ["A"], "output_columns": ["A_scaled"]},
|
||||
{"columns": ["A"], "output_columns": ["A_scaled"]},
|
||||
),
|
||||
(
|
||||
MaxAbsScaler,
|
||||
{"columns": ["A"], "output_columns": ["A_scaled"]},
|
||||
{"columns": ["A"], "output_columns": ["A_scaled"]},
|
||||
),
|
||||
(
|
||||
RobustScaler,
|
||||
{
|
||||
"columns": ["A"],
|
||||
"output_columns": ["A_scaled"],
|
||||
"quantile_range": (0.1, 0.9),
|
||||
"quantile_precision": 1000,
|
||||
},
|
||||
{
|
||||
"columns": ["A"],
|
||||
"output_columns": ["A_scaled"],
|
||||
"quantile_range": (0.1, 0.9),
|
||||
"quantile_precision": 1000,
|
||||
},
|
||||
),
|
||||
(
|
||||
SimpleImputer,
|
||||
{
|
||||
"columns": ["A", "B"],
|
||||
"output_columns": ["A_imputed", "B_imputed"],
|
||||
"strategy": "median",
|
||||
"fill_value": 99.0,
|
||||
},
|
||||
{
|
||||
"columns": ["A", "B"],
|
||||
"output_columns": ["A_imputed", "B_imputed"],
|
||||
"strategy": "median",
|
||||
"fill_value": 99.0,
|
||||
},
|
||||
),
|
||||
(
|
||||
CustomKBinsDiscretizer,
|
||||
{
|
||||
"columns": ["A", "B"],
|
||||
"bins": [0, 1, 2, 3],
|
||||
"right": False,
|
||||
"include_lowest": True,
|
||||
"duplicates": "drop",
|
||||
"dtypes": None,
|
||||
"output_columns": ["A_binned", "B_binned"],
|
||||
},
|
||||
{
|
||||
"columns": ["A", "B"],
|
||||
"bins": [0, 1, 2, 3],
|
||||
"right": False,
|
||||
"include_lowest": True,
|
||||
"duplicates": "drop",
|
||||
"dtypes": None,
|
||||
"output_columns": ["A_binned", "B_binned"],
|
||||
},
|
||||
),
|
||||
(
|
||||
UniformKBinsDiscretizer,
|
||||
{
|
||||
"columns": ["A", "B"],
|
||||
"bins": 4,
|
||||
"right": False,
|
||||
"include_lowest": True,
|
||||
"duplicates": "drop",
|
||||
"dtypes": None,
|
||||
"output_columns": ["A_binned", "B_binned"],
|
||||
},
|
||||
{
|
||||
"columns": ["A", "B"],
|
||||
"bins": 4,
|
||||
"right": False,
|
||||
"include_lowest": True,
|
||||
"duplicates": "drop",
|
||||
"dtypes": None,
|
||||
"output_columns": ["A_binned", "B_binned"],
|
||||
},
|
||||
),
|
||||
],
|
||||
ids=lambda x: x.__name__ if hasattr(x, "__name__") else str(x)[:20],
|
||||
)
|
||||
def test_field_migration_from_old_public_names(
|
||||
preprocessor_class, old_state, expected_attrs
|
||||
):
|
||||
"""Verify old public field names are migrated to new private fields."""
|
||||
preprocessor = preprocessor_class.__new__(preprocessor_class)
|
||||
preprocessor.__setstate__(old_state)
|
||||
|
||||
for attr_name, expected_value in expected_attrs.items():
|
||||
actual_value = getattr(preprocessor, attr_name)
|
||||
if expected_value == "callable":
|
||||
assert callable(actual_value), f"{attr_name} should be callable"
|
||||
else:
|
||||
assert actual_value == expected_value, f"Mismatch in {attr_name}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"preprocessor_class,minimal_state,expected_defaults",
|
||||
[
|
||||
# Callable default: tokenization_fn must be stored as the function itself,
|
||||
# not called. This would have failed with the old callable() check.
|
||||
(
|
||||
Tokenizer,
|
||||
{"columns": ["text"]},
|
||||
{"tokenization_fn": "callable", "output_columns": ["text"]},
|
||||
),
|
||||
(
|
||||
HashingVectorizer,
|
||||
{"columns": ["text"], "num_features": 100},
|
||||
{"tokenization_fn": "callable", "output_columns": ["text"]},
|
||||
),
|
||||
(
|
||||
CountVectorizer,
|
||||
{"columns": ["text"]},
|
||||
{"tokenization_fn": "callable", "output_columns": ["text"]},
|
||||
),
|
||||
# _Computed default: output_columns derives from _columns.
|
||||
(
|
||||
StandardScaler,
|
||||
{"columns": ["A", "B"]},
|
||||
{"output_columns": ["A", "B"]},
|
||||
),
|
||||
# _Computed default deriving from a different source field.
|
||||
(
|
||||
LabelEncoder,
|
||||
{"label_column": "label"},
|
||||
{"output_column": "label"},
|
||||
),
|
||||
# Plain value default alongside a _Computed default.
|
||||
(
|
||||
Normalizer,
|
||||
{"columns": ["A", "B"]},
|
||||
{"norm": "l2", "output_columns": ["A", "B"]},
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"Tokenizer",
|
||||
"HashingVectorizer",
|
||||
"CountVectorizer",
|
||||
"StandardScaler",
|
||||
"LabelEncoder",
|
||||
"Normalizer",
|
||||
],
|
||||
)
|
||||
def test_missing_optional_fields_use_defaults(
|
||||
preprocessor_class, minimal_state, expected_defaults
|
||||
):
|
||||
"""
|
||||
Verify that absent optional fields are filled with their correct defaults.
|
||||
|
||||
This exercises the default-fallback branch of migrate_private_fields. The minimal_state
|
||||
deliberately omits optional fields to force the default path.
|
||||
"""
|
||||
preprocessor = preprocessor_class.__new__(preprocessor_class)
|
||||
preprocessor.__setstate__(minimal_state)
|
||||
|
||||
for attr_name, expected_value in expected_defaults.items():
|
||||
actual_value = getattr(preprocessor, attr_name)
|
||||
if expected_value == "callable":
|
||||
assert callable(
|
||||
actual_value
|
||||
), f"{attr_name} should be a stored callable, not the result of calling it"
|
||||
else:
|
||||
assert (
|
||||
actual_value == expected_value
|
||||
), f"Mismatch in {attr_name}: {actual_value!r} != {expected_value!r}"
|
||||
|
||||
|
||||
def test_torchvision_preprocessor_field_migration():
|
||||
try:
|
||||
from torchvision import transforms
|
||||
except ImportError:
|
||||
pytest.skip("torchvision not installed")
|
||||
|
||||
transform = transforms.Lambda(lambda x: x)
|
||||
preprocessor = TorchVisionPreprocessor.__new__(TorchVisionPreprocessor)
|
||||
state = {
|
||||
"columns": ["image"],
|
||||
"output_columns": ["image_out"],
|
||||
"torchvision_transform": transform,
|
||||
"batched": True,
|
||||
}
|
||||
preprocessor.__setstate__(state)
|
||||
|
||||
assert preprocessor.columns == ["image"]
|
||||
assert preprocessor.output_columns == ["image_out"]
|
||||
assert preprocessor.torchvision_transform == transform
|
||||
assert preprocessor.batched is True
|
||||
|
||||
|
||||
def test_chain_field_migration():
|
||||
scaler1 = StandardScaler(columns=["A"])
|
||||
scaler2 = StandardScaler(columns=["B"])
|
||||
chain = Chain.__new__(Chain)
|
||||
|
||||
state = {"preprocessors": (scaler1, scaler2)}
|
||||
chain.__setstate__(state)
|
||||
|
||||
assert len(chain.preprocessors) == 2
|
||||
assert chain.preprocessors[0] == scaler1
|
||||
assert chain.preprocessors[1] == scaler2
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Functional Test Helpers
|
||||
# =============================================================================
|
||||
|
||||
|
||||
def _simulate_old_format_deserialization(preprocessor, field_mapping):
|
||||
"""Simulate deserialization from old format by renaming private->public fields."""
|
||||
state = preprocessor.__dict__.copy()
|
||||
for public_name, private_name in field_mapping.items():
|
||||
if private_name in state:
|
||||
state[public_name] = state.pop(private_name)
|
||||
|
||||
new_preprocessor = preprocessor.__class__.__new__(preprocessor.__class__)
|
||||
new_preprocessor.__setstate__(state)
|
||||
return new_preprocessor
|
||||
|
||||
|
||||
def _test_functional_backwards_compat(preprocessor, test_ds, field_mapping):
|
||||
"""Generic functional test: verify deserialized preprocessor produces same output."""
|
||||
expected_result = preprocessor.transform(test_ds).to_pandas()
|
||||
new_preprocessor = _simulate_old_format_deserialization(preprocessor, field_mapping)
|
||||
result = new_preprocessor.transform(test_ds).to_pandas()
|
||||
pd.testing.assert_frame_equal(result, expected_result)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Functional Tests - Simple Preprocessors (No Fitting Required)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"setup_func,field_mapping",
|
||||
[
|
||||
(
|
||||
lambda: (
|
||||
Concatenator(columns=["A", "B"], output_column_name="C"),
|
||||
pd.DataFrame({"A": [1, 2], "B": [3, 4]}),
|
||||
{
|
||||
"columns": "_columns",
|
||||
"output_column_name": "_output_column_name",
|
||||
"dtype": "_dtype",
|
||||
"raise_if_missing": "_raise_if_missing",
|
||||
"flatten": "_flatten",
|
||||
},
|
||||
),
|
||||
None,
|
||||
),
|
||||
(
|
||||
lambda: (
|
||||
Normalizer(columns=["A", "B"], norm="l2"),
|
||||
pd.DataFrame({"A": [1.0, 2.0], "B": [3.0, 4.0]}),
|
||||
{
|
||||
"columns": "_columns",
|
||||
"norm": "_norm",
|
||||
"output_columns": "_output_columns",
|
||||
},
|
||||
),
|
||||
None,
|
||||
),
|
||||
(
|
||||
lambda: (
|
||||
Tokenizer(columns=["text"]),
|
||||
pd.DataFrame({"text": ["hello world", "foo bar"]}),
|
||||
{
|
||||
"columns": "_columns",
|
||||
"tokenization_fn": "_tokenization_fn",
|
||||
"output_columns": "_output_columns",
|
||||
},
|
||||
),
|
||||
None,
|
||||
),
|
||||
(
|
||||
lambda: (
|
||||
PowerTransformer(columns=["A", "B"], power=2),
|
||||
pd.DataFrame({"A": [1.0, 2.0, 3.0], "B": [4.0, 5.0, 6.0]}),
|
||||
{
|
||||
"columns": "_columns",
|
||||
"power": "_power",
|
||||
"method": "_method",
|
||||
"output_columns": "_output_columns",
|
||||
},
|
||||
),
|
||||
None,
|
||||
),
|
||||
(
|
||||
lambda: (
|
||||
HashingVectorizer(columns=["text"], num_features=10),
|
||||
pd.DataFrame({"text": ["hello world", "foo bar"]}),
|
||||
{
|
||||
"columns": "_columns",
|
||||
"num_features": "_num_features",
|
||||
"tokenization_fn": "_tokenization_fn",
|
||||
"output_columns": "_output_columns",
|
||||
},
|
||||
),
|
||||
None,
|
||||
),
|
||||
(
|
||||
lambda: (
|
||||
FeatureHasher(
|
||||
columns=["token_a", "token_b"],
|
||||
num_features=5,
|
||||
output_column="hashed",
|
||||
),
|
||||
pd.DataFrame({"token_a": [1, 2], "token_b": [3, 4]}),
|
||||
{
|
||||
"columns": "_columns",
|
||||
"num_features": "_num_features",
|
||||
"output_column": "_output_column",
|
||||
},
|
||||
),
|
||||
None,
|
||||
),
|
||||
(
|
||||
lambda: (
|
||||
CustomKBinsDiscretizer(
|
||||
columns=["A", "B"],
|
||||
bins=[0, 1, 2, 3, 4],
|
||||
output_columns=["A_binned", "B_binned"],
|
||||
),
|
||||
pd.DataFrame({"A": [0.5, 1.5, 2.5, 3.5], "B": [0.2, 1.2, 2.2, 3.2]}),
|
||||
{
|
||||
"columns": "_columns",
|
||||
"bins": "_bins",
|
||||
"right": "_right",
|
||||
"include_lowest": "_include_lowest",
|
||||
"duplicates": "_duplicates",
|
||||
"dtypes": "_dtypes",
|
||||
"output_columns": "_output_columns",
|
||||
},
|
||||
),
|
||||
None,
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"Concatenator",
|
||||
"Normalizer",
|
||||
"Tokenizer",
|
||||
"PowerTransformer",
|
||||
"HashingVectorizer",
|
||||
"FeatureHasher",
|
||||
"CustomKBinsDiscretizer",
|
||||
],
|
||||
)
|
||||
def test_simple_functional_backwards_compat(setup_func, field_mapping):
|
||||
"""Verify preprocessors that don't need fitting work after deserialization."""
|
||||
preprocessor, test_data, field_mapping = setup_func()
|
||||
test_ds = ray.data.from_pandas(test_data)
|
||||
_test_functional_backwards_compat(preprocessor, test_ds, field_mapping)
|
||||
|
||||
|
||||
# =============================================================================
|
||||
# Functional Tests - Stateful Preprocessors (Require Fitting)
|
||||
# =============================================================================
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"setup_func",
|
||||
[
|
||||
lambda: (
|
||||
OrdinalEncoder(columns=["color"]),
|
||||
pd.DataFrame({"color": ["red", "green", "blue", "red", "green"]}),
|
||||
{
|
||||
"columns": "_columns",
|
||||
"output_columns": "_output_columns",
|
||||
"encode_lists": "_encode_lists",
|
||||
},
|
||||
),
|
||||
lambda: (
|
||||
OneHotEncoder(columns=["color"]),
|
||||
pd.DataFrame({"color": ["red", "green", "blue", "red", "green", "blue"]}),
|
||||
{
|
||||
"columns": "_columns",
|
||||
"output_columns": "_output_columns",
|
||||
"max_categories": "_max_categories",
|
||||
},
|
||||
),
|
||||
lambda: (
|
||||
LabelEncoder(label_column="label"),
|
||||
pd.DataFrame(
|
||||
{
|
||||
"feature": [1.0, 2.0, 3.0, 4.0],
|
||||
"label": ["cat", "dog", "cat", "bird"],
|
||||
}
|
||||
),
|
||||
{"label_column": "_label_column", "output_column": "_output_column"},
|
||||
),
|
||||
lambda: (
|
||||
StandardScaler(columns=["A", "B"]),
|
||||
pd.DataFrame(
|
||||
{"A": [1.0, 2.0, 3.0, 4.0, 5.0], "B": [10.0, 20.0, 30.0, 40.0, 50.0]}
|
||||
),
|
||||
{"columns": "_columns", "output_columns": "_output_columns"},
|
||||
),
|
||||
lambda: (
|
||||
MinMaxScaler(columns=["A", "B"]),
|
||||
pd.DataFrame(
|
||||
{"A": [1.0, 2.0, 3.0, 4.0, 5.0], "B": [10.0, 20.0, 30.0, 40.0, 50.0]}
|
||||
),
|
||||
{"columns": "_columns", "output_columns": "_output_columns"},
|
||||
),
|
||||
lambda: (
|
||||
RobustScaler(columns=["A"]),
|
||||
pd.DataFrame({"A": [1.0, 2.0, 3.0, 4.0, 5.0, 100.0]}),
|
||||
{
|
||||
"columns": "_columns",
|
||||
"output_columns": "_output_columns",
|
||||
"quantile_range": "_quantile_range",
|
||||
"quantile_precision": "_quantile_precision",
|
||||
},
|
||||
),
|
||||
lambda: (
|
||||
SimpleImputer(columns=["A", "B"], strategy="mean"),
|
||||
pd.DataFrame(
|
||||
{"A": [1.0, 2.0, None, 4.0, 5.0], "B": [10.0, None, 30.0, 40.0, 50.0]}
|
||||
),
|
||||
{
|
||||
"columns": "_columns",
|
||||
"output_columns": "_output_columns",
|
||||
"strategy": "_strategy",
|
||||
"fill_value": "_fill_value",
|
||||
},
|
||||
),
|
||||
lambda: (
|
||||
CountVectorizer(columns=["text"]),
|
||||
pd.DataFrame({"text": ["hello world", "foo bar", "hello foo"]}),
|
||||
{
|
||||
"columns": "_columns",
|
||||
"tokenization_fn": "_tokenization_fn",
|
||||
"max_features": "_max_features",
|
||||
"output_columns": "_output_columns",
|
||||
},
|
||||
),
|
||||
lambda: (
|
||||
UniformKBinsDiscretizer(
|
||||
columns=["A", "B"], bins=3, output_columns=["A_binned", "B_binned"]
|
||||
),
|
||||
pd.DataFrame(
|
||||
{"A": [1.0, 2.0, 3.0, 4.0, 5.0], "B": [10.0, 20.0, 30.0, 40.0, 50.0]}
|
||||
),
|
||||
{
|
||||
"columns": "_columns",
|
||||
"bins": "_bins",
|
||||
"right": "_right",
|
||||
"include_lowest": "_include_lowest",
|
||||
"duplicates": "_duplicates",
|
||||
"dtypes": "_dtypes",
|
||||
"output_columns": "_output_columns",
|
||||
},
|
||||
),
|
||||
lambda: (
|
||||
MultiHotEncoder(columns=["genre"]),
|
||||
pd.DataFrame(
|
||||
{
|
||||
"genre": [
|
||||
["comedy", "action"],
|
||||
["drama", "action"],
|
||||
["comedy", "drama"],
|
||||
]
|
||||
}
|
||||
),
|
||||
{
|
||||
"columns": "_columns",
|
||||
"output_columns": "_output_columns",
|
||||
"max_categories": "_max_categories",
|
||||
},
|
||||
),
|
||||
lambda: (
|
||||
MaxAbsScaler(columns=["A", "B"]),
|
||||
pd.DataFrame({"A": [-6.0, 3.0, -3.0], "B": [2.0, -4.0, 1.0]}),
|
||||
{"columns": "_columns", "output_columns": "_output_columns"},
|
||||
),
|
||||
lambda: (
|
||||
Categorizer(columns=["color"]),
|
||||
pd.DataFrame({"color": ["red", "green", "blue", "red", "green"]}),
|
||||
{
|
||||
"columns": "_columns",
|
||||
"output_columns": "_output_columns",
|
||||
"dtypes": "_dtypes",
|
||||
},
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"OrdinalEncoder",
|
||||
"OneHotEncoder",
|
||||
"LabelEncoder",
|
||||
"StandardScaler",
|
||||
"MinMaxScaler",
|
||||
"RobustScaler",
|
||||
"SimpleImputer",
|
||||
"CountVectorizer",
|
||||
"UniformKBinsDiscretizer",
|
||||
"MultiHotEncoder",
|
||||
"MaxAbsScaler",
|
||||
"Categorizer",
|
||||
],
|
||||
)
|
||||
def test_stateful_functional_backwards_compat(setup_func):
|
||||
"""Verify fitted preprocessors work after deserialization."""
|
||||
preprocessor, test_data, field_mapping = setup_func()
|
||||
test_ds = ray.data.from_pandas(test_data)
|
||||
preprocessor = preprocessor.fit(test_ds)
|
||||
_test_functional_backwards_compat(preprocessor, test_ds, field_mapping)
|
||||
|
||||
|
||||
def test_chain_functional_backwards_compat():
|
||||
df = pd.DataFrame({"A": [1.0, 2.0, 3.0]})
|
||||
ds = ray.data.from_pandas(df)
|
||||
|
||||
scaler = StandardScaler(columns=["A"])
|
||||
normalizer = Normalizer(columns=["A"])
|
||||
chain = Chain(scaler, normalizer)
|
||||
chain = chain.fit(ds)
|
||||
|
||||
expected_result = chain.transform(ds).to_pandas()
|
||||
|
||||
state = chain.__dict__.copy()
|
||||
state["preprocessors"] = state.pop("_preprocessors")
|
||||
|
||||
new_chain = Chain.__new__(Chain)
|
||||
new_chain.__setstate__(state)
|
||||
|
||||
result = new_chain.transform(ds).to_pandas()
|
||||
pd.testing.assert_frame_equal(result, expected_result)
|
||||
|
||||
|
||||
def test_torchvision_functional_backwards_compat():
|
||||
try:
|
||||
import torch
|
||||
from torchvision import transforms
|
||||
except ImportError:
|
||||
pytest.skip("torchvision not installed")
|
||||
|
||||
transform = transforms.Lambda(lambda x: torch.as_tensor(x, dtype=torch.float32))
|
||||
df = pd.DataFrame(
|
||||
{
|
||||
"image": [
|
||||
np.array([[1, 2], [3, 4]], dtype=np.uint8),
|
||||
np.array([[5, 6], [7, 8]], dtype=np.uint8),
|
||||
]
|
||||
}
|
||||
)
|
||||
ds = ray.data.from_pandas(df)
|
||||
|
||||
preprocessor = TorchVisionPreprocessor(
|
||||
columns=["image"], transform=transform, batched=False
|
||||
)
|
||||
expected_result = preprocessor.transform(ds).to_pandas()
|
||||
|
||||
state = preprocessor.__dict__.copy()
|
||||
state["columns"] = state.pop("_columns")
|
||||
state["output_columns"] = state.pop("_output_columns")
|
||||
state["torchvision_transform"] = state.pop("_torchvision_transform")
|
||||
state["batched"] = state.pop("_batched")
|
||||
|
||||
new_preprocessor = TorchVisionPreprocessor.__new__(TorchVisionPreprocessor)
|
||||
new_preprocessor.__setstate__(state)
|
||||
|
||||
result = new_preprocessor.transform(ds).to_pandas()
|
||||
assert len(result) == len(expected_result)
|
||||
assert "image" in result.columns
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,243 @@
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data.preprocessor import Preprocessor
|
||||
from ray.data.preprocessors import Chain, LabelEncoder, SimpleImputer, StandardScaler
|
||||
from ray.data.util.data_batch_conversion import BatchFormat
|
||||
|
||||
|
||||
def test_chain():
|
||||
"""Tests basic Chain functionality."""
|
||||
col_a = [-1, -1, 1, 1]
|
||||
col_b = [1, 1, 1, None]
|
||||
col_c = ["sunday", "monday", "tuesday", "tuesday"]
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
||||
ds = ray.data.from_pandas(in_df)
|
||||
|
||||
imputer = SimpleImputer(["B"])
|
||||
scaler = StandardScaler(["A", "B"])
|
||||
encoder = LabelEncoder("C")
|
||||
chain = Chain(scaler, imputer, encoder)
|
||||
|
||||
# Fit data.
|
||||
chain.fit(ds)
|
||||
# Transform data.
|
||||
transformed = chain.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
assert imputer.stats_ == {
|
||||
"mean(B)": 0.0,
|
||||
}
|
||||
assert scaler.stats_ == {
|
||||
"mean(A)": 0.0,
|
||||
"mean(B)": 1.0,
|
||||
"std(A)": 1.0,
|
||||
"std(B)": 0.0,
|
||||
}
|
||||
assert encoder.stats_ == {
|
||||
"unique_values(C)": {"monday": 0, "sunday": 1, "tuesday": 2}
|
||||
}
|
||||
|
||||
processed_col_a = [-1.0, -1.0, 1.0, 1.0]
|
||||
processed_col_b = [0.0, 0.0, 0.0, 0.0]
|
||||
processed_col_c = [1, 0, 2, 2]
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{"A": processed_col_a, "B": processed_col_b, "C": processed_col_c}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
# Transform batch.
|
||||
pred_col_a = [1, 2, None]
|
||||
pred_col_b = [0, None, 2]
|
||||
pred_col_c = ["monday", "tuesday", "wednesday"]
|
||||
pred_in_df = pd.DataFrame.from_dict(
|
||||
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c}
|
||||
)
|
||||
|
||||
pred_out_df = chain.transform_batch(pred_in_df)
|
||||
|
||||
pred_processed_col_a = [1, 2, None]
|
||||
pred_processed_col_b = [-1.0, 0.0, 1.0]
|
||||
pred_processed_col_c = [0, 2, None]
|
||||
pred_expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": pred_processed_col_a,
|
||||
"B": pred_processed_col_b,
|
||||
"C": pred_processed_col_c,
|
||||
}
|
||||
).astype(pred_out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(pred_out_df, pred_expected_df, check_like=True)
|
||||
|
||||
|
||||
def test_nested_chain_state():
|
||||
col_a = [-1, -1, 1, 1]
|
||||
col_b = [1, 1, 1, None]
|
||||
col_c = ["sunday", "monday", "tuesday", "tuesday"]
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
||||
ds = ray.data.from_pandas(in_df)
|
||||
|
||||
def create_chain():
|
||||
imputer = SimpleImputer(["B"])
|
||||
scaler = StandardScaler(["A", "B"])
|
||||
encoder = LabelEncoder("C")
|
||||
return Chain(Chain(scaler, imputer), encoder)
|
||||
|
||||
chain = create_chain()
|
||||
assert chain.fit_status() == Preprocessor.FitStatus.NOT_FITTED
|
||||
|
||||
chain = create_chain()
|
||||
chain.preprocessors[1].fit(ds)
|
||||
assert chain.fit_status() == Preprocessor.FitStatus.PARTIALLY_FITTED
|
||||
|
||||
chain = create_chain()
|
||||
chain.preprocessors[0].fit(ds)
|
||||
assert chain.fit_status() == Preprocessor.FitStatus.PARTIALLY_FITTED
|
||||
|
||||
chain.preprocessors[1].fit(ds)
|
||||
assert chain.fit_status() == Preprocessor.FitStatus.FITTED
|
||||
|
||||
chain = create_chain()
|
||||
chain.fit(ds)
|
||||
assert chain.fit_status() == Preprocessor.FitStatus.FITTED
|
||||
|
||||
|
||||
def test_nested_chain():
|
||||
"""Tests Chain-inside-Chain functionality."""
|
||||
col_a = [-1, -1, 1, 1]
|
||||
col_b = [1, 1, 1, None]
|
||||
col_c = ["sunday", "monday", "tuesday", "tuesday"]
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
||||
ds = ray.data.from_pandas(in_df)
|
||||
|
||||
imputer = SimpleImputer(["B"])
|
||||
scaler = StandardScaler(["A", "B"])
|
||||
encoder = LabelEncoder("C")
|
||||
chain = Chain(Chain(scaler, imputer), encoder)
|
||||
|
||||
# Fit data.
|
||||
chain.fit(ds)
|
||||
# Transform data.
|
||||
transformed = chain.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
assert imputer.stats_ == {
|
||||
"mean(B)": 0.0,
|
||||
}
|
||||
assert scaler.stats_ == {
|
||||
"mean(A)": 0.0,
|
||||
"mean(B)": 1.0,
|
||||
"std(A)": 1.0,
|
||||
"std(B)": 0.0,
|
||||
}
|
||||
assert encoder.stats_ == {
|
||||
"unique_values(C)": {"monday": 0, "sunday": 1, "tuesday": 2}
|
||||
}
|
||||
|
||||
processed_col_a = [-1.0, -1.0, 1.0, 1.0]
|
||||
processed_col_b = [0.0, 0.0, 0.0, 0.0]
|
||||
processed_col_c = [1, 0, 2, 2]
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{"A": processed_col_a, "B": processed_col_b, "C": processed_col_c}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
# Transform batch.
|
||||
pred_col_a = [1, 2, None]
|
||||
pred_col_b = [0, None, 2]
|
||||
pred_col_c = ["monday", "tuesday", "wednesday"]
|
||||
pred_in_df = pd.DataFrame.from_dict(
|
||||
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c}
|
||||
)
|
||||
|
||||
pred_out_df = chain.transform_batch(pred_in_df)
|
||||
|
||||
pred_processed_col_a = [1, 2, None]
|
||||
pred_processed_col_b = [-1.0, 0.0, 1.0]
|
||||
pred_processed_col_c = [0, 2, None]
|
||||
pred_expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": pred_processed_col_a,
|
||||
"B": pred_processed_col_b,
|
||||
"C": pred_processed_col_c,
|
||||
}
|
||||
).astype(pred_out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(pred_out_df, pred_expected_df, check_like=True)
|
||||
|
||||
|
||||
class PreprocessorWithoutTransform(Preprocessor):
|
||||
pass
|
||||
|
||||
|
||||
def test_determine_transform_to_use():
|
||||
# Test that _determine_transform_to_use doesn't throw any exceptions
|
||||
# and selects the transform function of the underlying preprocessor
|
||||
# while dealing with the nested Chain case.
|
||||
|
||||
# Check that error is propagated correctly
|
||||
with pytest.raises(NotImplementedError):
|
||||
chain = Chain(PreprocessorWithoutTransform())
|
||||
chain._determine_transform_to_use()
|
||||
|
||||
# Should have no errors from here on
|
||||
preprocessor = SimpleImputer(["A"])
|
||||
chain1 = Chain(preprocessor)
|
||||
format1 = chain1._determine_transform_to_use()
|
||||
assert format1 == BatchFormat.PANDAS
|
||||
|
||||
chain2 = Chain(chain1)
|
||||
format2 = chain2._determine_transform_to_use()
|
||||
|
||||
assert format1 == format2
|
||||
|
||||
|
||||
def test_chain_serialization():
|
||||
"""Test Chain serialization and deserialization functionality."""
|
||||
import ray
|
||||
from ray.data.preprocessor import SerializablePreprocessorBase
|
||||
from ray.data.preprocessors import Normalizer, StandardScaler
|
||||
|
||||
# Create and fit chain
|
||||
scaler = StandardScaler(columns=["A"])
|
||||
normalizer = Normalizer(columns=["A"])
|
||||
chain = Chain(scaler, normalizer)
|
||||
|
||||
df = pd.DataFrame({"A": [1.0, 2.0, 3.0]})
|
||||
ds = ray.data.from_pandas(df)
|
||||
fitted_chain = chain.fit(ds)
|
||||
|
||||
# Serialize using CloudPickle
|
||||
serialized = fitted_chain.serialize()
|
||||
|
||||
# Verify it's binary CloudPickle format
|
||||
assert isinstance(serialized, bytes)
|
||||
assert serialized.startswith(SerializablePreprocessorBase.MAGIC_CLOUDPICKLE)
|
||||
|
||||
# Deserialize
|
||||
deserialized = Chain.deserialize(serialized)
|
||||
|
||||
# Verify type and field values
|
||||
assert isinstance(deserialized, Chain)
|
||||
assert len(deserialized._preprocessors) == 2
|
||||
assert isinstance(deserialized._preprocessors[0], StandardScaler)
|
||||
assert isinstance(deserialized._preprocessors[1], Normalizer)
|
||||
# Verify the StandardScaler is fitted (Normalizer is stateless)
|
||||
assert deserialized._preprocessors[0]._fitted
|
||||
|
||||
# Verify it works correctly
|
||||
test_df = pd.DataFrame({"A": [1.5, 2.5]})
|
||||
result = deserialized.transform_batch(test_df)
|
||||
|
||||
# Result should have been transformed by both preprocessors
|
||||
assert "A" in result.columns
|
||||
assert len(result) == 2
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,227 @@
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
from pandas.testing import assert_frame_equal
|
||||
|
||||
import ray
|
||||
from ray.data.exceptions import UserCodeException
|
||||
from ray.data.preprocessors import Concatenator, OneHotEncoder
|
||||
|
||||
|
||||
class TestConcatenator:
|
||||
def test_basic(self):
|
||||
df = pd.DataFrame(
|
||||
{
|
||||
"a": [1, 2, 3, 4],
|
||||
"b": [5, 6, 7, 8],
|
||||
}
|
||||
)
|
||||
ds = ray.data.from_pandas(df)
|
||||
prep = Concatenator(columns=["a", "b"], output_column_name="c")
|
||||
new_ds = prep.transform(ds)
|
||||
for i, row in enumerate(new_ds.take()):
|
||||
assert np.array_equal(row["c"], np.array([i + 1, i + 5]))
|
||||
|
||||
def test_raise_if_missing(self):
|
||||
df = pd.DataFrame({"a": [1, 2, 3, 4]})
|
||||
ds = ray.data.from_pandas(df)
|
||||
prep = Concatenator(
|
||||
columns=["a", "b"], output_column_name="c", raise_if_missing=True
|
||||
)
|
||||
|
||||
with pytest.raises(UserCodeException):
|
||||
with pytest.raises(ValueError, match="'b'"):
|
||||
prep.transform(ds).materialize()
|
||||
|
||||
def test_exclude_column(self):
|
||||
df = pd.DataFrame({"a": [1, 2, 3, 4], "b": [2, 3, 4, 5], "c": [3, 4, 5, 6]})
|
||||
ds = ray.data.from_pandas(df)
|
||||
prep = Concatenator(columns=["a", "c"])
|
||||
new_ds = prep.transform(ds)
|
||||
for _, row in enumerate(new_ds.take()):
|
||||
assert set(row) == {"concat_out", "b"}
|
||||
|
||||
def test_include_columns(self):
|
||||
df = pd.DataFrame({"a": [1, 2, 3, 4], "b": [2, 3, 4, 5], "c": [3, 4, 5, 6]})
|
||||
ds = ray.data.from_pandas(df)
|
||||
prep = Concatenator(columns=["a", "b"])
|
||||
new_ds = prep.transform(ds)
|
||||
for _, row in enumerate(new_ds.take()):
|
||||
assert set(row) == {"concat_out", "c"}
|
||||
|
||||
def test_change_column_order(self):
|
||||
df = pd.DataFrame({"a": [1, 2, 3, 4], "b": [2, 3, 4, 5]})
|
||||
ds = ray.data.from_pandas(df)
|
||||
prep = Concatenator(columns=["b", "a"])
|
||||
new_ds = prep.transform(ds)
|
||||
expected_df = pd.DataFrame({"concat_out": [[2, 1], [3, 2], [4, 3], [5, 4]]})
|
||||
print(new_ds.to_pandas())
|
||||
assert_frame_equal(new_ds.to_pandas(), expected_df)
|
||||
|
||||
def test_strings(self):
|
||||
df = pd.DataFrame({"a": ["string", "string2", "string3"]})
|
||||
ds = ray.data.from_pandas(df)
|
||||
prep = Concatenator(columns=["a"], output_column_name="huh")
|
||||
new_ds = prep.transform(ds)
|
||||
assert "huh" in set(new_ds.schema().names)
|
||||
|
||||
def test_preserves_order(self):
|
||||
df = pd.DataFrame({"a": [1, 2, 3, 4], "b": [2, 3, 4, 5]})
|
||||
ds = ray.data.from_pandas(df)
|
||||
prep = Concatenator(columns=["a", "b"], output_column_name="c")
|
||||
prep = prep.fit(ds)
|
||||
|
||||
df = pd.DataFrame({"a": [5, 6, 7, 8], "b": [6, 7, 8, 9]})
|
||||
concatenated_df = prep.transform_batch(df)
|
||||
expected_df = pd.DataFrame({"c": [[5, 6], [6, 7], [7, 8], [8, 9]]})
|
||||
assert_frame_equal(concatenated_df, expected_df)
|
||||
|
||||
other_df = pd.DataFrame({"a": [9, 10, 11, 12], "b": [10, 11, 12, 13]})
|
||||
concatenated_other_df = prep.transform_batch(other_df)
|
||||
expected_df = pd.DataFrame(
|
||||
{
|
||||
"c": [
|
||||
[9, 10],
|
||||
[10, 11],
|
||||
[11, 12],
|
||||
[12, 13],
|
||||
]
|
||||
}
|
||||
)
|
||||
assert_frame_equal(concatenated_other_df, expected_df)
|
||||
|
||||
@pytest.mark.parametrize("col_b", [[[2, 3], [3, 4], [4, 5], [5, 6]], [2, 3, 4, 5]])
|
||||
@pytest.mark.parametrize("flatten", [True, False])
|
||||
def test_flatten(self, col_b, flatten):
|
||||
col_a = [1, 2, 3, 4]
|
||||
col_b = [np.array(v) for v in col_b] if isinstance(col_b[0], list) else col_b
|
||||
df = pd.DataFrame({"a": col_a, "b": col_b})
|
||||
ds = ray.data.from_pandas(df)
|
||||
prep = Concatenator(columns=["a", "b"], flatten=flatten)
|
||||
new_ds = prep.transform(ds)
|
||||
|
||||
for i, row in enumerate(new_ds.take()):
|
||||
if flatten or not isinstance(col_b[i], np.ndarray):
|
||||
# When flatten=True or when col_b contains simple values
|
||||
if isinstance(col_b[i], np.ndarray):
|
||||
expected = np.concatenate([np.array([col_a[i]]), col_b[i]])
|
||||
else:
|
||||
expected = np.array([col_a[i], col_b[i]])
|
||||
assert np.array_equal(row["concat_out"], expected)
|
||||
else:
|
||||
# When flatten=False and col_b contains numpy arrays
|
||||
# The output should be a list containing the scalar and the array
|
||||
assert len(row["concat_out"]) == 2
|
||||
assert row["concat_out"][0] == col_a[i]
|
||||
assert np.array_equal(row["concat_out"][1], col_b[i])
|
||||
|
||||
@pytest.mark.parametrize("flatten", [True, False])
|
||||
def test_concatenate_with_onehotencoder(self, flatten):
|
||||
df = pd.DataFrame(
|
||||
{
|
||||
"color": ["red", "green", "blue", "red"],
|
||||
"value": [1, 2, 3, 4],
|
||||
}
|
||||
)
|
||||
ds = ray.data.from_pandas(df)
|
||||
# OneHot encode the color column
|
||||
encoder = OneHotEncoder(columns=["color"], output_columns=["color_encoded"])
|
||||
encoder = encoder.fit(ds)
|
||||
encoded_ds = encoder.transform(ds)
|
||||
# Concatenate the one-hot encoded column with the value column
|
||||
prep = Concatenator(
|
||||
columns=["color_encoded", "value"],
|
||||
output_column_name="features",
|
||||
flatten=flatten,
|
||||
)
|
||||
new_ds = prep.transform(encoded_ds)
|
||||
# Get the expected one-hot vectors
|
||||
color_map = {"blue": [1, 0, 0], "green": [0, 1, 0], "red": [0, 0, 1]}
|
||||
for i, row in enumerate(new_ds.take()):
|
||||
if flatten:
|
||||
expected = color_map[df["color"][i]] + [df["value"][i]]
|
||||
assert np.array_equal(row["features"], np.array(expected))
|
||||
else:
|
||||
expected = [np.array(color_map[df["color"][i]]), df["value"][i]]
|
||||
assert np.array_equal(row["features"][0], expected[0])
|
||||
assert row["features"][1] == expected[1]
|
||||
|
||||
@pytest.mark.parametrize("flatten", [True, False])
|
||||
def test_nested_list_with_dtype(self, flatten: bool):
|
||||
# Tests Concatenator with nested lists and dtype: flattens and coerces when flatten=True,
|
||||
# raises ValueError when flatten=False.
|
||||
output_column = "c"
|
||||
df = pd.DataFrame(
|
||||
{
|
||||
"a": [12.0],
|
||||
"b": [[1, 0, 0, 0]],
|
||||
}
|
||||
)
|
||||
prep = Concatenator(
|
||||
columns=["a", "b"],
|
||||
output_column_name=output_column,
|
||||
dtype=np.float32,
|
||||
flatten=flatten,
|
||||
)
|
||||
|
||||
if flatten:
|
||||
pd_ds = prep._transform_pandas(df)
|
||||
expected_pd = pd.DataFrame(
|
||||
{output_column: pd.Series([[12.0, 1.0, 0.0, 0.0, 0.0]])}
|
||||
)
|
||||
assert_frame_equal(pd_ds, expected_pd)
|
||||
else:
|
||||
# Only for flattened output do we expect the dtype coercion to apply
|
||||
with pytest.raises(ValueError):
|
||||
pd_ds = prep._transform_pandas(df)
|
||||
|
||||
def test_concatenator_deserialize_backward_compat(self):
|
||||
p1 = Concatenator(columns=["A"], flatten=True)
|
||||
delattr(p1, "_flatten")
|
||||
data = p1.serialize()
|
||||
p2 = Concatenator.deserialize(data)
|
||||
assert isinstance(p2, Concatenator)
|
||||
assert p2.flatten is False
|
||||
|
||||
def test_concatenator_serialization(self):
|
||||
"""Test Concatenator serialization and deserialization functionality."""
|
||||
from ray.data.preprocessor import SerializablePreprocessorBase
|
||||
|
||||
# Create concatenator
|
||||
concatenator = Concatenator(
|
||||
columns=["A", "B"],
|
||||
output_column_name="combined",
|
||||
dtype=np.float32,
|
||||
flatten=True,
|
||||
)
|
||||
|
||||
# Serialize using CloudPickle
|
||||
serialized = concatenator.serialize()
|
||||
|
||||
# Verify it's binary CloudPickle format
|
||||
assert isinstance(serialized, bytes)
|
||||
assert serialized.startswith(SerializablePreprocessorBase.MAGIC_CLOUDPICKLE)
|
||||
|
||||
# Deserialize
|
||||
deserialized = Concatenator.deserialize(serialized)
|
||||
|
||||
# Verify type and field values
|
||||
assert isinstance(deserialized, Concatenator)
|
||||
assert deserialized.columns == ["A", "B"]
|
||||
assert deserialized.output_column_name == "combined"
|
||||
assert deserialized.dtype == np.float32
|
||||
assert deserialized.flatten is True
|
||||
|
||||
# Verify it works correctly
|
||||
df = pd.DataFrame({"A": [[1, 2]], "B": [[3, 4]]})
|
||||
result = deserialized.transform_batch(df)
|
||||
|
||||
# Verify concatenation was applied correctly
|
||||
assert "combined" in result.columns
|
||||
assert len(result["combined"][0]) == 4
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,312 @@
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data._internal.util import rows_same
|
||||
from ray.data.preprocessors import CustomKBinsDiscretizer, UniformKBinsDiscretizer
|
||||
|
||||
|
||||
@pytest.mark.parametrize("bins", (3, {"A": 4, "B": 3}))
|
||||
@pytest.mark.parametrize(
|
||||
"dtypes",
|
||||
(
|
||||
None,
|
||||
{"A": int, "B": int},
|
||||
{"A": int, "B": pd.CategoricalDtype(["cat1", "cat2", "cat3"], ordered=True)},
|
||||
),
|
||||
)
|
||||
@pytest.mark.parametrize("right", (True, False))
|
||||
@pytest.mark.parametrize("include_lowest", (True, False))
|
||||
def test_uniform_kbins_discretizer(
|
||||
bins,
|
||||
dtypes,
|
||||
right,
|
||||
include_lowest,
|
||||
):
|
||||
"""Tests basic UniformKBinsDiscretizer functionality."""
|
||||
|
||||
col_a = [0.2, 1.4, 2.5, 6.2, 9.7, 2.1]
|
||||
col_b = col_a.copy()
|
||||
col_c = col_a.copy()
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
||||
ds = ray.data.from_pandas(in_df).repartition(2)
|
||||
|
||||
discretizer = UniformKBinsDiscretizer(
|
||||
["A", "B"], bins=bins, dtypes=dtypes, right=right, include_lowest=include_lowest
|
||||
)
|
||||
|
||||
transformed = discretizer.fit_transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
if isinstance(bins, dict):
|
||||
bins_A = bins["A"]
|
||||
bins_B = bins["B"]
|
||||
else:
|
||||
bins_A = bins_B = bins
|
||||
|
||||
labels_A = False
|
||||
ordered_A = True
|
||||
labels_B = False
|
||||
ordered_B = True
|
||||
if isinstance(dtypes, dict):
|
||||
if isinstance(dtypes.get("A"), pd.CategoricalDtype):
|
||||
labels_A = dtypes.get("A").categories
|
||||
ordered_A = dtypes.get("A").ordered
|
||||
if isinstance(dtypes.get("B"), pd.CategoricalDtype):
|
||||
labels_B = dtypes.get("B").categories
|
||||
ordered_B = dtypes.get("B").ordered
|
||||
|
||||
# Create expected dataframe with transformed columns
|
||||
expected_df = in_df.copy()
|
||||
expected_df["A"] = pd.cut(
|
||||
in_df["A"],
|
||||
bins_A,
|
||||
labels=labels_A,
|
||||
ordered=ordered_A,
|
||||
right=right,
|
||||
include_lowest=include_lowest,
|
||||
)
|
||||
expected_df["B"] = pd.cut(
|
||||
in_df["B"],
|
||||
bins_B,
|
||||
labels=labels_B,
|
||||
ordered=ordered_B,
|
||||
right=right,
|
||||
include_lowest=include_lowest,
|
||||
)
|
||||
|
||||
# Use rows_same to compare regardless of row ordering
|
||||
assert rows_same(out_df, expected_df)
|
||||
|
||||
# append mode
|
||||
expected_message = "The length of columns and output_columns must match."
|
||||
with pytest.raises(ValueError, match=expected_message):
|
||||
UniformKBinsDiscretizer(["A", "B"], bins=bins, output_columns=["A_discretized"])
|
||||
|
||||
discretizer = UniformKBinsDiscretizer(
|
||||
["A", "B"],
|
||||
bins=bins,
|
||||
dtypes=dtypes,
|
||||
right=right,
|
||||
include_lowest=include_lowest,
|
||||
output_columns=["A_discretized", "B_discretized"],
|
||||
)
|
||||
|
||||
transformed = discretizer.fit_transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
# Create expected dataframe with appended columns
|
||||
expected_df = in_df.copy()
|
||||
expected_df["A_discretized"] = pd.cut(
|
||||
in_df["A"],
|
||||
bins_A,
|
||||
labels=labels_A,
|
||||
ordered=ordered_A,
|
||||
right=right,
|
||||
include_lowest=include_lowest,
|
||||
)
|
||||
expected_df["B_discretized"] = pd.cut(
|
||||
in_df["B"],
|
||||
bins_B,
|
||||
labels=labels_B,
|
||||
ordered=ordered_B,
|
||||
right=right,
|
||||
include_lowest=include_lowest,
|
||||
)
|
||||
|
||||
# Use rows_same to compare regardless of row ordering
|
||||
assert rows_same(out_df, expected_df)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"bins", ([3, 4, 6, 9], {"A": [3, 4, 6, 8, 9], "B": [3, 4, 6, 9]})
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"dtypes",
|
||||
(
|
||||
None,
|
||||
{"A": int, "B": int},
|
||||
{"A": int, "B": pd.CategoricalDtype(["cat1", "cat2", "cat3"], ordered=True)},
|
||||
),
|
||||
)
|
||||
@pytest.mark.parametrize("right", (True, False))
|
||||
@pytest.mark.parametrize("include_lowest", (True, False))
|
||||
def test_custom_kbins_discretizer(
|
||||
bins,
|
||||
dtypes,
|
||||
right,
|
||||
include_lowest,
|
||||
):
|
||||
"""Tests basic CustomKBinsDiscretizer functionality."""
|
||||
|
||||
col_a = [0.2, 1.4, 2.5, 6.2, 9.7, 2.1]
|
||||
col_b = col_a.copy()
|
||||
col_c = col_a.copy()
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
||||
ds = ray.data.from_pandas(in_df).repartition(2)
|
||||
|
||||
discretizer = CustomKBinsDiscretizer(
|
||||
["A", "B"], bins=bins, dtypes=dtypes, right=right, include_lowest=include_lowest
|
||||
)
|
||||
|
||||
transformed = discretizer.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
if isinstance(bins, dict):
|
||||
bins_A = bins["A"]
|
||||
bins_B = bins["B"]
|
||||
else:
|
||||
bins_A = bins_B = bins
|
||||
|
||||
labels_A = False
|
||||
ordered_A = True
|
||||
labels_B = False
|
||||
ordered_B = True
|
||||
if isinstance(dtypes, dict):
|
||||
if isinstance(dtypes.get("A"), pd.CategoricalDtype):
|
||||
labels_A = dtypes.get("A").categories
|
||||
ordered_A = dtypes.get("A").ordered
|
||||
if isinstance(dtypes.get("B"), pd.CategoricalDtype):
|
||||
labels_B = dtypes.get("B").categories
|
||||
ordered_B = dtypes.get("B").ordered
|
||||
|
||||
# Create expected dataframe with transformed columns
|
||||
expected_df = in_df.copy()
|
||||
expected_df["A"] = pd.cut(
|
||||
in_df["A"],
|
||||
bins_A,
|
||||
labels=labels_A,
|
||||
ordered=ordered_A,
|
||||
right=right,
|
||||
include_lowest=include_lowest,
|
||||
)
|
||||
expected_df["B"] = pd.cut(
|
||||
in_df["B"],
|
||||
bins_B,
|
||||
labels=labels_B,
|
||||
ordered=ordered_B,
|
||||
right=right,
|
||||
include_lowest=include_lowest,
|
||||
)
|
||||
|
||||
# Use rows_same to compare regardless of row ordering
|
||||
assert rows_same(out_df, expected_df)
|
||||
|
||||
# append mode
|
||||
expected_message = "The length of columns and output_columns must match."
|
||||
with pytest.raises(ValueError, match=expected_message):
|
||||
CustomKBinsDiscretizer(["A", "B"], bins=bins, output_columns=["A_discretized"])
|
||||
|
||||
discretizer = CustomKBinsDiscretizer(
|
||||
["A", "B"],
|
||||
bins=bins,
|
||||
dtypes=dtypes,
|
||||
right=right,
|
||||
include_lowest=include_lowest,
|
||||
output_columns=["A_discretized", "B_discretized"],
|
||||
)
|
||||
|
||||
transformed = discretizer.fit_transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
# Create expected dataframe with appended columns
|
||||
expected_df = in_df.copy()
|
||||
expected_df["A_discretized"] = pd.cut(
|
||||
in_df["A"],
|
||||
bins_A,
|
||||
labels=labels_A,
|
||||
ordered=ordered_A,
|
||||
right=right,
|
||||
include_lowest=include_lowest,
|
||||
)
|
||||
expected_df["B_discretized"] = pd.cut(
|
||||
in_df["B"],
|
||||
bins_B,
|
||||
labels=labels_B,
|
||||
ordered=ordered_B,
|
||||
right=right,
|
||||
include_lowest=include_lowest,
|
||||
)
|
||||
|
||||
# Use rows_same to compare regardless of row ordering
|
||||
assert rows_same(out_df, expected_df)
|
||||
|
||||
|
||||
def test_custom_kbins_discretizer_serialization():
|
||||
"""Test CustomKBinsDiscretizer serialization and deserialization functionality."""
|
||||
from ray.data.preprocessor import SerializablePreprocessorBase
|
||||
|
||||
# Create discretizer
|
||||
discretizer = CustomKBinsDiscretizer(
|
||||
columns=["A"], bins={"A": [0, 1, 2, 3]}, right=True
|
||||
)
|
||||
|
||||
# Serialize using CloudPickle
|
||||
serialized = discretizer.serialize()
|
||||
|
||||
# Verify it's binary CloudPickle format
|
||||
assert isinstance(serialized, bytes)
|
||||
assert serialized.startswith(SerializablePreprocessorBase.MAGIC_CLOUDPICKLE)
|
||||
|
||||
# Deserialize
|
||||
deserialized = CustomKBinsDiscretizer.deserialize(serialized)
|
||||
|
||||
# Verify type and field values
|
||||
assert isinstance(deserialized, CustomKBinsDiscretizer)
|
||||
assert deserialized.columns == ["A"]
|
||||
assert deserialized.bins == {"A": [0, 1, 2, 3]}
|
||||
assert deserialized.right is True
|
||||
|
||||
# Verify it works correctly
|
||||
df = pd.DataFrame({"A": [0.5, 1.5, 2.5]})
|
||||
result = deserialized.transform_batch(df)
|
||||
|
||||
# Verify discretization was applied correctly
|
||||
assert "A" in result.columns
|
||||
assert len(result) == 3
|
||||
|
||||
|
||||
def test_uniform_kbins_discretizer_serialization():
|
||||
"""Test UniformKBinsDiscretizer serialization and deserialization functionality."""
|
||||
import ray
|
||||
from ray.data.preprocessor import SerializablePreprocessorBase
|
||||
|
||||
# Create and fit discretizer
|
||||
discretizer = UniformKBinsDiscretizer(columns=["A"], bins=3)
|
||||
df = pd.DataFrame({"A": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]})
|
||||
ds = ray.data.from_pandas(df)
|
||||
fitted_discretizer = discretizer.fit(ds)
|
||||
|
||||
# Serialize using CloudPickle
|
||||
serialized = fitted_discretizer.serialize()
|
||||
|
||||
# Verify it's binary CloudPickle format
|
||||
assert isinstance(serialized, bytes)
|
||||
assert serialized.startswith(SerializablePreprocessorBase.MAGIC_CLOUDPICKLE)
|
||||
|
||||
# Deserialize
|
||||
deserialized = UniformKBinsDiscretizer.deserialize(serialized)
|
||||
|
||||
# Verify type and field values
|
||||
assert isinstance(deserialized, UniformKBinsDiscretizer)
|
||||
assert deserialized._fitted
|
||||
assert deserialized.columns == ["A"]
|
||||
assert deserialized.bins == 3
|
||||
|
||||
# Verify stats are preserved: bin edges for 3 bins = 4 edge values
|
||||
assert "A" in deserialized.stats_
|
||||
assert len(deserialized.stats_["A"]) == 4
|
||||
|
||||
# Verify it works correctly
|
||||
test_df = pd.DataFrame({"A": [1.5, 3.5, 5.5]})
|
||||
result = deserialized.transform_batch(test_df)
|
||||
|
||||
# Verify discretization was applied correctly
|
||||
assert "A" in result.columns
|
||||
assert len(result) == 3
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,80 @@
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data.preprocessors import FeatureHasher
|
||||
|
||||
|
||||
def test_feature_hasher():
|
||||
"""Tests basic FeatureHasher functionality."""
|
||||
# This dataframe represents the counts from the documents "I like Python" and "I
|
||||
# dislike Python".
|
||||
token_counts = pd.DataFrame(
|
||||
{"I": [1, 1], "like": [1, 0], "dislike": [0, 1], "Python": [1, 1]}
|
||||
)
|
||||
|
||||
hasher = FeatureHasher(
|
||||
["I", "like", "dislike", "Python"],
|
||||
num_features=256,
|
||||
output_column="hashed_features",
|
||||
)
|
||||
document_term_matrix = hasher.fit_transform(
|
||||
ray.data.from_pandas(token_counts)
|
||||
).to_pandas()
|
||||
|
||||
hashed_features = document_term_matrix["hashed_features"]
|
||||
# Document-term matrix should have shape (# documents, # features)
|
||||
assert hashed_features.shape == (2,)
|
||||
|
||||
# The tokens tokens "I", "like", and "Python" should be hashed to distinct indices
|
||||
# for adequately large `num_features`.
|
||||
assert len(hashed_features.iloc[0]) == 256
|
||||
assert hashed_features.iloc[0].sum() == 3
|
||||
assert all(hashed_features.iloc[0] <= 1)
|
||||
|
||||
# The tokens tokens "I", "dislike", and "Python" should be hashed to distinct
|
||||
# indices for adequately large `num_features`.
|
||||
assert len(hashed_features.iloc[1]) == 256
|
||||
assert hashed_features.iloc[1].sum() == 3
|
||||
assert all(hashed_features.iloc[1] <= 1)
|
||||
|
||||
|
||||
def test_feature_hasher_serialization():
|
||||
"""Test FeatureHasher serialization and deserialization functionality."""
|
||||
from ray.data.preprocessor import SerializablePreprocessorBase
|
||||
|
||||
# Create hasher
|
||||
hasher = FeatureHasher(
|
||||
columns=["I", "like", "Python"], num_features=8, output_column="hashed"
|
||||
)
|
||||
|
||||
# Serialize using CloudPickle
|
||||
serialized = hasher.serialize()
|
||||
|
||||
# Verify it's binary CloudPickle format
|
||||
assert isinstance(serialized, bytes)
|
||||
assert serialized.startswith(SerializablePreprocessorBase.MAGIC_CLOUDPICKLE)
|
||||
|
||||
# Deserialize
|
||||
deserialized = FeatureHasher.deserialize(serialized)
|
||||
|
||||
# Verify type and field values
|
||||
assert isinstance(deserialized, FeatureHasher)
|
||||
assert deserialized.columns == ["I", "like", "Python"]
|
||||
assert deserialized.num_features == 8
|
||||
assert deserialized.output_column == "hashed"
|
||||
|
||||
# Verify it works correctly
|
||||
df = pd.DataFrame({"I": [1, 1], "like": [1, 0], "Python": [1, 1]})
|
||||
result = deserialized.transform_batch(df)
|
||||
|
||||
# Verify hashing was applied correctly
|
||||
assert "hashed" in result.columns
|
||||
assert len(result["hashed"][0]) == 8
|
||||
assert len(result["hashed"][1]) == 8
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,629 @@
|
||||
"""
|
||||
Tests for SimpleImputer functionality and serialization.
|
||||
|
||||
This file contains:
|
||||
1. Basic functional tests for SimpleImputer operations
|
||||
2. Comprehensive serialization/deserialization tests
|
||||
"""
|
||||
import tempfile
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data._internal.util import rows_same
|
||||
from ray.data.preprocessor import (
|
||||
PreprocessorNotFittedException,
|
||||
SerializablePreprocessorBase,
|
||||
)
|
||||
from ray.data.preprocessors import SimpleImputer
|
||||
from ray.data.preprocessors.version_support import UnknownPreprocessorError
|
||||
|
||||
|
||||
def test_simple_imputer():
|
||||
col_a = [1, 1, 1, np.nan]
|
||||
col_b = [1, 3, None, np.nan]
|
||||
col_c = [1, 1, 1, 1]
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
||||
|
||||
ds = ray.data.from_pandas(in_df)
|
||||
|
||||
imputer = SimpleImputer(["B", "C"])
|
||||
|
||||
# Transform with unfitted preprocessor.
|
||||
with pytest.raises(PreprocessorNotFittedException):
|
||||
imputer.transform(ds)
|
||||
|
||||
# Fit data.
|
||||
imputer.fit(ds)
|
||||
assert imputer.stats_ == {"mean(B)": 2.0, "mean(C)": 1.0}
|
||||
|
||||
# Transform data.
|
||||
transformed = imputer.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
processed_col_a = col_a
|
||||
processed_col_b = [1.0, 3.0, 2.0, 2.0]
|
||||
processed_col_c = [1, 1, 1, 1]
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{"A": processed_col_a, "B": processed_col_b, "C": processed_col_c}
|
||||
)
|
||||
expected_df = expected_df.astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df)
|
||||
|
||||
# Transform batch.
|
||||
pred_col_a = [1, 2, np.nan]
|
||||
pred_col_b = [1, 2, np.nan]
|
||||
pred_col_c = [None, None, None]
|
||||
pred_in_df = pd.DataFrame.from_dict(
|
||||
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c}
|
||||
)
|
||||
|
||||
pred_out_df = imputer.transform_batch(pred_in_df)
|
||||
|
||||
pred_processed_col_a = pred_col_a
|
||||
pred_processed_col_b = [1.0, 2.0, 2.0]
|
||||
pred_processed_col_c = [1.0, 1.0, 1.0]
|
||||
pred_expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": pred_processed_col_a,
|
||||
"B": pred_processed_col_b,
|
||||
"C": pred_processed_col_c,
|
||||
}
|
||||
)
|
||||
|
||||
pd.testing.assert_frame_equal(pred_out_df, pred_expected_df, check_like=True)
|
||||
|
||||
# with missing column
|
||||
pred_in_df = pd.DataFrame.from_dict({"A": pred_col_a, "B": pred_col_b})
|
||||
pred_out_df = imputer.transform_batch(pred_in_df)
|
||||
pred_expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": pred_processed_col_a,
|
||||
"B": pred_processed_col_b,
|
||||
"C": pred_processed_col_c,
|
||||
}
|
||||
)
|
||||
pd.testing.assert_frame_equal(pred_out_df, pred_expected_df, check_like=True)
|
||||
|
||||
# append mode
|
||||
with pytest.raises(ValueError):
|
||||
SimpleImputer(columns=["B", "C"], output_columns=["B_encoded"])
|
||||
|
||||
imputer = SimpleImputer(
|
||||
columns=["B", "C"],
|
||||
output_columns=["B_imputed", "C_imputed"],
|
||||
)
|
||||
imputer.fit(ds)
|
||||
|
||||
pred_col_a = [1, 2, np.nan]
|
||||
pred_col_b = [1, 2, np.nan]
|
||||
pred_col_c = [None, None, None]
|
||||
pred_in_df = pd.DataFrame.from_dict(
|
||||
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c}
|
||||
)
|
||||
pred_out_df = imputer.transform_batch(pred_in_df)
|
||||
|
||||
pred_processed_col_b = [1.0, 2.0, 2.0]
|
||||
pred_processed_col_c = [1.0, 1.0, 1.0]
|
||||
pred_expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": pred_col_a,
|
||||
"B": pred_col_b,
|
||||
"C": pred_col_c,
|
||||
"B_imputed": pred_processed_col_b,
|
||||
"C_imputed": pred_processed_col_c,
|
||||
}
|
||||
)
|
||||
|
||||
pd.testing.assert_frame_equal(pred_out_df, pred_expected_df, check_like=True)
|
||||
|
||||
# Test "most_frequent" strategy.
|
||||
most_frequent_col_a = [1, 2, 2, None, None, None]
|
||||
# Use 3 "c"s to ensure it's clearly the most frequent (no tie with "b")
|
||||
most_frequent_col_b = [None, "c", "c", "c", "b", "a"]
|
||||
most_frequent_df = pd.DataFrame.from_dict(
|
||||
{"A": most_frequent_col_a, "B": most_frequent_col_b}
|
||||
)
|
||||
most_frequent_ds = ray.data.from_pandas(most_frequent_df).repartition(3)
|
||||
|
||||
most_frequent_imputer = SimpleImputer(["A", "B"], strategy="most_frequent")
|
||||
most_frequent_imputer.fit(most_frequent_ds)
|
||||
assert most_frequent_imputer.stats_ == {
|
||||
"most_frequent(A)": 2.0,
|
||||
"most_frequent(B)": "c",
|
||||
}
|
||||
|
||||
most_frequent_transformed = most_frequent_imputer.transform(most_frequent_ds)
|
||||
most_frequent_out_df = most_frequent_transformed.to_pandas()
|
||||
|
||||
most_frequent_processed_col_a = [1.0, 2.0, 2.0, 2.0, 2.0, 2.0]
|
||||
most_frequent_processed_col_b = ["c", "c", "c", "c", "b", "a"]
|
||||
most_frequent_expected_df = pd.DataFrame.from_dict(
|
||||
{"A": most_frequent_processed_col_a, "B": most_frequent_processed_col_b}
|
||||
)
|
||||
|
||||
assert rows_same(most_frequent_out_df, most_frequent_expected_df)
|
||||
|
||||
# Test "constant" strategy.
|
||||
constant_col_a = ["apple", None]
|
||||
constant_col_b = constant_col_a.copy()
|
||||
constant_df = pd.DataFrame.from_dict({"A": constant_col_a, "B": constant_col_b})
|
||||
# category dtype requires special handling
|
||||
constant_df["B"] = constant_df["B"].astype("category")
|
||||
constant_ds = ray.data.from_pandas(constant_df)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
SimpleImputer(["A", "B"], strategy="constant")
|
||||
|
||||
constant_imputer = SimpleImputer(
|
||||
["A", "B"], strategy="constant", fill_value="missing"
|
||||
)
|
||||
constant_transformed = constant_imputer.transform(constant_ds)
|
||||
constant_out_df = constant_transformed.to_pandas()
|
||||
|
||||
constant_processed_col_a = ["apple", "missing"]
|
||||
constant_processed_col_b = constant_processed_col_a.copy()
|
||||
constant_expected_df = pd.DataFrame.from_dict(
|
||||
{"A": constant_processed_col_a, "B": constant_processed_col_b}
|
||||
)
|
||||
constant_expected_df["B"] = constant_expected_df["B"].astype("category")
|
||||
constant_expected_df = constant_expected_df.astype(constant_out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(
|
||||
constant_out_df, constant_expected_df, check_like=True
|
||||
)
|
||||
|
||||
|
||||
def test_imputer_all_nan_raise_error():
|
||||
data = {
|
||||
"A": [np.nan, np.nan, np.nan, np.nan],
|
||||
}
|
||||
df = pd.DataFrame(data)
|
||||
dataset = ray.data.from_pandas(df)
|
||||
|
||||
imputer = SimpleImputer(columns=["A"], strategy="mean")
|
||||
imputer.fit(dataset)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
imputer.transform_batch(df)
|
||||
|
||||
|
||||
def test_imputer_constant_categorical():
|
||||
data = {
|
||||
"A_cat": ["one", "two", None, "four"],
|
||||
}
|
||||
df = pd.DataFrame(data)
|
||||
df["A_cat"] = df["A_cat"].astype("category")
|
||||
dataset = ray.data.from_pandas(df)
|
||||
|
||||
imputer = SimpleImputer(columns=["A_cat"], strategy="constant", fill_value="three")
|
||||
imputer.fit(dataset)
|
||||
|
||||
transformed_df = imputer.transform_batch(df)
|
||||
|
||||
expected = {
|
||||
"A_cat": ["one", "two", "three", "four"],
|
||||
}
|
||||
|
||||
for column in data.keys():
|
||||
np.testing.assert_array_equal(transformed_df[column].values, expected[column])
|
||||
|
||||
df = pd.DataFrame({"A": [1, 2, 3, 4]})
|
||||
transformed_df = imputer.transform_batch(df)
|
||||
|
||||
expected = {
|
||||
"A": [1, 2, 3, 4],
|
||||
"A_cat": ["three", "three", "three", "three"],
|
||||
}
|
||||
|
||||
for column in df:
|
||||
np.testing.assert_array_equal(transformed_df[column].values, expected[column])
|
||||
|
||||
|
||||
class TestSimpleImputerSerialization:
|
||||
"""Test CloudPickle-based serialization/deserialization functionality for SimpleImputer."""
|
||||
|
||||
def setup_method(self):
|
||||
"""Set up test data."""
|
||||
self.df_numeric = pd.DataFrame(
|
||||
{
|
||||
"temp": [20.0, 25.0, None, 30.0, None],
|
||||
"humidity": [60.0, None, 70.0, 80.0, 65.0],
|
||||
"other": ["a", "b", "c", "d", "e"], # Non-processed column
|
||||
}
|
||||
)
|
||||
|
||||
def test_basic_serialization(self):
|
||||
"""Test basic serialization and deserialization functionality."""
|
||||
# Create and fit a simple imputer
|
||||
imputer = SimpleImputer(columns=["temp", "humidity"], strategy="mean")
|
||||
|
||||
# Create test data
|
||||
df = pd.DataFrame(
|
||||
{
|
||||
"temp": [1.0, 2.0, None, 4.0],
|
||||
"humidity": [None, 2.0, 3.0, 4.0],
|
||||
"other": [1, 2, 3, 4],
|
||||
}
|
||||
)
|
||||
|
||||
# Fit the imputer
|
||||
dataset = ray.data.from_pandas(df)
|
||||
fitted_imputer = imputer.fit(dataset)
|
||||
|
||||
# Serialize using CloudPickle (primary format)
|
||||
serialized = fitted_imputer.serialize()
|
||||
|
||||
# Verify it's binary CloudPickle format
|
||||
assert isinstance(serialized, bytes)
|
||||
assert serialized.startswith(SerializablePreprocessorBase.MAGIC_CLOUDPICKLE)
|
||||
|
||||
# Deserialize
|
||||
deserialized = SimpleImputer.deserialize(serialized)
|
||||
|
||||
# Verify type and state
|
||||
assert isinstance(deserialized, SimpleImputer)
|
||||
assert deserialized._fitted
|
||||
assert deserialized.columns == ["temp", "humidity"]
|
||||
assert deserialized.strategy == "mean"
|
||||
|
||||
# Verify stats are preserved
|
||||
assert "mean(temp)" in deserialized.stats_
|
||||
assert "mean(humidity)" in deserialized.stats_
|
||||
assert abs(deserialized.stats_["mean(temp)"] - 2.333333) < 0.001
|
||||
assert abs(deserialized.stats_["mean(humidity)"] - 3.0) < 0.001
|
||||
|
||||
def test_serialization_formats(self):
|
||||
"""Test serialization and deserialization."""
|
||||
imputer = SimpleImputer(columns=["temp"], strategy="mean")
|
||||
dataset = ray.data.from_pandas(self.df_numeric)
|
||||
fitted_imputer = imputer.fit(dataset)
|
||||
|
||||
# Test CloudPickle format (default)
|
||||
serialized = fitted_imputer.serialize()
|
||||
assert isinstance(serialized, bytes)
|
||||
assert serialized.startswith(SerializablePreprocessorBase.MAGIC_CLOUDPICKLE)
|
||||
|
||||
# Deserialize and verify it works
|
||||
deserialized = SimpleImputer.deserialize(serialized)
|
||||
|
||||
# Verify it works correctly
|
||||
test_df = pd.DataFrame({"temp": [None, 35.0], "other": [1, 2]})
|
||||
result = deserialized.transform_batch(test_df.copy())
|
||||
|
||||
# Verify the result has the expected structure
|
||||
assert "temp" in result.columns
|
||||
assert "other" in result.columns
|
||||
|
||||
def test_functional_equivalence(self):
|
||||
"""Test that deserialized SimpleImputer works identically to original."""
|
||||
# Create and fit original
|
||||
imputer = SimpleImputer(columns=["value"], strategy="mean")
|
||||
train_df = pd.DataFrame({"value": [10, 20, None, 40], "id": [1, 2, 3, 4]})
|
||||
train_dataset = ray.data.from_pandas(train_df)
|
||||
fitted_imputer = imputer.fit(train_dataset)
|
||||
|
||||
# Test data
|
||||
test_df = pd.DataFrame({"value": [None, 50, None], "id": [5, 6, 7]})
|
||||
|
||||
# Transform with original
|
||||
original_result = fitted_imputer.transform_batch(test_df.copy())
|
||||
|
||||
# Serialize, deserialize, and transform (using CloudPickle)
|
||||
serialized = fitted_imputer.serialize()
|
||||
deserialized = SerializablePreprocessorBase.deserialize(serialized)
|
||||
deserialized_result = deserialized.transform_batch(test_df.copy())
|
||||
|
||||
# Results should be identical
|
||||
pd.testing.assert_frame_equal(original_result, deserialized_result)
|
||||
|
||||
# Verify specific values
|
||||
expected_mean = (10 + 20 + 40) / 3 # 23.333...
|
||||
assert abs(original_result.iloc[0]["value"] - expected_mean) < 1e-10
|
||||
assert abs(deserialized_result.iloc[0]["value"] - expected_mean) < 1e-10
|
||||
|
||||
def test_complex_stats_preservation(self):
|
||||
"""Test that CloudPickle perfectly preserves complex stats with various key types."""
|
||||
imputer = SimpleImputer(columns=["A"], strategy="mean")
|
||||
|
||||
# Manually set complex stats that would be problematic for other formats
|
||||
imputer.stats_ = {
|
||||
# Simple stats
|
||||
"mean(A)": 5.0,
|
||||
"count(A)": 100,
|
||||
# Complex key types that CloudPickle handles natively
|
||||
"unique_values(ints)": {1: 0, 2: 1, 3: 2, 4: 3, 5: 4}, # int keys
|
||||
"unique_values(floats)": {1.1: 0, 2.2: 1, 3.3: 2}, # float keys
|
||||
"unique_values(bools)": {True: 0, False: 1}, # bool keys
|
||||
"unique_values(none)": {None: 0}, # None keys
|
||||
"unique_values(tuples)": {
|
||||
("red", "car"): 0,
|
||||
("blue", "bike"): 1,
|
||||
(1, 2, 3): 2,
|
||||
("nested", ("inner", "tuple")): 3,
|
||||
},
|
||||
"unique_values(sets)": {
|
||||
frozenset([1, 2, 3]): 0,
|
||||
frozenset(["a", "b"]): 1,
|
||||
},
|
||||
"unique_values(mixed)": {
|
||||
"string": 0,
|
||||
42: 1,
|
||||
(1, 2): 2,
|
||||
frozenset([3, 4]): 3,
|
||||
None: 4,
|
||||
True: 5,
|
||||
},
|
||||
}
|
||||
imputer._fitted = True
|
||||
|
||||
# Serialize and deserialize (using CloudPickle)
|
||||
serialized = imputer.serialize()
|
||||
deserialized = SimpleImputer.deserialize(serialized)
|
||||
|
||||
# Verify ALL stats are perfectly preserved
|
||||
assert deserialized.stats_ == imputer.stats_
|
||||
|
||||
# Verify specific complex key preservation
|
||||
for stat_name, stat_dict in imputer.stats_.items():
|
||||
if isinstance(stat_dict, dict):
|
||||
original_keys = set(stat_dict.keys())
|
||||
restored_keys = set(deserialized.stats_[stat_name].keys())
|
||||
|
||||
# Keys should be identical (including types)
|
||||
assert original_keys == restored_keys
|
||||
|
||||
# Values should be identical
|
||||
for key in original_keys:
|
||||
assert stat_dict[key] == deserialized.stats_[stat_name][key]
|
||||
|
||||
# Key types should be preserved
|
||||
for orig_key, rest_key in zip(original_keys, restored_keys):
|
||||
if orig_key == rest_key: # Same key
|
||||
assert type(orig_key) is type(rest_key)
|
||||
|
||||
def test_performance_comparison(self):
|
||||
"""Test CloudPickle performance and simplicity."""
|
||||
# Create a large imputer with many stats
|
||||
imputer = SimpleImputer(
|
||||
columns=[f"col_{i}" for i in range(10)], strategy="mean"
|
||||
)
|
||||
|
||||
# Create large stats dictionary
|
||||
large_stats = {}
|
||||
for i in range(10):
|
||||
large_stats[f"mean(col_{i})"] = float(i)
|
||||
large_stats[f"count(col_{i})"] = 1000 + i
|
||||
|
||||
# Add complex key stats that CloudPickle handles natively
|
||||
large_stats[f"unique_values(col_{i})"] = {
|
||||
(f"key_{j}", j): j for j in range(100) # 100 tuple keys per column
|
||||
}
|
||||
|
||||
imputer.stats_ = large_stats
|
||||
imputer._fitted = True
|
||||
|
||||
# Test serialization performance and correctness (using CloudPickle)
|
||||
start_time = time.time()
|
||||
serialized = imputer.serialize()
|
||||
serialize_time = time.time() - start_time
|
||||
|
||||
start_time = time.time()
|
||||
deserialized = SimpleImputer.deserialize(serialized)
|
||||
deserialize_time = time.time() - start_time
|
||||
|
||||
# Verify correctness
|
||||
assert deserialized.stats_ == imputer.stats_
|
||||
assert len(deserialized.stats_) == len(imputer.stats_)
|
||||
|
||||
# Performance should be reasonable (less than 1 second for this size)
|
||||
assert serialize_time < 1.0
|
||||
assert deserialize_time < 1.0
|
||||
|
||||
# Verify no data loss with complex keys
|
||||
for stat_name in large_stats:
|
||||
if "unique_values" in stat_name:
|
||||
original_keys = set(large_stats[stat_name].keys())
|
||||
restored_keys = set(deserialized.stats_[stat_name].keys())
|
||||
assert original_keys == restored_keys
|
||||
|
||||
def test_cloudpickle_native_support(self):
|
||||
"""Test that CloudPickle handles all Python types natively without transformation."""
|
||||
imputer = SimpleImputer(columns=["A"], strategy="mean")
|
||||
|
||||
# Test all the key types that used to require custom transformation
|
||||
test_keys = [
|
||||
# Basic types
|
||||
"string_key",
|
||||
42, # int
|
||||
3.14, # float
|
||||
True, # bool
|
||||
False, # bool
|
||||
None, # None
|
||||
# Complex types that CloudPickle handles natively
|
||||
(1, 2, 3), # tuple
|
||||
("nested", ("inner", "tuple")), # nested tuple
|
||||
frozenset([1, 2, 3]), # frozenset
|
||||
frozenset(["a", "b"]), # frozenset with strings
|
||||
]
|
||||
|
||||
# Create stats with all these key types
|
||||
imputer.stats_ = {
|
||||
"test_dict": {key: f"value_{i}" for i, key in enumerate(test_keys)}
|
||||
}
|
||||
imputer._fitted = True
|
||||
|
||||
# Serialize and deserialize (using CloudPickle)
|
||||
serialized = imputer.serialize()
|
||||
deserialized = SimpleImputer.deserialize(serialized)
|
||||
|
||||
# Verify perfect preservation
|
||||
original_dict = imputer.stats_["test_dict"]
|
||||
restored_dict = deserialized.stats_["test_dict"]
|
||||
|
||||
assert len(original_dict) == len(restored_dict)
|
||||
|
||||
# Check each key-value pair and key type preservation
|
||||
for orig_key, orig_value in original_dict.items():
|
||||
# Key should exist and have same value
|
||||
assert orig_key in restored_dict
|
||||
assert restored_dict[orig_key] == orig_value
|
||||
|
||||
# Find the corresponding restored key to check type
|
||||
for rest_key in restored_dict.keys():
|
||||
if rest_key == orig_key:
|
||||
assert type(orig_key) is type(rest_key)
|
||||
break
|
||||
|
||||
def test_edge_case_empty_stats(self):
|
||||
"""Test serialization with empty stats."""
|
||||
imputer = SimpleImputer(columns=["A"], strategy="constant", fill_value=0)
|
||||
# Constant strategy doesn't need fitting, so stats will be empty
|
||||
|
||||
serialized = imputer.serialize()
|
||||
deserialized = SimpleImputer.deserialize(serialized)
|
||||
|
||||
assert deserialized.stats_ == {}
|
||||
assert deserialized.strategy == "constant"
|
||||
assert deserialized.fill_value == 0
|
||||
assert deserialized._is_fittable is False
|
||||
|
||||
def test_edge_case_none_values(self):
|
||||
"""Test serialization with None values in stats."""
|
||||
imputer = SimpleImputer(columns=["A"], strategy="mean")
|
||||
imputer._fitted = True
|
||||
imputer.stats_ = {
|
||||
"mean(A)": None,
|
||||
"count(A)": 0,
|
||||
"complex_dict": {
|
||||
None: "none_key",
|
||||
"none_value": None,
|
||||
(None, "tuple"): "tuple_with_none",
|
||||
},
|
||||
}
|
||||
|
||||
serialized = imputer.serialize()
|
||||
deserialized = SimpleImputer.deserialize(serialized)
|
||||
|
||||
assert deserialized.stats_ == imputer.stats_
|
||||
assert deserialized.stats_["mean(A)"] is None
|
||||
assert None in deserialized.stats_["complex_dict"]
|
||||
|
||||
def test_nested_complex_structures(self):
|
||||
"""Test deeply nested complex data structures."""
|
||||
imputer = SimpleImputer(columns=["A"], strategy="mean")
|
||||
imputer._fitted = True
|
||||
|
||||
# Create deeply nested structure with various key types
|
||||
imputer.stats_ = {
|
||||
"nested_structure": {
|
||||
("level1", "tuple"): {
|
||||
frozenset([1, 2]): "frozenset_key",
|
||||
42: {"nested_dict": "value"},
|
||||
None: [1, 2, 3],
|
||||
True: {"another": {"level": "deep"}},
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
serialized = imputer.serialize()
|
||||
deserialized = SimpleImputer.deserialize(serialized)
|
||||
|
||||
assert deserialized.stats_ == imputer.stats_
|
||||
|
||||
# Verify specific nested access works
|
||||
nested = deserialized.stats_["nested_structure"]
|
||||
tuple_key = ("level1", "tuple")
|
||||
assert tuple_key in nested
|
||||
assert frozenset([1, 2]) in nested[tuple_key]
|
||||
|
||||
def test_unknown_preprocessor_type(self):
|
||||
"""Test error when trying to deserialize unknown preprocessor type."""
|
||||
import cloudpickle
|
||||
|
||||
# Create fake serialized data with unknown type
|
||||
unknown_data = {
|
||||
"type": "NonExistentPreprocessor",
|
||||
"version": 1,
|
||||
"fields": {"columns": ["test"]},
|
||||
"stats": {},
|
||||
"stats_type": "cloudpickle",
|
||||
}
|
||||
|
||||
fake_serialized = (
|
||||
SerializablePreprocessorBase.MAGIC_CLOUDPICKLE
|
||||
+ cloudpickle.dumps(unknown_data)
|
||||
)
|
||||
|
||||
with pytest.raises(UnknownPreprocessorError) as exc_info:
|
||||
SerializablePreprocessorBase.deserialize(fake_serialized)
|
||||
|
||||
# Verify the exception contains the correct preprocessor type
|
||||
assert exc_info.value.preprocessor_type == "NonExistentPreprocessor"
|
||||
assert "Unknown preprocessor type: NonExistentPreprocessor" in str(
|
||||
exc_info.value
|
||||
)
|
||||
|
||||
def test_file_system_integration(self):
|
||||
"""Test integration with file system operations."""
|
||||
imputer = SimpleImputer(columns=["value"], strategy="mean")
|
||||
df = pd.DataFrame({"value": [1, 2, None, 4]})
|
||||
dataset = ray.data.from_pandas(df)
|
||||
fitted = imputer.fit(dataset)
|
||||
|
||||
# Test with binary files (CloudPickle)
|
||||
with tempfile.NamedTemporaryFile(mode="wb", suffix=".cloudpickle") as f:
|
||||
# Save as CloudPickle
|
||||
serialized = fitted.serialize()
|
||||
f.write(serialized)
|
||||
f.flush()
|
||||
|
||||
# Load from file
|
||||
with open(f.name, "rb") as read_f:
|
||||
loaded_data = read_f.read()
|
||||
|
||||
deserialized = SerializablePreprocessorBase.deserialize(loaded_data)
|
||||
assert isinstance(deserialized, SimpleImputer)
|
||||
assert abs(deserialized.stats_["mean(value)"] - 2.333333333333333) < 1e-10
|
||||
|
||||
def test_special_numeric_values(self):
|
||||
"""Test serialization with inf, -inf, and NaN values."""
|
||||
# Test with inf fill_value
|
||||
imputer1 = SimpleImputer(columns=["col"], strategy="mean")
|
||||
imputer1.stats_ = {"mean(col)": float("inf")}
|
||||
imputer1._fitted = True
|
||||
|
||||
serialized = imputer1.serialize()
|
||||
deserialized = SerializablePreprocessorBase.deserialize(serialized)
|
||||
assert np.isinf(deserialized.stats_["mean(col)"])
|
||||
|
||||
# Test with -inf fill_value
|
||||
imputer2 = SimpleImputer(columns=["col"], strategy="mean")
|
||||
imputer2.stats_ = {"mean(col)": float("-inf")}
|
||||
imputer2._fitted = True
|
||||
|
||||
serialized = imputer2.serialize()
|
||||
deserialized = SerializablePreprocessorBase.deserialize(serialized)
|
||||
assert (
|
||||
np.isinf(deserialized.stats_["mean(col)"])
|
||||
and deserialized.stats_["mean(col)"] < 0
|
||||
)
|
||||
|
||||
# Test with NaN fill_value
|
||||
imputer3 = SimpleImputer(columns=["col"], strategy="mean")
|
||||
imputer3.stats_ = {"mean(col)": float("nan")}
|
||||
imputer3._fitted = True
|
||||
|
||||
serialized = imputer3.serialize()
|
||||
deserialized = SerializablePreprocessorBase.deserialize(serialized)
|
||||
assert np.isnan(deserialized.stats_["mean(col)"])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,121 @@
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data.preprocessors import Normalizer
|
||||
|
||||
|
||||
def test_normalizer():
|
||||
"""Tests basic Normalizer functionality."""
|
||||
|
||||
col_a = [10, 10, 10]
|
||||
col_b = [1, 3, 3]
|
||||
col_c = [2, 4, -4]
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
||||
ds = ray.data.from_pandas(in_df)
|
||||
|
||||
# l2 norm
|
||||
normalizer = Normalizer(["B", "C"])
|
||||
transformed = normalizer.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
processed_col_a = col_a
|
||||
processed_col_b = [1 / np.sqrt(5), 0.6, 0.6]
|
||||
processed_col_c = [2 / np.sqrt(5), 0.8, -0.8]
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{"A": processed_col_a, "B": processed_col_b, "C": processed_col_c}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
# l1 norm
|
||||
normalizer = Normalizer(["B", "C"], norm="l1")
|
||||
|
||||
transformed = normalizer.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
processed_col_a = col_a
|
||||
processed_col_b = [1 / 3, 3 / 7, 3 / 7]
|
||||
processed_col_c = [2 / 3, 4 / 7, -4 / 7]
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{"A": processed_col_a, "B": processed_col_b, "C": processed_col_c}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
# max norm
|
||||
normalizer = Normalizer(["B", "C"], norm="max")
|
||||
|
||||
transformed = normalizer.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
processed_col_a = col_a
|
||||
processed_col_b = [0.5, 0.75, 0.75]
|
||||
processed_col_c = [1.0, 1.0, -1.0]
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{"A": processed_col_a, "B": processed_col_b, "C": processed_col_c}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
# append mode
|
||||
with pytest.raises(ValueError):
|
||||
Normalizer(columns=["B", "C"], output_columns=["B_encoded"])
|
||||
|
||||
normalizer = Normalizer(["B", "C"], output_columns=["B_normalized", "C_normalized"])
|
||||
transformed = normalizer.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
processed_col_a = col_a
|
||||
processed_col_b = [1 / np.sqrt(5), 0.6, 0.6]
|
||||
processed_col_c = [2 / np.sqrt(5), 0.8, -0.8]
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": col_a,
|
||||
"B": col_b,
|
||||
"C": col_c,
|
||||
"B_normalized": processed_col_b,
|
||||
"C_normalized": processed_col_c,
|
||||
}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
|
||||
def test_normalizer_serialization():
|
||||
"""Test Normalizer serialization and deserialization functionality."""
|
||||
from ray.data.preprocessor import SerializablePreprocessorBase
|
||||
|
||||
# Create normalizer with test data
|
||||
normalizer = Normalizer(columns=["A", "B"], norm="l1")
|
||||
|
||||
# Serialize using CloudPickle
|
||||
serialized = normalizer.serialize()
|
||||
|
||||
# Verify it's binary CloudPickle format
|
||||
assert isinstance(serialized, bytes)
|
||||
assert serialized.startswith(SerializablePreprocessorBase.MAGIC_CLOUDPICKLE)
|
||||
|
||||
# Deserialize
|
||||
deserialized = Normalizer.deserialize(serialized)
|
||||
|
||||
# Verify type and field values
|
||||
assert isinstance(deserialized, Normalizer)
|
||||
assert deserialized.columns == ["A", "B"]
|
||||
assert deserialized.norm == "l1"
|
||||
assert deserialized.output_columns == ["A", "B"]
|
||||
|
||||
# Verify it works correctly
|
||||
df = pd.DataFrame({"A": [3.0, 4.0], "B": [4.0, 3.0]})
|
||||
result = deserialized.transform_batch(df)
|
||||
|
||||
# For l1 norm, values should sum to 1 for each row
|
||||
assert abs(result["A"][0] + result["B"][0] - 1.0) < 1e-10
|
||||
assert abs(result["A"][1] + result["B"][1] - 1.0) < 1e-10
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,570 @@
|
||||
import re
|
||||
import warnings
|
||||
from typing import Dict, Union
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pyarrow
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data.aggregate import Mean
|
||||
from ray.data.constants import MAX_REPR_LENGTH
|
||||
from ray.data.preprocessor import Preprocessor
|
||||
from ray.data.preprocessors import (
|
||||
Categorizer,
|
||||
Chain,
|
||||
Concatenator,
|
||||
CountVectorizer,
|
||||
FeatureHasher,
|
||||
HashingVectorizer,
|
||||
LabelEncoder,
|
||||
MaxAbsScaler,
|
||||
MinMaxScaler,
|
||||
MultiHotEncoder,
|
||||
Normalizer,
|
||||
OneHotEncoder,
|
||||
OrdinalEncoder,
|
||||
PowerTransformer,
|
||||
RobustScaler,
|
||||
SimpleImputer,
|
||||
StandardScaler,
|
||||
Tokenizer,
|
||||
TorchVisionPreprocessor,
|
||||
)
|
||||
from ray.data.util.data_batch_conversion import BatchFormat
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def create_dummy_preprocessors():
|
||||
class DummyPreprocessorWithNothing(Preprocessor):
|
||||
_is_fittable = False
|
||||
|
||||
class DummyPreprocessorWithPandas(DummyPreprocessorWithNothing):
|
||||
def _transform_pandas(self, df: "pd.DataFrame") -> "pd.DataFrame":
|
||||
return df
|
||||
|
||||
class DummyPreprocessorWithNumpy(DummyPreprocessorWithNothing):
|
||||
batch_format = "numpy"
|
||||
|
||||
def _transform_numpy(
|
||||
self, np_data: Union[np.ndarray, Dict[str, np.ndarray]]
|
||||
) -> Union[np.ndarray, Dict[str, np.ndarray]]:
|
||||
return np_data
|
||||
|
||||
class DummyPreprocessorWithPandasAndNumpy(DummyPreprocessorWithNothing):
|
||||
def _transform_pandas(self, df: "pd.DataFrame") -> "pd.DataFrame":
|
||||
return df
|
||||
|
||||
def _transform_numpy(
|
||||
self, np_data: Union[np.ndarray, Dict[str, np.ndarray]]
|
||||
) -> Union[np.ndarray, Dict[str, np.ndarray]]:
|
||||
return np_data
|
||||
|
||||
class DummyPreprocessorWithPandasAndNumpyPreferred(DummyPreprocessorWithNothing):
|
||||
def _transform_pandas(self, df: "pd.DataFrame") -> "pd.DataFrame":
|
||||
return df
|
||||
|
||||
def _transform_numpy(
|
||||
self, np_data: Union[np.ndarray, Dict[str, np.ndarray]]
|
||||
) -> Union[np.ndarray, Dict[str, np.ndarray]]:
|
||||
return np_data
|
||||
|
||||
def preferred_batch_format(cls) -> BatchFormat:
|
||||
return BatchFormat.NUMPY
|
||||
|
||||
yield (
|
||||
DummyPreprocessorWithNothing(),
|
||||
DummyPreprocessorWithPandas(),
|
||||
DummyPreprocessorWithNumpy(),
|
||||
DummyPreprocessorWithPandasAndNumpy(),
|
||||
DummyPreprocessorWithPandasAndNumpyPreferred(),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"preprocessor",
|
||||
[
|
||||
Categorizer(columns=["X"]),
|
||||
CountVectorizer(columns=["X"]),
|
||||
Chain(StandardScaler(columns=["X"]), MinMaxScaler(columns=["X"])),
|
||||
FeatureHasher(columns=["X"], num_features=1, output_column="X_transformed"),
|
||||
HashingVectorizer(columns=["X"], num_features=1),
|
||||
LabelEncoder(label_column="X"),
|
||||
MaxAbsScaler(columns=["X"]),
|
||||
MinMaxScaler(columns=["X"]),
|
||||
MultiHotEncoder(columns=["X"]),
|
||||
Normalizer(columns=["X"]),
|
||||
OneHotEncoder(columns=["X"]),
|
||||
OrdinalEncoder(columns=["X"]),
|
||||
PowerTransformer(columns=["X"], power=1),
|
||||
RobustScaler(columns=["X"]),
|
||||
SimpleImputer(columns=["X"]),
|
||||
StandardScaler(columns=["X"]),
|
||||
Concatenator(columns=["X"]),
|
||||
Tokenizer(columns=["X"]),
|
||||
],
|
||||
)
|
||||
def test_repr(preprocessor):
|
||||
representation = repr(preprocessor)
|
||||
|
||||
assert len(representation) < MAX_REPR_LENGTH
|
||||
pattern = re.compile(f"^{preprocessor.__class__.__name__}\\((.*)\\)$")
|
||||
assert pattern.match(representation)
|
||||
|
||||
|
||||
def test_fitted_preprocessor_without_stats():
|
||||
"""Tests that Preprocessors can be fitted without needing to set self.stats_."""
|
||||
|
||||
class FittablePreprocessor(Preprocessor):
|
||||
def _fit(self, ds):
|
||||
return self
|
||||
|
||||
preprocessor = FittablePreprocessor()
|
||||
ds = ray.data.from_items([1])
|
||||
_ = preprocessor.fit(ds)
|
||||
assert preprocessor.fit_status() == Preprocessor.FitStatus.FITTED
|
||||
|
||||
|
||||
def test_fitted_preprocessor_with_stats():
|
||||
"""Tests that Preprocessors can be fitted by setting an attribute that ends
|
||||
with _."""
|
||||
|
||||
class FittablePreprocessor(Preprocessor):
|
||||
...
|
||||
|
||||
preprocessor = FittablePreprocessor()
|
||||
preprocessor.stats_ = True
|
||||
assert preprocessor.fit_status() == Preprocessor.FitStatus.FITTED
|
||||
|
||||
|
||||
@patch.object(warnings, "warn")
|
||||
def test_fit_twice(mocked_warn):
|
||||
"""Tests that a warning msg should be printed."""
|
||||
col_a = [-1, 0, 1]
|
||||
col_b = [1, 3, 5]
|
||||
col_c = [1, 1, None]
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
||||
ds = ray.data.from_pandas(in_df)
|
||||
|
||||
scaler = MinMaxScaler(["B", "C"])
|
||||
|
||||
# Fit data.
|
||||
scaler.fit(ds)
|
||||
assert scaler.stats_ == {"min(B)": 1, "max(B)": 5, "min(C)": 1, "max(C)": 1}
|
||||
|
||||
ds = ds.map_batches(lambda x: {k: v * 2 for k, v in x.items()})
|
||||
# Fit again
|
||||
scaler.fit(ds)
|
||||
# Assert that the fitted state is corresponding to the second ds.
|
||||
assert scaler.stats_ == {"min(B)": 2, "max(B)": 10, "min(C)": 2, "max(C)": 2}
|
||||
msg = (
|
||||
"`fit` has already been called on the preprocessor (or at least one "
|
||||
"contained preprocessors if this is a chain). "
|
||||
"All previously fitted state will be overwritten!"
|
||||
)
|
||||
mocked_warn.assert_called_once_with(msg)
|
||||
|
||||
|
||||
def test_fit_twice_clears_stale_stats():
|
||||
"""Tests that fit() clears stale stats when stat keys are data-dependent.
|
||||
|
||||
When a preprocessor's stat keys depend on the data (e.g., auto-detected columns),
|
||||
calling fit() again on a different dataset should not retain stale stats from
|
||||
the previous fit. This ensures that fit(A).fit(B) is equivalent to fit(B).
|
||||
"""
|
||||
|
||||
class DataDependentPreprocessor(Preprocessor):
|
||||
"""A preprocessor whose stat keys depend on the data columns present."""
|
||||
|
||||
_is_fittable = True
|
||||
|
||||
def _fit(self, ds):
|
||||
# Dynamically detect columns from the dataset schema
|
||||
schema = ds.schema()
|
||||
column_names = list(schema.names)
|
||||
self.stat_computation_plan.add_aggregator(
|
||||
aggregator_fn=Mean,
|
||||
columns=column_names,
|
||||
)
|
||||
return self
|
||||
|
||||
def _transform_pandas(self, df):
|
||||
return df
|
||||
|
||||
# Dataset A has columns: "a", "b"
|
||||
dataset_a = ray.data.from_items(
|
||||
[
|
||||
{"a": 1.0, "b": 10.0},
|
||||
{"a": 2.0, "b": 20.0},
|
||||
{"a": 3.0, "b": 30.0},
|
||||
]
|
||||
)
|
||||
|
||||
# Dataset B has columns: "b", "c" (note: "a" is missing, "c" is new)
|
||||
dataset_b = ray.data.from_items(
|
||||
[
|
||||
{"b": 100.0, "c": 1000.0},
|
||||
{"b": 200.0, "c": 2000.0},
|
||||
{"b": 300.0, "c": 3000.0},
|
||||
]
|
||||
)
|
||||
|
||||
preprocessor = DataDependentPreprocessor()
|
||||
|
||||
# First fit on dataset A
|
||||
preprocessor.fit(dataset_a)
|
||||
assert preprocessor.stats_ == {"mean(a)": 2.0, "mean(b)": 20.0}
|
||||
|
||||
# Second fit on dataset B - stale stats should be cleared
|
||||
preprocessor.fit(dataset_b)
|
||||
|
||||
# Verify stale stat "mean(a)" is NOT present
|
||||
# Verify stats are correct after refit, and stale stats are cleared.
|
||||
expected_stats = {"mean(b)": 200.0, "mean(c)": 2000.0}
|
||||
assert preprocessor.stats_ == expected_stats, (
|
||||
f"Stats after refit are incorrect. "
|
||||
f"Expected: {expected_stats}, Got: {preprocessor.stats_}"
|
||||
)
|
||||
|
||||
|
||||
def test_transform_all_configs():
|
||||
batch_size = 2
|
||||
num_cpus = 2
|
||||
concurrency = 2
|
||||
memory = 1024
|
||||
|
||||
class DummyPreprocessor(Preprocessor):
|
||||
_is_fittable = False
|
||||
|
||||
def _get_transform_config(self):
|
||||
return {"batch_size": batch_size}
|
||||
|
||||
def _transform_numpy(self, data):
|
||||
assert ray.get_runtime_context().get_assigned_resources()["CPU"] == num_cpus
|
||||
assert (
|
||||
ray.get_runtime_context().get_assigned_resources()["memory"] == memory
|
||||
)
|
||||
# Read(10 rows) → Limit(5) → Transform(batch_size=2)
|
||||
assert (
|
||||
len(data["value"]) <= batch_size
|
||||
) # The last batch is size 1, and limit pushdown resulted in the transform occurring for fewer rows.
|
||||
return data
|
||||
|
||||
def _transform_pandas(self, data):
|
||||
raise RuntimeError(
|
||||
"Pandas transform should not be called with numpy batch format."
|
||||
)
|
||||
|
||||
def _determine_transform_to_use(self):
|
||||
return "numpy"
|
||||
|
||||
prep = DummyPreprocessor()
|
||||
ds = ray.data.from_pandas(pd.DataFrame({"value": list(range(10))}))
|
||||
ds = prep.transform(
|
||||
ds,
|
||||
num_cpus=num_cpus,
|
||||
memory=memory,
|
||||
concurrency=concurrency,
|
||||
)
|
||||
assert [x["value"] for x in ds.take(5)] == [0, 1, 2, 3, 4]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dataset_format", ["simple", "pandas", "arrow"])
|
||||
def test_transform_all_formats(create_dummy_preprocessors, dataset_format):
|
||||
(
|
||||
with_nothing,
|
||||
with_pandas,
|
||||
with_numpy,
|
||||
with_pandas_and_numpy,
|
||||
with_pandas_and_numpy_preferred,
|
||||
) = create_dummy_preprocessors
|
||||
|
||||
if dataset_format == "simple":
|
||||
ds = ray.data.range(10)
|
||||
elif dataset_format == "pandas":
|
||||
df = pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=["A", "B", "C"])
|
||||
ds = ray.data.from_pandas(df)
|
||||
elif dataset_format == "arrow":
|
||||
df = pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=["A", "B", "C"])
|
||||
ds = ray.data.from_arrow(pyarrow.Table.from_pandas(df))
|
||||
else:
|
||||
raise ValueError(f"Untested dataset_format configuration: {dataset_format}.")
|
||||
|
||||
with pytest.raises(NotImplementedError):
|
||||
with_nothing.transform(ds)
|
||||
|
||||
patcher = patch.object(ray.data.dataset.Dataset, "map_batches")
|
||||
|
||||
with patcher as mock_map_batches:
|
||||
with_pandas.transform(ds)
|
||||
mock_map_batches.assert_called_once_with(
|
||||
with_pandas._transform_pandas,
|
||||
batch_format=BatchFormat.PANDAS,
|
||||
zero_copy_batch=True,
|
||||
)
|
||||
|
||||
with patcher as mock_map_batches:
|
||||
with_numpy.transform(ds)
|
||||
mock_map_batches.assert_called_once_with(
|
||||
with_numpy._transform_numpy,
|
||||
batch_format=BatchFormat.NUMPY,
|
||||
zero_copy_batch=True,
|
||||
)
|
||||
|
||||
# Pandas preferred by default.
|
||||
with patcher as mock_map_batches:
|
||||
with_pandas_and_numpy.transform(ds)
|
||||
mock_map_batches.assert_called_once_with(
|
||||
with_pandas_and_numpy._transform_pandas,
|
||||
batch_format=BatchFormat.PANDAS,
|
||||
zero_copy_batch=True,
|
||||
)
|
||||
|
||||
with patcher as mock_map_batches:
|
||||
with_pandas_and_numpy_preferred.transform(ds)
|
||||
mock_map_batches.assert_called_once_with(
|
||||
with_pandas_and_numpy_preferred._transform_numpy,
|
||||
batch_format=BatchFormat.NUMPY,
|
||||
zero_copy_batch=True,
|
||||
)
|
||||
|
||||
|
||||
def test_numpy_pandas_support_transform_batch_wrong_format(create_dummy_preprocessors):
|
||||
# Case 1: simple dataset. No support
|
||||
(
|
||||
with_nothing,
|
||||
with_pandas,
|
||||
with_numpy,
|
||||
with_pandas_and_numpy,
|
||||
with_pandas_and_numpy_preferred,
|
||||
) = create_dummy_preprocessors
|
||||
|
||||
batch = [1, 2, 3]
|
||||
with pytest.raises(ValueError):
|
||||
with_nothing.transform_batch(batch)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
with_pandas.transform_batch(batch)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
with_numpy.transform_batch(batch)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
with_pandas_and_numpy.transform_batch(batch)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
with_pandas_and_numpy_preferred.transform_batch(batch)
|
||||
|
||||
|
||||
def test_numpy_pandas_support_transform_batch_pandas(create_dummy_preprocessors):
|
||||
# Case 2: pandas dataset
|
||||
(
|
||||
with_nothing,
|
||||
with_pandas,
|
||||
with_numpy,
|
||||
with_pandas_and_numpy,
|
||||
with_pandas_and_numpy_preferred,
|
||||
) = create_dummy_preprocessors
|
||||
|
||||
df = pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=["A", "B", "C"])
|
||||
df_single_column = pd.DataFrame([1, 2, 3], columns=["A"])
|
||||
with pytest.raises(NotImplementedError):
|
||||
with_nothing.transform_batch(df)
|
||||
with pytest.raises(NotImplementedError):
|
||||
with_nothing.transform_batch(df_single_column)
|
||||
|
||||
assert isinstance(with_pandas.transform_batch(df), pd.DataFrame)
|
||||
assert isinstance(with_pandas.transform_batch(df_single_column), pd.DataFrame)
|
||||
|
||||
assert isinstance(with_numpy.transform_batch(df), (np.ndarray, dict))
|
||||
# We can get pd.DataFrame after returning numpy data from UDF
|
||||
assert isinstance(with_numpy.transform_batch(df_single_column), (np.ndarray, dict))
|
||||
|
||||
assert isinstance(with_pandas_and_numpy.transform_batch(df), pd.DataFrame)
|
||||
assert isinstance(
|
||||
with_pandas_and_numpy.transform_batch(df_single_column), pd.DataFrame
|
||||
)
|
||||
|
||||
assert isinstance(
|
||||
with_pandas_and_numpy_preferred.transform_batch(df), (np.ndarray, dict)
|
||||
)
|
||||
assert isinstance(
|
||||
with_pandas_and_numpy_preferred.transform_batch(df_single_column),
|
||||
(np.ndarray, dict),
|
||||
)
|
||||
|
||||
|
||||
def test_numpy_pandas_support_transform_batch_arrow(create_dummy_preprocessors):
|
||||
# Case 3: arrow dataset
|
||||
(
|
||||
with_nothing,
|
||||
with_pandas,
|
||||
with_numpy,
|
||||
with_pandas_and_numpy,
|
||||
with_pandas_and_numpy_preferred,
|
||||
) = create_dummy_preprocessors
|
||||
|
||||
df = pd.DataFrame([[1, 2, 3], [4, 5, 6]], columns=["A", "B", "C"])
|
||||
df_single_column = pd.DataFrame([1, 2, 3], columns=["A"])
|
||||
|
||||
table = pyarrow.Table.from_pandas(df)
|
||||
table_single_column = pyarrow.Table.from_pandas(df_single_column)
|
||||
with pytest.raises(NotImplementedError):
|
||||
with_nothing.transform_batch(table)
|
||||
with pytest.raises(NotImplementedError):
|
||||
with_nothing.transform_batch(table_single_column)
|
||||
|
||||
assert isinstance(with_pandas.transform_batch(table), pd.DataFrame)
|
||||
assert isinstance(with_pandas.transform_batch(table_single_column), pd.DataFrame)
|
||||
|
||||
assert isinstance(with_numpy.transform_batch(table), (np.ndarray, dict))
|
||||
# We can get pyarrow.Table after returning numpy data from UDF
|
||||
assert isinstance(
|
||||
with_numpy.transform_batch(table_single_column), (np.ndarray, dict)
|
||||
)
|
||||
|
||||
assert isinstance(with_pandas_and_numpy.transform_batch(table), pd.DataFrame)
|
||||
assert isinstance(
|
||||
with_pandas_and_numpy.transform_batch(table_single_column), pd.DataFrame
|
||||
)
|
||||
|
||||
assert isinstance(
|
||||
with_pandas_and_numpy_preferred.transform_batch(table), (np.ndarray, dict)
|
||||
)
|
||||
assert isinstance(
|
||||
with_pandas_and_numpy_preferred.transform_batch(table_single_column),
|
||||
(np.ndarray, dict),
|
||||
)
|
||||
|
||||
|
||||
def test_numpy_pandas_support_transform_batch_tensor(create_dummy_preprocessors):
|
||||
# Case 4: tensor dataset created by from numpy data directly
|
||||
(
|
||||
with_nothing,
|
||||
with_pandas,
|
||||
with_numpy,
|
||||
with_pandas_and_numpy,
|
||||
with_pandas_and_numpy_preferred,
|
||||
) = create_dummy_preprocessors
|
||||
np_data = np.arange(12).reshape(3, 2, 2)
|
||||
np_single_column = {"A": np.arange(12).reshape(3, 2, 2)}
|
||||
np_multi_column = {
|
||||
"A": np.arange(12).reshape(3, 2, 2),
|
||||
"B": np.arange(12, 24).reshape(3, 2, 2),
|
||||
}
|
||||
|
||||
with pytest.raises(NotImplementedError):
|
||||
with_nothing.transform_batch(np_data)
|
||||
with pytest.raises(NotImplementedError):
|
||||
with_nothing.transform_batch(np_single_column)
|
||||
with pytest.raises(NotImplementedError):
|
||||
with_nothing.transform_batch(np_multi_column)
|
||||
|
||||
assert isinstance(with_pandas.transform_batch(np_data), pd.DataFrame)
|
||||
assert isinstance(with_pandas.transform_batch(np_single_column), pd.DataFrame)
|
||||
assert isinstance(with_pandas.transform_batch(np_multi_column), pd.DataFrame)
|
||||
|
||||
assert isinstance(with_numpy.transform_batch(np_data), np.ndarray)
|
||||
assert isinstance(with_numpy.transform_batch(np_single_column), dict)
|
||||
assert isinstance(with_numpy.transform_batch(np_multi_column), dict)
|
||||
|
||||
assert isinstance(with_pandas_and_numpy.transform_batch(np_data), pd.DataFrame)
|
||||
assert isinstance(
|
||||
with_pandas_and_numpy.transform_batch(np_single_column), pd.DataFrame
|
||||
)
|
||||
assert isinstance(
|
||||
with_pandas_and_numpy.transform_batch(np_multi_column), pd.DataFrame
|
||||
)
|
||||
|
||||
assert isinstance(
|
||||
with_pandas_and_numpy_preferred.transform_batch(np_data), np.ndarray
|
||||
)
|
||||
assert isinstance(
|
||||
with_pandas_and_numpy_preferred.transform_batch(np_single_column), dict
|
||||
)
|
||||
assert isinstance(
|
||||
with_pandas_and_numpy_preferred.transform_batch(np_multi_column), dict
|
||||
)
|
||||
|
||||
|
||||
def test_get_input_output_columns():
|
||||
"""Tests get_input_columns() and get_output_columns() methods."""
|
||||
# Test with preprocessors that have columns attribute
|
||||
scaler = StandardScaler(columns=["A", "B"])
|
||||
assert scaler.get_input_columns() == ["A", "B"]
|
||||
assert scaler.get_output_columns() == ["A", "B"]
|
||||
|
||||
# Test with output_columns specified
|
||||
scaler_with_output = StandardScaler(
|
||||
columns=["A", "B"], output_columns=["A_scaled", "B_scaled"]
|
||||
)
|
||||
assert scaler_with_output.get_input_columns() == ["A", "B"]
|
||||
assert scaler_with_output.get_output_columns() == ["A_scaled", "B_scaled"]
|
||||
|
||||
# Test with encoders
|
||||
encoder = OneHotEncoder(columns=["X", "Y"])
|
||||
assert encoder.get_input_columns() == ["X", "Y"]
|
||||
assert encoder.get_output_columns() == ["X", "Y"]
|
||||
|
||||
encoder_with_output = OneHotEncoder(
|
||||
columns=["X", "Y"], output_columns=["X_encoded", "Y_encoded"]
|
||||
)
|
||||
assert encoder_with_output.get_input_columns() == ["X", "Y"]
|
||||
assert encoder_with_output.get_output_columns() == ["X_encoded", "Y_encoded"]
|
||||
|
||||
# Test LabelEncoder without output_column (in-place transformation)
|
||||
label_encoder = LabelEncoder(label_column="target")
|
||||
assert label_encoder.get_input_columns() == ["target"]
|
||||
assert label_encoder.get_output_columns() == ["target"]
|
||||
|
||||
# Test LabelEncoder with output_column (append mode)
|
||||
label_encoder = LabelEncoder(label_column="target", output_column="target_encoded")
|
||||
assert label_encoder.get_input_columns() == ["target"]
|
||||
assert label_encoder.get_output_columns() == ["target_encoded"]
|
||||
|
||||
# Test Concatenator (uses output_column_name instead of output_columns)
|
||||
concatenator = Concatenator(columns=["A", "B"])
|
||||
assert concatenator.get_input_columns() == ["A", "B"]
|
||||
assert concatenator.get_output_columns() == ["concat_out"]
|
||||
|
||||
concatenator_with_output = Concatenator(
|
||||
columns=["A", "B"], output_column_name="AB_concat"
|
||||
)
|
||||
assert concatenator_with_output.get_input_columns() == ["A", "B"]
|
||||
assert concatenator_with_output.get_output_columns() == ["AB_concat"]
|
||||
|
||||
# Test FeatureHasher (uses output_column instead of output_columns)
|
||||
feature_hasher = FeatureHasher(
|
||||
columns=["token1", "token2"], num_features=8, output_column="hashed"
|
||||
)
|
||||
assert feature_hasher.get_input_columns() == ["token1", "token2"]
|
||||
assert feature_hasher.get_output_columns() == ["hashed"]
|
||||
|
||||
# Test TorchVisionPreprocessor (uses _columns and _output_columns)
|
||||
torch_preprocessor = TorchVisionPreprocessor(
|
||||
columns=["image"], transform=lambda x: x
|
||||
)
|
||||
assert torch_preprocessor.get_input_columns() == ["image"]
|
||||
assert torch_preprocessor.get_output_columns() == ["image"]
|
||||
|
||||
torch_preprocessor_with_output = TorchVisionPreprocessor(
|
||||
columns=["image"], transform=lambda x: x, output_columns=["image_transformed"]
|
||||
)
|
||||
assert torch_preprocessor_with_output.get_input_columns() == ["image"]
|
||||
assert torch_preprocessor_with_output.get_output_columns() == ["image_transformed"]
|
||||
|
||||
# Test with preprocessor without columns attribute
|
||||
class CustomPreprocessor(Preprocessor):
|
||||
_is_fittable = False
|
||||
|
||||
custom = CustomPreprocessor()
|
||||
assert custom.get_input_columns() == []
|
||||
assert custom.get_output_columns() == []
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,891 @@
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data.preprocessor import (
|
||||
PreprocessorNotFittedException,
|
||||
SerializablePreprocessorBase,
|
||||
)
|
||||
from ray.data.preprocessors import (
|
||||
MaxAbsScaler,
|
||||
MinMaxScaler,
|
||||
RobustScaler,
|
||||
StandardScaler,
|
||||
)
|
||||
|
||||
|
||||
def test_min_max_scaler():
|
||||
"""Tests basic MinMaxScaler functionality."""
|
||||
col_a = [-1, 0, 1]
|
||||
col_b = [1, 3, 5]
|
||||
col_c = [1, 1, None]
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
||||
ds = ray.data.from_pandas(in_df)
|
||||
|
||||
scaler = MinMaxScaler(["B", "C"])
|
||||
|
||||
# Transform with unfitted preprocessor.
|
||||
with pytest.raises(PreprocessorNotFittedException):
|
||||
scaler.transform(ds)
|
||||
|
||||
# Fit data.
|
||||
scaler.fit(ds)
|
||||
assert scaler.stats_ == {"min(B)": 1, "max(B)": 5, "min(C)": 1, "max(C)": 1}
|
||||
|
||||
transformed = scaler.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
processed_col_a = col_a
|
||||
processed_col_b = [0.0, 0.5, 1.0]
|
||||
processed_col_c = [0.0, 0.0, None]
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{"A": processed_col_a, "B": processed_col_b, "C": processed_col_c}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df)
|
||||
|
||||
# Transform batch.
|
||||
pred_col_a = [1, 2, 3]
|
||||
pred_col_b = [3, 5, 7]
|
||||
pred_col_c = [0, 1, 2]
|
||||
pred_in_df = pd.DataFrame.from_dict(
|
||||
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c}
|
||||
)
|
||||
|
||||
pred_out_df = scaler.transform_batch(pred_in_df)
|
||||
|
||||
pred_processed_col_a = pred_col_a
|
||||
pred_processed_col_b = [0.5, 1.0, 1.5]
|
||||
pred_processed_col_c = [-1.0, 0.0, 1.0]
|
||||
pred_expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": pred_processed_col_a,
|
||||
"B": pred_processed_col_b,
|
||||
"C": pred_processed_col_c,
|
||||
}
|
||||
).astype(pred_out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(pred_out_df, pred_expected_df)
|
||||
|
||||
# append mode
|
||||
with pytest.raises(ValueError):
|
||||
MinMaxScaler(columns=["B", "C"], output_columns=["B_mm_scaled"])
|
||||
|
||||
scaler = MinMaxScaler(
|
||||
columns=["B", "C"], output_columns=["B_mm_scaled", "C_mm_scaled"]
|
||||
)
|
||||
scaler.fit(ds)
|
||||
|
||||
pred_in_df = pd.DataFrame.from_dict(
|
||||
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c}
|
||||
)
|
||||
pred_out_df = scaler.transform_batch(pred_in_df)
|
||||
|
||||
pred_expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": pred_col_a,
|
||||
"B": pred_col_b,
|
||||
"C": pred_col_c,
|
||||
"B_mm_scaled": pred_processed_col_b,
|
||||
"C_mm_scaled": pred_processed_col_c,
|
||||
}
|
||||
).astype(pred_out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(pred_out_df, pred_expected_df, check_like=True)
|
||||
|
||||
|
||||
def test_max_abs_scaler():
|
||||
"""Tests basic MaxAbsScaler functionality."""
|
||||
col_a = [-1, 0, 1]
|
||||
col_b = [1, 3, -5]
|
||||
col_c = [1, 1, None]
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
||||
ds = ray.data.from_pandas(in_df)
|
||||
|
||||
scaler = MaxAbsScaler(["B", "C"])
|
||||
|
||||
# Transform with unfitted preprocessor.
|
||||
with pytest.raises(PreprocessorNotFittedException):
|
||||
scaler.transform(ds)
|
||||
|
||||
# Fit data.
|
||||
scaler.fit(ds)
|
||||
assert scaler.stats_ == {"abs_max(B)": 5, "abs_max(C)": 1}
|
||||
|
||||
transformed = scaler.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
processed_col_a = col_a
|
||||
processed_col_b = [0.2, 0.6, -1.0]
|
||||
processed_col_c = [1.0, 1.0, None]
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{"A": processed_col_a, "B": processed_col_b, "C": processed_col_c}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
# Transform batch.
|
||||
pred_col_a = [1, 2, 3]
|
||||
pred_col_b = [3, 5, 7]
|
||||
pred_col_c = [0, 1, -2]
|
||||
pred_in_df = pd.DataFrame.from_dict(
|
||||
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c}
|
||||
)
|
||||
|
||||
pred_out_df = scaler.transform_batch(pred_in_df)
|
||||
|
||||
pred_processed_col_a = pred_col_a
|
||||
pred_processed_col_b = [0.6, 1.0, 1.4]
|
||||
pred_processed_col_c = [0.0, 1.0, -2.0]
|
||||
pred_expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": pred_processed_col_a,
|
||||
"B": pred_processed_col_b,
|
||||
"C": pred_processed_col_c,
|
||||
}
|
||||
).astype(pred_out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(pred_out_df, pred_expected_df, check_like=True)
|
||||
|
||||
# append mode
|
||||
with pytest.raises(ValueError):
|
||||
MaxAbsScaler(columns=["B", "C"], output_columns=["B_ma_scaled"])
|
||||
|
||||
scaler = MaxAbsScaler(
|
||||
columns=["B", "C"], output_columns=["B_ma_scaled", "C_ma_scaled"]
|
||||
)
|
||||
scaler.fit(ds)
|
||||
|
||||
pred_in_df = pd.DataFrame.from_dict(
|
||||
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c}
|
||||
)
|
||||
pred_out_df = scaler.transform_batch(pred_in_df)
|
||||
|
||||
pred_expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": pred_col_a,
|
||||
"B": pred_col_b,
|
||||
"C": pred_col_c,
|
||||
"B_ma_scaled": pred_processed_col_b,
|
||||
"C_ma_scaled": pred_processed_col_c,
|
||||
}
|
||||
).astype(pred_out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(pred_out_df, pred_expected_df, check_like=True)
|
||||
|
||||
|
||||
def test_robust_scaler():
|
||||
"""Tests basic RobustScaler functionality."""
|
||||
col_a = [-2, -1, 0, 1, 2]
|
||||
col_b = [-2, -1, 0, 1, 2]
|
||||
col_c = [-10, 1, 2, 3, 10]
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
||||
ds = ray.data.from_pandas(in_df)
|
||||
|
||||
scaler = RobustScaler(["B", "C"])
|
||||
|
||||
# Transform with unfitted preprocessor.
|
||||
with pytest.raises(PreprocessorNotFittedException):
|
||||
scaler.transform(ds)
|
||||
|
||||
# Fit data.
|
||||
scaler.fit(ds)
|
||||
assert scaler.stats_ == {
|
||||
"low_quantile(B)": -1,
|
||||
"median(B)": 0,
|
||||
"high_quantile(B)": 1,
|
||||
"low_quantile(C)": 1,
|
||||
"median(C)": 2,
|
||||
"high_quantile(C)": 3,
|
||||
}
|
||||
|
||||
transformed = scaler.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
processed_col_a = col_a
|
||||
processed_col_b = [-1.0, -0.5, 0, 0.5, 1.0]
|
||||
processed_col_c = [-6, -0.5, 0, 0.5, 4]
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{"A": processed_col_a, "B": processed_col_b, "C": processed_col_c}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
# Transform batch.
|
||||
pred_col_a = [1, 2, 3]
|
||||
pred_col_b = [3, 5, 7]
|
||||
pred_col_c = [0, 1, 2]
|
||||
pred_in_df = pd.DataFrame.from_dict(
|
||||
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c}
|
||||
)
|
||||
|
||||
pred_out_df = scaler.transform_batch(pred_in_df)
|
||||
|
||||
pred_processed_col_a = pred_col_a
|
||||
pred_processed_col_b = [1.5, 2.5, 3.5]
|
||||
pred_processed_col_c = [-1.0, -0.5, 0.0]
|
||||
pred_expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": pred_processed_col_a,
|
||||
"B": pred_processed_col_b,
|
||||
"C": pred_processed_col_c,
|
||||
}
|
||||
).astype(pred_out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(pred_out_df, pred_expected_df, check_like=True)
|
||||
|
||||
# append mode
|
||||
with pytest.raises(ValueError):
|
||||
RobustScaler(columns=["B", "C"], output_columns=["B_r_scaled"])
|
||||
|
||||
scaler = RobustScaler(
|
||||
columns=["B", "C"], output_columns=["B_r_scaled", "C_r_scaled"]
|
||||
)
|
||||
scaler.fit(ds)
|
||||
|
||||
pred_in_df = pd.DataFrame.from_dict(
|
||||
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c}
|
||||
)
|
||||
pred_out_df = scaler.transform_batch(pred_in_df)
|
||||
|
||||
pred_expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": pred_col_a,
|
||||
"B": pred_col_b,
|
||||
"C": pred_col_c,
|
||||
"B_r_scaled": pred_processed_col_b,
|
||||
"C_r_scaled": pred_processed_col_c,
|
||||
}
|
||||
).astype(pred_out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(pred_out_df, pred_expected_df, check_like=True)
|
||||
|
||||
|
||||
def test_standard_scaler():
|
||||
"""Tests basic StandardScaler functionality."""
|
||||
col_a = [-1, 0, 1, 2]
|
||||
col_b = [1, 1, 5, 5]
|
||||
col_c = [1, 1, 1, None]
|
||||
col_d = [None, None, None, None]
|
||||
sample_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c, "D": col_d})
|
||||
ds = ray.data.from_pandas(sample_df)
|
||||
|
||||
scaler = StandardScaler(["B", "C", "D"])
|
||||
|
||||
# Transform with unfitted preprocessor.
|
||||
with pytest.raises(PreprocessorNotFittedException):
|
||||
scaler.transform(ds)
|
||||
|
||||
# Fit data.
|
||||
scaler = scaler.fit(ds)
|
||||
assert scaler.stats_ == {
|
||||
"mean(B)": 3.0,
|
||||
"mean(C)": 1.0,
|
||||
"mean(D)": None,
|
||||
"std(B)": 2.0,
|
||||
"std(C)": 0.0,
|
||||
"std(D)": None,
|
||||
}
|
||||
|
||||
# Transform data.
|
||||
in_col_a = [-1, 0, 1, 2]
|
||||
in_col_b = [1, 1, 5, 5]
|
||||
in_col_c = [1, 1, 1, None]
|
||||
in_col_d = [0, None, None, None]
|
||||
in_df = pd.DataFrame.from_dict(
|
||||
{"A": in_col_a, "B": in_col_b, "C": in_col_c, "D": in_col_d}
|
||||
)
|
||||
in_ds = ray.data.from_pandas(in_df)
|
||||
transformed = scaler.transform(in_ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
processed_col_a = col_a
|
||||
processed_col_b = [-1.0, -1.0, 1.0, 1.0]
|
||||
processed_col_c = [0.0, 0.0, 0.0, None]
|
||||
processed_col_d = [np.nan, np.nan, np.nan, np.nan]
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": processed_col_a,
|
||||
"B": processed_col_b,
|
||||
"C": processed_col_c,
|
||||
"D": processed_col_d,
|
||||
}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
# Transform batch.
|
||||
pred_col_a = [1, 2, 3]
|
||||
pred_col_b = [3, 5, 7]
|
||||
pred_col_c = [0, 1, 2]
|
||||
pred_col_d = [None, None, None]
|
||||
pred_in_df = pd.DataFrame.from_dict(
|
||||
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c, "D": pred_col_d}
|
||||
)
|
||||
|
||||
pred_out_df = scaler.transform_batch(pred_in_df)
|
||||
|
||||
pred_processed_col_a = pred_col_a
|
||||
pred_processed_col_b = [0.0, 1.0, 2.0]
|
||||
pred_processed_col_c = [-1.0, 0.0, 1.0]
|
||||
pred_processed_col_d = [None, None, None]
|
||||
pred_expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": pred_processed_col_a,
|
||||
"B": pred_processed_col_b,
|
||||
"C": pred_processed_col_c,
|
||||
"D": pred_processed_col_d,
|
||||
}
|
||||
).astype(pred_out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(pred_out_df, pred_expected_df, check_like=True)
|
||||
|
||||
# append mode
|
||||
with pytest.raises(ValueError):
|
||||
StandardScaler(columns=["B", "C"], output_columns=["B_s_scaled"])
|
||||
|
||||
scaler = StandardScaler(
|
||||
columns=["B", "C"], output_columns=["B_s_scaled", "C_s_scaled"]
|
||||
)
|
||||
scaler.fit(ds)
|
||||
|
||||
pred_in_df = pd.DataFrame.from_dict(
|
||||
{"A": pred_col_a, "B": pred_col_b, "C": pred_col_c}
|
||||
)
|
||||
pred_out_df = scaler.transform_batch(pred_in_df)
|
||||
|
||||
pred_expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": pred_col_a,
|
||||
"B": pred_col_b,
|
||||
"C": pred_col_c,
|
||||
"B_s_scaled": pred_processed_col_b,
|
||||
"C_s_scaled": pred_processed_col_c,
|
||||
}
|
||||
).astype(pred_out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(pred_out_df, pred_expected_df, check_like=True)
|
||||
|
||||
|
||||
def test_standard_scaler_arrow_transform():
|
||||
"""Test the StandardScaler _transform_arrow method directly."""
|
||||
# Create test data
|
||||
col_a = ["red", "green", "blue", "red"]
|
||||
col_b = [1.0, 3.0, 5.0, 7.0] # mean=4, std=2.236
|
||||
col_c = [10.0, 10.0, 10.0, 10.0] # constant column, std=0
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b, "C": col_c})
|
||||
|
||||
scaler = StandardScaler(["B", "C"])
|
||||
scaler.fit(ray.data.from_pandas(in_df))
|
||||
|
||||
# Create Arrow table for transformation
|
||||
table = pa.Table.from_pandas(in_df)
|
||||
|
||||
# Transform using Arrow
|
||||
result_table = scaler._transform_arrow(table)
|
||||
|
||||
# Verify result is an Arrow table
|
||||
assert isinstance(result_table, pa.Table)
|
||||
|
||||
# Convert to pandas for easier comparison
|
||||
result_df = result_table.to_pandas()
|
||||
|
||||
# Expected encoding:
|
||||
# B: (x - mean(B)) / std(B)
|
||||
# C: std(C)=0 -> std becomes 1 -> (x - mean(C)) / 1 = 0 for all
|
||||
b_mean = scaler.stats_["mean(B)"]
|
||||
b_std = scaler.stats_["std(B)"] or 0.0
|
||||
if b_std == 0:
|
||||
b_std = 1
|
||||
expected_col_b = [(x - b_mean) / b_std for x in col_b]
|
||||
|
||||
c_mean = scaler.stats_["mean(C)"]
|
||||
c_std = scaler.stats_["std(C)"] or 0.0
|
||||
if c_std == 0:
|
||||
c_std = 1
|
||||
expected_col_c = [(x - c_mean) / c_std for x in col_c]
|
||||
|
||||
assert result_df["A"].tolist() == col_a, "Column A should be unchanged"
|
||||
assert np.allclose(
|
||||
result_df["B"].tolist(), expected_col_b
|
||||
), f"Column B mismatch: {result_df['B'].tolist()}"
|
||||
assert np.allclose(
|
||||
result_df["C"].tolist(), expected_col_c
|
||||
), f"Column C mismatch: {result_df['C'].tolist()}"
|
||||
|
||||
|
||||
def test_standard_scaler_arrow_transform_append_mode():
|
||||
"""Test the StandardScaler _transform_arrow method in append mode."""
|
||||
col_a = ["red", "green", "blue"]
|
||||
col_b = [1.0, 3.0, 5.0]
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b})
|
||||
|
||||
scaler = StandardScaler(["B"], output_columns=["B_scaled"])
|
||||
scaler.fit(ray.data.from_pandas(in_df))
|
||||
|
||||
table = pa.Table.from_pandas(in_df)
|
||||
result_table = scaler._transform_arrow(table)
|
||||
result_df = result_table.to_pandas()
|
||||
|
||||
# Original columns should be unchanged
|
||||
assert result_df["A"].tolist() == col_a
|
||||
assert result_df["B"].tolist() == col_b
|
||||
|
||||
# New column should have scaled values: (x - 3) / 2
|
||||
b_mean = scaler.stats_["mean(B)"]
|
||||
b_std = scaler.stats_["std(B)"] or 0.0
|
||||
if b_std == 0:
|
||||
b_std = 1
|
||||
expected_b_scaled = [(x - b_mean) / b_std for x in col_b]
|
||||
assert np.allclose(result_df["B_scaled"].tolist(), expected_b_scaled)
|
||||
|
||||
|
||||
def test_standard_scaler_arrow_transform_null_stats():
|
||||
"""Test the StandardScaler _transform_arrow method with null mean/std."""
|
||||
# Use an all-null column to produce null mean/std during fit.
|
||||
in_df = pd.DataFrame.from_dict({"A": [None, None, None]})
|
||||
|
||||
scaler = StandardScaler(["A"])
|
||||
scaler.fit(ray.data.from_pandas(in_df))
|
||||
|
||||
table = pa.Table.from_pandas(in_df)
|
||||
result_table = scaler._transform_arrow(table)
|
||||
result_df = result_table.to_pandas()
|
||||
|
||||
# All values should be null when mean/std is None
|
||||
assert result_df["A"].isna().all(), "All values should be null when stats are None"
|
||||
|
||||
|
||||
def test_standard_scaler_arrow_transform_overlapping_columns():
|
||||
"""Test StandardScaler _transform_arrow with overlapping input/output columns.
|
||||
|
||||
This tests the case where output_columns[i] == columns[j] for i < j.
|
||||
The Arrow implementation must read all input columns before writing any output
|
||||
to avoid corrupting data that will be read later.
|
||||
"""
|
||||
# columns=['A', 'B'], output_columns=['B', 'C']
|
||||
# Without the fix, B would be overwritten before being read as input
|
||||
col_a = [2.0, 4.0, 6.0] # mean=4, std=2 -> scaled: [-1, 0, 1]
|
||||
col_b = [10.0, 20.0, 30.0] # mean=20, std=10 -> scaled: [-1, 0, 1]
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b})
|
||||
|
||||
scaler = StandardScaler(["A", "B"], output_columns=["B", "C"])
|
||||
scaler.fit(ray.data.from_pandas(in_df))
|
||||
|
||||
# Test Arrow transform
|
||||
table = pa.Table.from_pandas(in_df)
|
||||
result_table = scaler._transform_arrow(table)
|
||||
result_df = result_table.to_pandas()
|
||||
|
||||
# Test pandas transform for comparison
|
||||
pandas_result = scaler._transform_pandas(in_df.copy())
|
||||
|
||||
# Column A should be unchanged (not in output_columns with same index)
|
||||
assert result_df["A"].tolist() == col_a, "Column A should be unchanged"
|
||||
|
||||
# Column B should contain scaled A: (A - 4) / 2 = [-1, 0, 1]
|
||||
a_mean = scaler.stats_["mean(A)"]
|
||||
a_std = scaler.stats_["std(A)"] or 0.0
|
||||
if a_std == 0:
|
||||
a_std = 1
|
||||
expected_b = [(x - a_mean) / a_std for x in col_a]
|
||||
assert np.allclose(result_df["B"].tolist(), expected_b), (
|
||||
f"Column B should contain scaled A. Expected {expected_b}, "
|
||||
f"got {result_df['B'].tolist()}"
|
||||
)
|
||||
|
||||
# Column C should contain scaled B: (B - 20) / 10 = [-1, 0, 1]
|
||||
b_mean = scaler.stats_["mean(B)"]
|
||||
b_std = scaler.stats_["std(B)"] or 0.0
|
||||
if b_std == 0:
|
||||
b_std = 1
|
||||
expected_c = [(x - b_mean) / b_std for x in col_b]
|
||||
assert np.allclose(result_df["C"].tolist(), expected_c), (
|
||||
f"Column C should contain scaled B. Expected {expected_c}, "
|
||||
f"got {result_df['C'].tolist()}"
|
||||
)
|
||||
|
||||
# Arrow and pandas results should match
|
||||
pd.testing.assert_frame_equal(
|
||||
result_df,
|
||||
pandas_result,
|
||||
check_like=True,
|
||||
obj="Arrow vs Pandas transform results should match",
|
||||
)
|
||||
|
||||
|
||||
class TestScalerSerialization:
|
||||
"""Test serialization/deserialization functionality for scaler preprocessors."""
|
||||
|
||||
def setup_method(self):
|
||||
"""Set up test data."""
|
||||
self.test_df = pd.DataFrame(
|
||||
{
|
||||
"feature1": [1, 2, 3, 4, 5],
|
||||
"feature2": [10, 20, 30, 40, 50],
|
||||
"feature3": [100, 200, 300, 400, 500],
|
||||
"other": ["a", "b", "c", "d", "e"],
|
||||
}
|
||||
)
|
||||
self.test_dataset = ray.data.from_pandas(self.test_df)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"scaler_class,fit_data,expected_stats,transform_data",
|
||||
[
|
||||
(
|
||||
StandardScaler,
|
||||
None, # Use default self.test_df
|
||||
{
|
||||
"mean(feature1)": 3.0,
|
||||
"mean(feature2)": 30.0,
|
||||
"std(feature1)": np.sqrt(2.0),
|
||||
"std(feature2)": np.sqrt(200.0),
|
||||
},
|
||||
pd.DataFrame(
|
||||
{
|
||||
"feature1": [6, 7, 8],
|
||||
"feature2": [60, 70, 80],
|
||||
"other": ["f", "g", "h"],
|
||||
}
|
||||
),
|
||||
),
|
||||
(
|
||||
MinMaxScaler,
|
||||
None, # Use default self.test_df
|
||||
{
|
||||
"min(feature1)": 1,
|
||||
"min(feature2)": 10,
|
||||
"max(feature1)": 5,
|
||||
"max(feature2)": 50,
|
||||
},
|
||||
pd.DataFrame(
|
||||
{
|
||||
"feature1": [6, 7, 8],
|
||||
"feature2": [60, 70, 80],
|
||||
"other": ["f", "g", "h"],
|
||||
}
|
||||
),
|
||||
),
|
||||
(
|
||||
MaxAbsScaler,
|
||||
pd.DataFrame(
|
||||
{
|
||||
"feature1": [-5, -2, 0, 2, 5],
|
||||
"feature2": [-50, -20, 0, 20, 50],
|
||||
"other": ["a", "b", "c", "d", "e"],
|
||||
}
|
||||
),
|
||||
{
|
||||
"abs_max(feature1)": 5,
|
||||
"abs_max(feature2)": 50,
|
||||
},
|
||||
pd.DataFrame(
|
||||
{
|
||||
"feature1": [-6, 0, 6],
|
||||
"feature2": [-60, 0, 60],
|
||||
"other": ["f", "g", "h"],
|
||||
}
|
||||
),
|
||||
),
|
||||
(
|
||||
RobustScaler,
|
||||
None, # Use default self.test_df
|
||||
{
|
||||
"low_quantile(feature1)": 2.0,
|
||||
"median(feature1)": 3.0,
|
||||
"high_quantile(feature1)": 4.0,
|
||||
"low_quantile(feature2)": 20.0,
|
||||
"median(feature2)": 30.0,
|
||||
"high_quantile(feature2)": 40.0,
|
||||
},
|
||||
pd.DataFrame(
|
||||
{
|
||||
"feature1": [6, 7, 8],
|
||||
"feature2": [60, 70, 80],
|
||||
"other": ["f", "g", "h"],
|
||||
}
|
||||
),
|
||||
),
|
||||
],
|
||||
ids=["StandardScaler", "MinMaxScaler", "MaxAbsScaler", "RobustScaler"],
|
||||
)
|
||||
def test_scaler_serialization(
|
||||
self, scaler_class, fit_data, expected_stats, transform_data
|
||||
):
|
||||
"""Test scaler serialization for all scaler types."""
|
||||
# Use custom fit data if provided, otherwise use default test dataset
|
||||
if fit_data is not None:
|
||||
fit_dataset = ray.data.from_pandas(fit_data)
|
||||
else:
|
||||
fit_dataset = self.test_dataset
|
||||
|
||||
# Create and fit scaler
|
||||
scaler = scaler_class(columns=["feature1", "feature2"])
|
||||
fitted_scaler = scaler.fit(fit_dataset)
|
||||
|
||||
# Verify fitted stats match expected values
|
||||
assert fitted_scaler.stats_ == expected_stats, (
|
||||
f"Stats mismatch for {scaler_class.__name__}:\n"
|
||||
f"Expected: {expected_stats}\n"
|
||||
f"Got: {fitted_scaler.stats_}"
|
||||
)
|
||||
|
||||
# Test CloudPickle serialization
|
||||
serialized = fitted_scaler.serialize()
|
||||
assert isinstance(serialized, bytes)
|
||||
assert serialized.startswith(SerializablePreprocessorBase.MAGIC_CLOUDPICKLE)
|
||||
|
||||
# Test deserialization
|
||||
deserialized = SerializablePreprocessorBase.deserialize(serialized)
|
||||
assert deserialized.__class__.__name__ == scaler_class.__name__
|
||||
assert deserialized.columns == ["feature1", "feature2"]
|
||||
assert deserialized._fitted
|
||||
|
||||
# Verify stats are preserved after deserialization
|
||||
assert deserialized.stats_ == expected_stats, (
|
||||
f"Deserialized stats mismatch for {scaler_class.__name__}:\n"
|
||||
f"Expected: {expected_stats}\n"
|
||||
f"Got: {deserialized.stats_}"
|
||||
)
|
||||
|
||||
# Verify each stat key exists and has correct value
|
||||
for stat_key, stat_value in expected_stats.items():
|
||||
assert stat_key in deserialized.stats_
|
||||
if isinstance(stat_value, float):
|
||||
assert np.isclose(deserialized.stats_[stat_key], stat_value)
|
||||
else:
|
||||
assert deserialized.stats_[stat_key] == stat_value
|
||||
|
||||
# Test functional equivalence
|
||||
original_result = fitted_scaler.transform_batch(transform_data.copy())
|
||||
deserialized_result = deserialized.transform_batch(transform_data.copy())
|
||||
|
||||
pd.testing.assert_frame_equal(original_result, deserialized_result)
|
||||
|
||||
def test_scaler_with_output_columns_serialization(self):
|
||||
"""Test scaler serialization with custom output columns."""
|
||||
# Test with StandardScaler and output columns
|
||||
scaler = StandardScaler(
|
||||
columns=["feature1", "feature2"],
|
||||
output_columns=["scaled_feature1", "scaled_feature2"],
|
||||
)
|
||||
fitted_scaler = scaler.fit(self.test_dataset)
|
||||
|
||||
# Serialize and deserialize
|
||||
serialized = fitted_scaler.serialize()
|
||||
deserialized = SerializablePreprocessorBase.deserialize(serialized)
|
||||
|
||||
# Verify output columns are preserved
|
||||
assert deserialized.output_columns == ["scaled_feature1", "scaled_feature2"]
|
||||
|
||||
# Test functional equivalence
|
||||
test_df = pd.DataFrame(
|
||||
{"feature1": [6, 7, 8], "feature2": [60, 70, 80], "other": ["f", "g", "h"]}
|
||||
)
|
||||
|
||||
original_result = fitted_scaler.transform_batch(test_df.copy())
|
||||
deserialized_result = deserialized.transform_batch(test_df.copy())
|
||||
|
||||
pd.testing.assert_frame_equal(original_result, deserialized_result)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"scaler_class",
|
||||
[StandardScaler, MinMaxScaler, MaxAbsScaler, RobustScaler],
|
||||
ids=["StandardScaler", "MinMaxScaler", "MaxAbsScaler", "RobustScaler"],
|
||||
)
|
||||
def test_unfitted_scaler_serialization(self, scaler_class):
|
||||
"""Test serialization of unfitted scalers."""
|
||||
# Test unfitted scaler
|
||||
scaler = scaler_class(columns=["feature1", "feature2"])
|
||||
|
||||
# Serialize unfitted scaler
|
||||
serialized = scaler.serialize()
|
||||
deserialized = SerializablePreprocessorBase.deserialize(serialized)
|
||||
|
||||
# Verify it's still unfitted
|
||||
assert not deserialized._fitted
|
||||
assert deserialized.columns == ["feature1", "feature2"]
|
||||
assert deserialized.__class__.__name__ == scaler_class.__name__
|
||||
|
||||
# Should raise error when trying to transform
|
||||
test_df = pd.DataFrame({"feature1": [1, 2, 3], "feature2": [10, 20, 30]})
|
||||
with pytest.raises(PreprocessorNotFittedException):
|
||||
deserialized.transform_batch(test_df)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"scaler_class,expected_stats",
|
||||
[
|
||||
(
|
||||
StandardScaler,
|
||||
{
|
||||
"mean(feature1)": 3.0,
|
||||
"std(feature1)": np.sqrt(2.0),
|
||||
},
|
||||
),
|
||||
(
|
||||
MinMaxScaler,
|
||||
{
|
||||
"min(feature1)": 1,
|
||||
"max(feature1)": 5,
|
||||
},
|
||||
),
|
||||
(
|
||||
MaxAbsScaler,
|
||||
{
|
||||
"abs_max(feature1)": 5,
|
||||
},
|
||||
),
|
||||
(
|
||||
RobustScaler,
|
||||
{
|
||||
"low_quantile(feature1)": 2.0,
|
||||
"median(feature1)": 3.0,
|
||||
"high_quantile(feature1)": 4.0,
|
||||
},
|
||||
),
|
||||
],
|
||||
ids=["StandardScaler", "MinMaxScaler", "MaxAbsScaler", "RobustScaler"],
|
||||
)
|
||||
def test_scaler_stats_preservation(self, scaler_class, expected_stats):
|
||||
"""Test that scaler statistics are perfectly preserved during serialization."""
|
||||
# Create scaler with known stats
|
||||
scaler = scaler_class(columns=["feature1"])
|
||||
fitted_scaler = scaler.fit(self.test_dataset)
|
||||
|
||||
# Verify fitted stats match expected values
|
||||
for stat_key, stat_value in expected_stats.items():
|
||||
assert stat_key in fitted_scaler.stats_
|
||||
if isinstance(stat_value, float):
|
||||
assert np.isclose(fitted_scaler.stats_[stat_key], stat_value)
|
||||
else:
|
||||
assert fitted_scaler.stats_[stat_key] == stat_value
|
||||
|
||||
# Get original stats
|
||||
original_stats = fitted_scaler.stats_.copy()
|
||||
|
||||
# Serialize and deserialize
|
||||
serialized = fitted_scaler.serialize()
|
||||
deserialized = SerializablePreprocessorBase.deserialize(serialized)
|
||||
|
||||
# Verify stats are identical
|
||||
assert deserialized.stats_ == original_stats
|
||||
|
||||
# Verify expected stat values are preserved
|
||||
for stat_key, stat_value in expected_stats.items():
|
||||
assert stat_key in deserialized.stats_
|
||||
if isinstance(stat_value, float):
|
||||
assert np.isclose(deserialized.stats_[stat_key], stat_value)
|
||||
else:
|
||||
assert deserialized.stats_[stat_key] == stat_value
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"scaler_class",
|
||||
[StandardScaler, MinMaxScaler, MaxAbsScaler, RobustScaler],
|
||||
ids=["StandardScaler", "MinMaxScaler", "MaxAbsScaler", "RobustScaler"],
|
||||
)
|
||||
def test_scaler_version_compatibility(self, scaler_class):
|
||||
"""Test that scalers can be deserialized with version support."""
|
||||
# Create and fit scaler
|
||||
scaler = scaler_class(columns=["feature1", "feature2"])
|
||||
fitted_scaler = scaler.fit(self.test_dataset)
|
||||
|
||||
# Serialize
|
||||
serialized = fitted_scaler.serialize()
|
||||
|
||||
# Deserialize and verify version handling
|
||||
deserialized = SerializablePreprocessorBase.deserialize(serialized)
|
||||
assert deserialized.__class__.__name__ == scaler_class.__name__
|
||||
assert deserialized._fitted
|
||||
|
||||
# Test that it works correctly
|
||||
test_df = pd.DataFrame({"feature1": [6, 7, 8], "feature2": [60, 70, 80]})
|
||||
|
||||
result = deserialized.transform_batch(test_df)
|
||||
assert len(result.columns) == 2 # Should have the scaled columns
|
||||
assert "feature1" in result.columns
|
||||
assert "feature2" in result.columns
|
||||
|
||||
|
||||
def test_standard_scaler_near_zero_std():
|
||||
"""Test StandardScaler handles near-zero standard deviation correctly."""
|
||||
# Create data with very small standard deviation (near-constant values)
|
||||
col_a = [1.0, 1.0 + 1e-10, 1.0]
|
||||
col_b = [5, 10, 15] # Normal column for comparison
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b})
|
||||
ds = ray.data.from_pandas(in_df)
|
||||
|
||||
scaler = StandardScaler(["A", "B"])
|
||||
scaler.fit(ds)
|
||||
transformed = scaler.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
# Column A should be scaled to zeros (near-constant)
|
||||
# Instead of NaN or inf values
|
||||
assert np.allclose(
|
||||
out_df["A"], 0.0, atol=1e-6
|
||||
), "Near-constant column should be scaled to zeros"
|
||||
|
||||
# Column B should be normally scaled
|
||||
assert not np.allclose(out_df["B"], 0.0), "Normal column should not be all zeros"
|
||||
|
||||
# No NaN or inf values should be present
|
||||
assert not out_df["A"].isna().any(), "Should not contain NaN values"
|
||||
assert not np.isinf(out_df["A"]).any(), "Should not contain inf values"
|
||||
|
||||
|
||||
def test_min_max_scaler_near_zero_range():
|
||||
"""Test MinMaxScaler handles near-zero range correctly."""
|
||||
# Create data with very small range (near-constant values)
|
||||
col_a = [2.0, 2.0 + 1e-10, 2.0]
|
||||
col_b = [1, 5, 10] # Normal column for comparison
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b})
|
||||
ds = ray.data.from_pandas(in_df)
|
||||
|
||||
scaler = MinMaxScaler(["A", "B"])
|
||||
scaler.fit(ds)
|
||||
transformed = scaler.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
# Column A should be scaled to zeros (near-constant)
|
||||
# Instead of NaN or inf values
|
||||
assert np.allclose(
|
||||
out_df["A"], 0.0, atol=1e-6
|
||||
), "Near-constant column should be scaled to zeros"
|
||||
|
||||
# Column B should be normally scaled
|
||||
expected_b = [0.0, 4 / 9, 1.0]
|
||||
assert np.allclose(
|
||||
out_df["B"], expected_b, atol=1e-6
|
||||
), "Normal column should be scaled correctly"
|
||||
|
||||
# No NaN or inf values should be present
|
||||
assert not out_df["A"].isna().any(), "Should not contain NaN values"
|
||||
assert not np.isinf(out_df["A"]).any(), "Should not contain inf values"
|
||||
|
||||
|
||||
def test_standard_scaler_exact_zero_std():
|
||||
"""Test StandardScaler still handles exact zero standard deviation.
|
||||
|
||||
This is a regression test to ensure the epsilon-based handling
|
||||
doesn't break the existing behavior for exact zero std.
|
||||
"""
|
||||
# Create constant column (exact zero std)
|
||||
col_c = [5, 5, 5]
|
||||
in_df = pd.DataFrame.from_dict({"C": col_c})
|
||||
ds = ray.data.from_pandas(in_df)
|
||||
|
||||
scaler = StandardScaler(["C"])
|
||||
scaler.fit(ds)
|
||||
transformed = scaler.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
# Should be all zeros
|
||||
assert np.allclose(out_df["C"], 0.0), "Constant column should be scaled to zeros"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,118 @@
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data.preprocessors import Tokenizer
|
||||
|
||||
|
||||
def test_tokenizer():
|
||||
"""Tests basic Tokenizer functionality."""
|
||||
|
||||
col_a = ["this is a test", "apple"]
|
||||
col_b = ["the quick brown fox jumps over the lazy dog", "banana banana"]
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b})
|
||||
ds = ray.data.from_pandas(in_df)
|
||||
|
||||
tokenizer = Tokenizer(["A", "B"])
|
||||
transformed = tokenizer.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
processed_col_a = [["this", "is", "a", "test"], ["apple"]]
|
||||
processed_col_b = [
|
||||
["the", "quick", "brown", "fox", "jumps", "over", "the", "lazy", "dog"],
|
||||
["banana", "banana"],
|
||||
]
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{"A": processed_col_a, "B": processed_col_b}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
# Test append mode
|
||||
with pytest.raises(
|
||||
ValueError, match="The length of columns and output_columns must match."
|
||||
):
|
||||
Tokenizer(columns=["A", "B"], output_columns=["A_tokenized"])
|
||||
|
||||
tokenizer = Tokenizer(
|
||||
columns=["A", "B"], output_columns=["A_tokenized", "B_tokenized"]
|
||||
)
|
||||
transformed = tokenizer.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
print(out_df)
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": col_a,
|
||||
"B": col_b,
|
||||
"A_tokenized": processed_col_a,
|
||||
"B_tokenized": processed_col_b,
|
||||
}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
# Test custom tokenization function
|
||||
def custom_tokenizer(s: str) -> list:
|
||||
return s.replace("banana", "fruit").split()
|
||||
|
||||
tokenizer = Tokenizer(
|
||||
columns=["A", "B"],
|
||||
tokenization_fn=custom_tokenizer,
|
||||
output_columns=["A_custom", "B_custom"],
|
||||
)
|
||||
transformed = tokenizer.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
custom_processed_col_a = [["this", "is", "a", "test"], ["apple"]]
|
||||
custom_processed_col_b = [
|
||||
["the", "quick", "brown", "fox", "jumps", "over", "the", "lazy", "dog"],
|
||||
["fruit", "fruit"],
|
||||
]
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": col_a,
|
||||
"B": col_b,
|
||||
"A_custom": custom_processed_col_a,
|
||||
"B_custom": custom_processed_col_b,
|
||||
}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
|
||||
def test_tokenizer_serialization():
|
||||
"""Test Tokenizer serialization and deserialization functionality."""
|
||||
from ray.data.preprocessor import SerializablePreprocessorBase
|
||||
|
||||
# Create tokenizer
|
||||
tokenizer = Tokenizer(columns=["text"])
|
||||
|
||||
# Serialize using CloudPickle
|
||||
serialized = tokenizer.serialize()
|
||||
|
||||
# Verify it's binary CloudPickle format
|
||||
assert isinstance(serialized, bytes)
|
||||
assert serialized.startswith(SerializablePreprocessorBase.MAGIC_CLOUDPICKLE)
|
||||
|
||||
# Deserialize
|
||||
deserialized = Tokenizer.deserialize(serialized)
|
||||
|
||||
# Verify type and field values
|
||||
assert isinstance(deserialized, Tokenizer)
|
||||
assert deserialized.columns == ["text"]
|
||||
assert callable(deserialized.tokenization_fn)
|
||||
assert deserialized.output_columns == ["text"]
|
||||
|
||||
# Verify it works correctly
|
||||
df = pd.DataFrame({"text": ["hello world", "foo bar"]})
|
||||
result = deserialized.transform_batch(df)
|
||||
|
||||
# Verify tokenization was applied correctly
|
||||
assert result["text"][0] == ["hello", "world"]
|
||||
assert result["text"][1] == ["foo", "bar"]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,161 @@
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from torchvision import transforms
|
||||
|
||||
import ray
|
||||
from ray.data.exceptions import UserCodeException
|
||||
from ray.data.preprocessors import TorchVisionPreprocessor
|
||||
|
||||
|
||||
class TestTorchVisionPreprocessor:
|
||||
def test_repr(self):
|
||||
class StubTransform:
|
||||
def __call__(self, tensor):
|
||||
return tensor
|
||||
|
||||
def __repr__(self):
|
||||
return "StubTransform()"
|
||||
|
||||
preprocessor = TorchVisionPreprocessor(
|
||||
columns=["spam"], transform=StubTransform()
|
||||
)
|
||||
assert repr(preprocessor) == (
|
||||
"TorchVisionPreprocessor(columns=['spam'], "
|
||||
"output_columns=['spam'], transform=StubTransform())"
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"transform",
|
||||
[
|
||||
transforms.ToTensor(), # `ToTensor` accepts an `np.ndarray` as input
|
||||
transforms.Lambda(lambda tensor: tensor.permute(2, 0, 1)),
|
||||
],
|
||||
)
|
||||
def test_transform_images(self, transform):
|
||||
dataset = ray.data.from_items(
|
||||
[
|
||||
{"image": np.zeros((32, 32, 3)), "label": 0},
|
||||
{"image": np.zeros((32, 32, 3)), "label": 1},
|
||||
]
|
||||
)
|
||||
preprocessor = TorchVisionPreprocessor(columns=["image"], transform=transform)
|
||||
|
||||
transformed_dataset = preprocessor.transform(dataset)
|
||||
|
||||
assert transformed_dataset.schema().names == ["image", "label"]
|
||||
transformed_images = [
|
||||
record["image"] for record in transformed_dataset.take_all()
|
||||
]
|
||||
assert all(image.shape == (3, 32, 32) for image in transformed_images)
|
||||
assert all(image.dtype == np.double for image in transformed_images)
|
||||
labels = {record["label"] for record in transformed_dataset.take_all()}
|
||||
assert labels == {0, 1}
|
||||
|
||||
def test_batch_transform_images(self):
|
||||
dataset = ray.data.from_items(
|
||||
[
|
||||
{"image": np.zeros((32, 32, 3)), "label": 0},
|
||||
{"image": np.zeros((32, 32, 3)), "label": 1},
|
||||
]
|
||||
)
|
||||
transform = transforms.Compose(
|
||||
[
|
||||
transforms.Lambda(
|
||||
lambda batch: torch.as_tensor(batch).permute(0, 3, 1, 2)
|
||||
),
|
||||
transforms.Resize(64),
|
||||
]
|
||||
)
|
||||
preprocessor = TorchVisionPreprocessor(
|
||||
columns=["image"], transform=transform, batched=True
|
||||
)
|
||||
|
||||
transformed_dataset = preprocessor.transform(dataset)
|
||||
|
||||
assert transformed_dataset.schema().names == ["image", "label"]
|
||||
transformed_images = [
|
||||
record["image"] for record in transformed_dataset.take_all()
|
||||
]
|
||||
assert all(image.shape == (3, 64, 64) for image in transformed_images)
|
||||
assert all(image.dtype == np.double for image in transformed_images)
|
||||
labels = {record["label"] for record in transformed_dataset.take_all()}
|
||||
assert labels == {0, 1}
|
||||
|
||||
def test_transform_ragged_images(self):
|
||||
dataset = ray.data.from_items(
|
||||
[
|
||||
{"image": np.zeros((16, 16, 3)), "label": 0},
|
||||
{"image": np.zeros((32, 32, 3)), "label": 1},
|
||||
]
|
||||
)
|
||||
transform = transforms.ToTensor()
|
||||
preprocessor = TorchVisionPreprocessor(columns=["image"], transform=transform)
|
||||
|
||||
transformed_dataset = preprocessor.transform(dataset)
|
||||
|
||||
assert transformed_dataset.schema().names == ["image", "label"]
|
||||
transformed_images = [
|
||||
record["image"] for record in transformed_dataset.take_all()
|
||||
]
|
||||
assert sorted(image.shape for image in transformed_images) == [
|
||||
(3, 16, 16),
|
||||
(3, 32, 32),
|
||||
]
|
||||
assert all(image.dtype == np.double for image in transformed_images)
|
||||
labels = {record["label"] for record in transformed_dataset.take_all()}
|
||||
assert labels == {0, 1}
|
||||
|
||||
def test_invalid_transform_raises_value_error(self):
|
||||
dataset = ray.data.from_items(
|
||||
[
|
||||
{"image": np.zeros((32, 32, 3)), "label": 0},
|
||||
{"image": np.zeros((32, 32, 3)), "label": 1},
|
||||
]
|
||||
)
|
||||
transform = transforms.Lambda(lambda tensor: "BLAH BLAH INVALID")
|
||||
preprocessor = TorchVisionPreprocessor(columns=["image"], transform=transform)
|
||||
|
||||
with pytest.raises((UserCodeException, ValueError)):
|
||||
preprocessor.transform(dataset).materialize()
|
||||
|
||||
|
||||
def test_torchvision_preprocessor_serialization():
|
||||
"""Test TorchVisionPreprocessor serialization and deserialization functionality."""
|
||||
from torchvision import transforms
|
||||
|
||||
from ray.data.preprocessor import SerializablePreprocessorBase
|
||||
|
||||
# Create preprocessor
|
||||
transform = transforms.Compose([transforms.ToTensor()])
|
||||
preprocessor = TorchVisionPreprocessor(columns=["image"], transform=transform)
|
||||
|
||||
# Serialize using CloudPickle
|
||||
serialized = preprocessor.serialize()
|
||||
|
||||
# Verify it's binary CloudPickle format
|
||||
assert isinstance(serialized, bytes)
|
||||
assert serialized.startswith(SerializablePreprocessorBase.MAGIC_CLOUDPICKLE)
|
||||
|
||||
# Deserialize
|
||||
deserialized = TorchVisionPreprocessor.deserialize(serialized)
|
||||
|
||||
# Verify type and field values
|
||||
assert isinstance(deserialized, TorchVisionPreprocessor)
|
||||
assert deserialized.columns == ["image"]
|
||||
assert isinstance(deserialized.torchvision_transform, type(transform))
|
||||
|
||||
# Verify it works correctly
|
||||
test_data = {"image": np.zeros((32, 32, 3), dtype=np.uint8)}
|
||||
result = deserialized.transform_batch(test_data)
|
||||
|
||||
# Verify transformation was applied - ToTensor converts uint8 [0,255] to float [0.0, 1.0]
|
||||
assert "image" in result
|
||||
assert result["image"].dtype in (np.float32, np.float64)
|
||||
assert result["image"].min() >= 0.0 and result["image"].max() <= 1.0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,146 @@
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data.preprocessors import PowerTransformer
|
||||
|
||||
|
||||
def test_power_transformer():
|
||||
"""Tests basic PowerTransformer functionality."""
|
||||
|
||||
# yeo-johnson
|
||||
col_a = [-1, 0]
|
||||
col_b = [0, 1]
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b})
|
||||
ds = ray.data.from_pandas(in_df)
|
||||
|
||||
# yeo-johnson power=0
|
||||
transformer = PowerTransformer(["A", "B"], power=0)
|
||||
transformed = transformer.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
processed_col_a = [-1.5, 0]
|
||||
processed_col_b = [0, np.log(2)]
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{"A": processed_col_a, "B": processed_col_b}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
# yeo-johnson power=2
|
||||
transformer = PowerTransformer(["A", "B"], power=2)
|
||||
transformed = transformer.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
processed_col_a = [-np.log(2), 0]
|
||||
processed_col_b = [0, 1.5]
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{"A": processed_col_a, "B": processed_col_b}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
# box-cox
|
||||
col_a = [1, 2]
|
||||
col_b = [3, 4]
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b})
|
||||
ds = ray.data.from_pandas(in_df)
|
||||
|
||||
# box-cox power=0
|
||||
transformer = PowerTransformer(["A", "B"], power=0, method="box-cox")
|
||||
transformed = transformer.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
processed_col_a = [0, np.log(2)]
|
||||
processed_col_b = [np.log(3), np.log(4)]
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{"A": processed_col_a, "B": processed_col_b}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
# box-cox power=2
|
||||
transformer = PowerTransformer(["A", "B"], power=2, method="box-cox")
|
||||
transformed = transformer.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
processed_col_a = [0, 1.5]
|
||||
processed_col_b = [4, 7.5]
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{"A": processed_col_a, "B": processed_col_b}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
# Test append mode
|
||||
# First test that providing wrong number of output columns raises error
|
||||
with pytest.raises(
|
||||
ValueError, match="The length of columns and output_columns must match."
|
||||
):
|
||||
PowerTransformer(columns=["A", "B"], power=2, output_columns=["A_transformed"])
|
||||
|
||||
# Test append mode with correct output columns
|
||||
transformer = PowerTransformer(
|
||||
columns=["A", "B"],
|
||||
power=2,
|
||||
method="box-cox",
|
||||
output_columns=["A_transformed", "B_transformed"],
|
||||
)
|
||||
transformed = transformer.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
# Transformed columns should have the expected values
|
||||
processed_col_a = [0, 1.5]
|
||||
processed_col_b = [4, 7.5]
|
||||
|
||||
expected_df = pd.DataFrame(
|
||||
{
|
||||
"A": col_a,
|
||||
"B": col_b,
|
||||
"A_transformed": processed_col_a,
|
||||
"B_transformed": processed_col_b,
|
||||
}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
|
||||
def test_power_transformer_serialization():
|
||||
"""Test PowerTransformer serialization and deserialization functionality."""
|
||||
from ray.data.preprocessor import SerializablePreprocessorBase
|
||||
|
||||
# Create transformer with test data
|
||||
transformer = PowerTransformer(columns=["A", "B"], power=2.0, method="yeo-johnson")
|
||||
|
||||
# Serialize using CloudPickle
|
||||
serialized = transformer.serialize()
|
||||
|
||||
# Verify it's binary CloudPickle format
|
||||
assert isinstance(serialized, bytes)
|
||||
assert serialized.startswith(SerializablePreprocessorBase.MAGIC_CLOUDPICKLE)
|
||||
|
||||
# Deserialize
|
||||
deserialized = PowerTransformer.deserialize(serialized)
|
||||
|
||||
# Verify type and field values
|
||||
assert isinstance(deserialized, PowerTransformer)
|
||||
assert deserialized.columns == ["A", "B"]
|
||||
assert deserialized.power == 2.0
|
||||
assert deserialized.method == "yeo-johnson"
|
||||
assert deserialized.output_columns == ["A", "B"]
|
||||
|
||||
# Verify it works correctly
|
||||
df = pd.DataFrame({"A": [1.0, 2.0, 3.0], "B": [4.0, 5.0, 6.0]})
|
||||
result = deserialized.transform_batch(df.copy())
|
||||
|
||||
# Verify transformation was applied
|
||||
# For power=2, yeo-johnson on positive values: ((x+1)^2 - 1) / 2
|
||||
expected_a_0 = ((1.0 + 1) ** 2.0 - 1) / 2.0
|
||||
assert abs(result["A"][0] - expected_a_0) < 1e-10
|
||||
assert "A" in result.columns
|
||||
assert "B" in result.columns
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,29 @@
|
||||
import pytest
|
||||
|
||||
from ray.data.preprocessors.utils import simple_hash, simple_split_tokenizer
|
||||
|
||||
|
||||
def test_simple_split_tokenizer():
|
||||
# Tests simple_split_tokenizer.
|
||||
assert simple_split_tokenizer("one_word") == ["one_word"]
|
||||
assert simple_split_tokenizer("two words") == ["two", "words"]
|
||||
assert simple_split_tokenizer("One fish. Two fish.") == [
|
||||
"One",
|
||||
"fish.",
|
||||
"Two",
|
||||
"fish.",
|
||||
]
|
||||
|
||||
|
||||
def test_simple_hash():
|
||||
# Tests simple_hash determinism.
|
||||
assert simple_hash(1, 100) == 15
|
||||
assert simple_hash("a", 100) == 99
|
||||
assert simple_hash("banana", 100) == 10
|
||||
assert simple_hash([1, 2, "apple"], 100) == 58
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,234 @@
|
||||
from collections import Counter
|
||||
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data.preprocessors import CountVectorizer, HashingVectorizer
|
||||
|
||||
|
||||
def test_count_vectorizer():
|
||||
"""Tests basic CountVectorizer functionality."""
|
||||
|
||||
# Increase data size & repartition to test for
|
||||
# discuss.ray.io/t/xgboost-ray-crashes-when-used-for-multiclass-text-classification
|
||||
row_multiplier = 100000
|
||||
|
||||
col_a = ["a b b c c c", "a a a a c"] * row_multiplier
|
||||
col_b = ["apple", "banana banana banana"] * row_multiplier
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b})
|
||||
ds = ray.data.from_pandas(in_df).repartition(10)
|
||||
|
||||
vectorizer = CountVectorizer(["A", "B"])
|
||||
vectorizer.fit(ds)
|
||||
assert vectorizer.stats_ == {
|
||||
"token_counts(A)": Counter(
|
||||
{"a": 5 * row_multiplier, "c": 4 * row_multiplier, "b": 2 * row_multiplier}
|
||||
),
|
||||
"token_counts(B)": Counter(
|
||||
{"banana": 3 * row_multiplier, "apple": 1 * row_multiplier}
|
||||
),
|
||||
}
|
||||
|
||||
transformed = vectorizer.transform(ds)
|
||||
out_df = transformed.to_pandas(limit=float("inf"))
|
||||
|
||||
processed_col_a = [[1, 3, 2], [4, 1, 0]] * row_multiplier
|
||||
processed_col_b = [[0, 1], [3, 0]] * row_multiplier
|
||||
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": processed_col_a,
|
||||
"B": processed_col_b,
|
||||
}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
# max_features
|
||||
vectorizer = CountVectorizer(["A", "B"], max_features=2)
|
||||
vectorizer.fit(ds)
|
||||
|
||||
assert vectorizer.stats_ == {
|
||||
"token_counts(A)": Counter({"a": 5 * row_multiplier, "c": 4 * row_multiplier}),
|
||||
"token_counts(B)": Counter(
|
||||
{"banana": 3 * row_multiplier, "apple": 1 * row_multiplier}
|
||||
),
|
||||
}
|
||||
|
||||
transformed = vectorizer.transform(ds)
|
||||
out_df = transformed.to_pandas(limit=float("inf"))
|
||||
|
||||
processed_col_a = [[1, 3], [4, 1]] * row_multiplier
|
||||
processed_col_b = [[0, 1], [3, 0]] * row_multiplier
|
||||
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": processed_col_a,
|
||||
"B": processed_col_b,
|
||||
}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
# Test append mode
|
||||
with pytest.raises(
|
||||
ValueError, match="The length of columns and output_columns must match."
|
||||
):
|
||||
CountVectorizer(
|
||||
columns=["A", "B"],
|
||||
output_columns=[
|
||||
"A_counts"
|
||||
], # Should provide same number of output columns as input
|
||||
)
|
||||
|
||||
vectorizer = CountVectorizer(["A", "B"], output_columns=["A_counts", "B_counts"])
|
||||
vectorizer.fit(ds)
|
||||
|
||||
transformed = vectorizer.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
processed_col_a = [[1, 3, 2], [4, 1, 0]] * row_multiplier
|
||||
processed_col_b = [[0, 1], [3, 0]] * row_multiplier
|
||||
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": col_a,
|
||||
"B": col_b,
|
||||
"A_counts": processed_col_a,
|
||||
"B_counts": processed_col_b,
|
||||
}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
|
||||
def test_hashing_vectorizer():
|
||||
"""Tests basic HashingVectorizer functionality."""
|
||||
|
||||
col_a = ["a b b c c c", "a a a a c"]
|
||||
col_b = ["apple", "banana banana banana"]
|
||||
in_df = pd.DataFrame.from_dict({"A": col_a, "B": col_b})
|
||||
ds = ray.data.from_pandas(in_df)
|
||||
|
||||
vectorizer = HashingVectorizer(["A", "B"], num_features=3)
|
||||
|
||||
transformed = vectorizer.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
processed_col_a = [[0, 4, 2], [0, 5, 0]]
|
||||
processed_col_b = [[0, 0, 1], [3, 0, 0]]
|
||||
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{"A": processed_col_a, "B": processed_col_b}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
# Test append mode
|
||||
with pytest.raises(
|
||||
ValueError, match="The length of columns and output_columns must match."
|
||||
):
|
||||
HashingVectorizer(
|
||||
columns=["A", "B"],
|
||||
num_features=3,
|
||||
output_columns=[
|
||||
"A_hashed"
|
||||
], # Should provide same number of output columns as input
|
||||
)
|
||||
|
||||
vectorizer = HashingVectorizer(
|
||||
["A", "B"], num_features=3, output_columns=["A_hashed", "B_hashed"]
|
||||
)
|
||||
|
||||
transformed = vectorizer.transform(ds)
|
||||
out_df = transformed.to_pandas()
|
||||
|
||||
expected_df = pd.DataFrame.from_dict(
|
||||
{
|
||||
"A": col_a,
|
||||
"B": col_b,
|
||||
"A_hashed": processed_col_a,
|
||||
"B_hashed": processed_col_b,
|
||||
}
|
||||
).astype(out_df.dtypes.to_dict())
|
||||
|
||||
pd.testing.assert_frame_equal(out_df, expected_df, check_like=True)
|
||||
|
||||
|
||||
def test_hashing_vectorizer_serialization():
|
||||
"""Test HashingVectorizer serialization and deserialization functionality."""
|
||||
from ray.data.preprocessor import SerializablePreprocessorBase
|
||||
|
||||
# Create vectorizer
|
||||
vectorizer = HashingVectorizer(columns=["text"], num_features=16)
|
||||
|
||||
# Serialize using CloudPickle
|
||||
serialized = vectorizer.serialize()
|
||||
|
||||
# Verify it's binary CloudPickle format
|
||||
assert isinstance(serialized, bytes)
|
||||
assert serialized.startswith(SerializablePreprocessorBase.MAGIC_CLOUDPICKLE)
|
||||
|
||||
# Deserialize
|
||||
deserialized = HashingVectorizer.deserialize(serialized)
|
||||
|
||||
# Verify type and field values
|
||||
assert isinstance(deserialized, HashingVectorizer)
|
||||
assert deserialized.columns == ["text"]
|
||||
assert deserialized.num_features == 16
|
||||
assert callable(deserialized.tokenization_fn)
|
||||
assert deserialized.output_columns == ["text"]
|
||||
|
||||
# Verify it works correctly
|
||||
df = pd.DataFrame({"text": ["hello world", "foo bar"]})
|
||||
result = deserialized.transform_batch(df)
|
||||
|
||||
# Verify vectorization was applied correctly
|
||||
assert "text" in result.columns
|
||||
assert len(result["text"][0]) == 16
|
||||
assert len(result["text"][1]) == 16
|
||||
|
||||
|
||||
def test_count_vectorizer_serialization():
|
||||
"""Test CountVectorizer serialization and deserialization functionality."""
|
||||
import ray
|
||||
from ray.data.preprocessor import SerializablePreprocessorBase
|
||||
|
||||
# Create and fit vectorizer
|
||||
vectorizer = CountVectorizer(columns=["text"], max_features=5)
|
||||
df = pd.DataFrame({"text": ["hello world", "foo bar", "hello foo"]})
|
||||
ds = ray.data.from_pandas(df)
|
||||
fitted_vectorizer = vectorizer.fit(ds)
|
||||
|
||||
# Serialize using CloudPickle
|
||||
serialized = fitted_vectorizer.serialize()
|
||||
|
||||
# Verify it's binary CloudPickle format
|
||||
assert isinstance(serialized, bytes)
|
||||
assert serialized.startswith(SerializablePreprocessorBase.MAGIC_CLOUDPICKLE)
|
||||
|
||||
# Deserialize
|
||||
deserialized = CountVectorizer.deserialize(serialized)
|
||||
|
||||
# Verify type and field values
|
||||
assert isinstance(deserialized, CountVectorizer)
|
||||
assert deserialized._fitted
|
||||
assert deserialized.columns == ["text"]
|
||||
assert deserialized.max_features == 5
|
||||
|
||||
# Verify stats are preserved
|
||||
assert "token_counts(text)" in deserialized.stats_
|
||||
|
||||
# Verify it works correctly
|
||||
test_df = pd.DataFrame({"text": ["hello world"]})
|
||||
result = deserialized.transform_batch(test_df)
|
||||
|
||||
# Verify vectorization was applied correctly
|
||||
assert "text" in result.columns
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,709 @@
|
||||
import time
|
||||
from contextlib import contextmanager
|
||||
from types import MethodType
|
||||
from typing import Optional
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data import ExecutionResources
|
||||
from ray.data._internal.actor_autoscaler import (
|
||||
ActorPoolScalingRequest,
|
||||
DefaultActorAutoscaler,
|
||||
)
|
||||
from ray.data._internal.actor_autoscaler.default_actor_autoscaler import (
|
||||
_get_max_scale_up,
|
||||
)
|
||||
from ray.data._internal.execution.operators.actor_pool_map_operator import _ActorPool
|
||||
from ray.data._internal.execution.operators.base_physical_operator import (
|
||||
InternalQueueOperatorMixin,
|
||||
)
|
||||
from ray.data._internal.execution.resource_manager import ResourceManager
|
||||
from ray.data._internal.execution.streaming_executor_state import OpState
|
||||
from ray.data.context import (
|
||||
AutoscalingConfig,
|
||||
)
|
||||
|
||||
|
||||
def test_actor_pool_scaling():
|
||||
"""Test `_actor_pool_should_scale_up` and `_actor_pool_should_scale_down`
|
||||
in `DefaultAutoscaler`"""
|
||||
|
||||
resource_manager = MagicMock(
|
||||
spec=ResourceManager,
|
||||
get_budget=MagicMock(return_value=None),
|
||||
get_allocation=MagicMock(return_value=None),
|
||||
)
|
||||
autoscaler = DefaultActorAutoscaler(
|
||||
topology=MagicMock(),
|
||||
resource_manager=resource_manager,
|
||||
config=AutoscalingConfig(
|
||||
actor_pool_util_upscaling_threshold=1.0,
|
||||
actor_pool_util_downscaling_threshold=0.5,
|
||||
actor_pool_max_upscaling_delta=None,
|
||||
),
|
||||
)
|
||||
|
||||
# Current actor pool utilization is 0.9, which is above the threshold.
|
||||
actor_pool: _ActorPool = MagicMock(
|
||||
spec=_ActorPool,
|
||||
min_size=MagicMock(return_value=5),
|
||||
max_size=MagicMock(return_value=15),
|
||||
current_size=MagicMock(return_value=10),
|
||||
num_active_actors=MagicMock(return_value=10),
|
||||
num_running_actors=MagicMock(return_value=10),
|
||||
num_pending_actors=MagicMock(return_value=0),
|
||||
num_tasks_in_flight=MagicMock(return_value=15),
|
||||
per_actor_resource_usage=MagicMock(return_value=ExecutionResources(cpu=1)),
|
||||
max_tasks_in_flight_per_actor=MagicMock(return_value=2),
|
||||
max_actor_concurrency=MagicMock(return_value=1),
|
||||
get_pool_util=MagicMock(
|
||||
# NOTE: Unittest mocking library doesn't support proxying to actual
|
||||
# non-mocked methods so we have emulate it by directly binding existing
|
||||
# method of `get_pool_util` to a mocked object
|
||||
side_effect=lambda: MethodType(_ActorPool.get_pool_util, actor_pool)()
|
||||
),
|
||||
)
|
||||
|
||||
op = MagicMock(
|
||||
spec=InternalQueueOperatorMixin,
|
||||
has_completed=MagicMock(return_value=False),
|
||||
_inputs_complete=False,
|
||||
input_dependencies=[MagicMock()],
|
||||
internal_input_queue_num_blocks=MagicMock(return_value=1),
|
||||
metrics=MagicMock(average_num_inputs_per_task=1, num_inputs_received=1),
|
||||
num_output_splits=MagicMock(return_value=1),
|
||||
)
|
||||
op_state = OpState(
|
||||
op, inqueues=[MagicMock(__len__=MagicMock(return_value=10), num_blocks=10)]
|
||||
)
|
||||
op_state._scheduling_status = MagicMock(under_resource_limits=True)
|
||||
|
||||
@contextmanager
|
||||
def patch(mock, attr, value, is_method=True):
|
||||
original = getattr(mock, attr)
|
||||
if is_method:
|
||||
value = MagicMock(return_value=value)
|
||||
setattr(mock, attr, value)
|
||||
yield
|
||||
setattr(mock, attr, original)
|
||||
|
||||
def assert_autoscaling_action(
|
||||
*, delta: int, expected_reason: Optional[str], force: bool = False
|
||||
):
|
||||
nonlocal actor_pool, op, op_state
|
||||
|
||||
assert autoscaler._derive_target_scaling_config(
|
||||
actor_pool=actor_pool,
|
||||
op=op,
|
||||
op_state=op_state,
|
||||
) == ActorPoolScalingRequest(delta=delta, force=force, reason=expected_reason)
|
||||
|
||||
# Should scale up since the util above the threshold.
|
||||
assert actor_pool.get_pool_util() == 1.5
|
||||
assert_autoscaling_action(
|
||||
delta=5,
|
||||
expected_reason="utilization of 1.5 >= 1.0",
|
||||
)
|
||||
|
||||
# Should scale up immediately when the actor pool has no running actors.
|
||||
with patch(actor_pool, "num_running_actors", 0):
|
||||
with patch(actor_pool, "get_pool_util", float("inf")):
|
||||
assert_autoscaling_action(
|
||||
delta=1,
|
||||
expected_reason="no running actors, scale up immediately",
|
||||
)
|
||||
|
||||
# Should be no-op since the util is below the threshold.
|
||||
with patch(actor_pool, "num_tasks_in_flight", 9):
|
||||
assert actor_pool.get_pool_util() == 0.9
|
||||
assert_autoscaling_action(
|
||||
delta=0, expected_reason="utilization of 0.9 w/in limits [0.5, 1.0]"
|
||||
)
|
||||
|
||||
# Should be no-op since there are pending actors (no downscaling while pending)
|
||||
with patch(actor_pool, "num_pending_actors", 1):
|
||||
with patch(actor_pool, "num_tasks_in_flight", 4):
|
||||
assert actor_pool.get_pool_util() == 0.4
|
||||
assert_autoscaling_action(
|
||||
delta=0,
|
||||
expected_reason="no downscaling while actors are pending",
|
||||
)
|
||||
|
||||
# Should be no-op since we have reached the max size (ie could not scale
|
||||
# up even though utilization > threshold)
|
||||
with patch(actor_pool, "current_size", 15):
|
||||
with patch(actor_pool, "num_tasks_in_flight", 20):
|
||||
assert_autoscaling_action(
|
||||
delta=0,
|
||||
expected_reason="reached max size",
|
||||
)
|
||||
|
||||
# Should be no-op since we have reached the min size (ie could not scale
|
||||
# down even though utilization < threshold)
|
||||
with patch(actor_pool, "current_size", 5):
|
||||
with patch(actor_pool, "num_tasks_in_flight", 2):
|
||||
assert_autoscaling_action(
|
||||
delta=0,
|
||||
expected_reason="reached min size",
|
||||
)
|
||||
|
||||
# Should scale up since the pool is below the min size.
|
||||
with patch(actor_pool, "current_size", 4):
|
||||
assert_autoscaling_action(
|
||||
delta=1,
|
||||
expected_reason="pool below min size",
|
||||
)
|
||||
|
||||
# Should scale down since if the op is completed, or
|
||||
# the op has no more inputs.
|
||||
with patch(op, "has_completed", True):
|
||||
# NOTE: We simulate actor pool dipping below min size upon
|
||||
# completion (to verify that it will be able to scale to 0)
|
||||
with patch(actor_pool, "current_size", 5):
|
||||
assert_autoscaling_action(
|
||||
delta=-1,
|
||||
expected_reason="consumed all inputs",
|
||||
force=True,
|
||||
)
|
||||
|
||||
# Should scale down only once all inputs have been already dispatched AND
|
||||
# no new inputs ar expected
|
||||
with patch(op_state.input_queues[0], "num_blocks", 0, is_method=False):
|
||||
with patch(op, "internal_input_queue_num_blocks", 0):
|
||||
with patch(op, "_inputs_complete", True, is_method=False):
|
||||
assert_autoscaling_action(
|
||||
delta=-1,
|
||||
force=True,
|
||||
expected_reason="consumed all inputs",
|
||||
)
|
||||
|
||||
# With no enqueued inputs but inputs not being complete still,
|
||||
# the autoscaler should still scale up based on utilization
|
||||
assert_autoscaling_action(
|
||||
delta=5,
|
||||
expected_reason="utilization of 1.5 >= 1.0",
|
||||
)
|
||||
|
||||
# Should be no-op since the op doesn't have enough resources.
|
||||
with patch(
|
||||
op_state._scheduling_status,
|
||||
"under_resource_limits",
|
||||
False,
|
||||
is_method=False,
|
||||
):
|
||||
assert_autoscaling_action(
|
||||
delta=0,
|
||||
expected_reason="operator exceeding resource quota",
|
||||
)
|
||||
|
||||
# Should be a no-op since the op has enough available concurrency slots for
|
||||
# the existing inputs.
|
||||
with patch(actor_pool, "num_tasks_in_flight", 7):
|
||||
assert_autoscaling_action(
|
||||
delta=0,
|
||||
expected_reason="utilization of 0.7 w/in limits [0.5, 1.0]",
|
||||
)
|
||||
|
||||
# Should scale down since the util is below the threshold.
|
||||
with patch(actor_pool, "num_tasks_in_flight", 4):
|
||||
assert actor_pool.get_pool_util() == 0.4
|
||||
assert_autoscaling_action(
|
||||
delta=-1,
|
||||
expected_reason="utilization of 0.4 <= 0.5",
|
||||
)
|
||||
|
||||
# Should scale down since the pool is above the max size.
|
||||
with patch(actor_pool, "current_size", 16):
|
||||
assert_autoscaling_action(
|
||||
delta=-1,
|
||||
expected_reason="pool exceeding max size",
|
||||
)
|
||||
|
||||
# Should no-op because the op has no budget.
|
||||
with patch(resource_manager, "get_budget", ExecutionResources.zero()):
|
||||
assert_autoscaling_action(
|
||||
delta=0,
|
||||
expected_reason="exceeded resource limits",
|
||||
)
|
||||
|
||||
# Should no-op because the op has not received any inputs.
|
||||
with patch(op.metrics, "num_inputs_received", 0, is_method=False):
|
||||
assert_autoscaling_action(
|
||||
delta=0,
|
||||
expected_reason="no inputs received",
|
||||
)
|
||||
|
||||
# --- Resource budget enforcement (downscaling) ---
|
||||
# get_allocation and get_op_usage are patched to simulate an operator that
|
||||
# has exceeded its total resource allocation. The over-budget check fires
|
||||
# before utilization logic, so even high utilization (1.5x) is overridden.
|
||||
|
||||
# CPU over-budget by 2 actors: allocation=8 CPUs, usage=10 CPUs, 1 CPU/actor.
|
||||
# allocation - usage = -2 → scale down by ceil(2/1) = 2.
|
||||
with patch(resource_manager, "get_allocation", ExecutionResources(cpu=8)):
|
||||
with patch(resource_manager, "get_op_usage", ExecutionResources(cpu=10)):
|
||||
assert_autoscaling_action(
|
||||
delta=-2,
|
||||
expected_reason="actor pool exceeds resource allocation",
|
||||
)
|
||||
|
||||
# Over-budget but current_size=6 (min_size+1): required=2 but can only
|
||||
# release 1 actor (max_can_release = 6 - 5 = 1).
|
||||
with patch(resource_manager, "get_allocation", ExecutionResources(cpu=8)):
|
||||
with patch(resource_manager, "get_op_usage", ExecutionResources(cpu=10)):
|
||||
with patch(actor_pool, "current_size", 6):
|
||||
assert_autoscaling_action(
|
||||
delta=-1,
|
||||
expected_reason="actor pool exceeds resource allocation",
|
||||
)
|
||||
|
||||
# Over-budget but pool is at min_size (current=5): cannot release any actors.
|
||||
with patch(resource_manager, "get_allocation", ExecutionResources(cpu=8)):
|
||||
with patch(resource_manager, "get_op_usage", ExecutionResources(cpu=10)):
|
||||
with patch(actor_pool, "current_size", 5):
|
||||
assert_autoscaling_action(
|
||||
delta=0,
|
||||
expected_reason="actor pool exceeds resource allocation "
|
||||
"but cannot scale below min size",
|
||||
)
|
||||
|
||||
# GPU pool: allocation=3 GPUs, usage=6 GPUs, 1 GPU/actor.
|
||||
# allocation - usage = -3 → scale down by 3.
|
||||
with patch(actor_pool, "per_actor_resource_usage", ExecutionResources(gpu=1)):
|
||||
with patch(resource_manager, "get_allocation", ExecutionResources(gpu=3)):
|
||||
with patch(resource_manager, "get_op_usage", ExecutionResources(gpu=6)):
|
||||
assert_autoscaling_action(
|
||||
delta=-3,
|
||||
expected_reason="actor pool exceeds resource allocation",
|
||||
)
|
||||
|
||||
# Cross-resource: GPU-only pool (per_actor.cpu=0) with negative CPU budget
|
||||
# but positive GPU budget. CPU over-budget doesn't trigger since the pool
|
||||
# doesn't consume CPU. GPU headroom = floor(5/1)=5, capped by
|
||||
# max_size(15)-current_size(10)=5.
|
||||
with patch(actor_pool, "per_actor_resource_usage", ExecutionResources(gpu=1)):
|
||||
with patch(
|
||||
resource_manager, "get_allocation", ExecutionResources(cpu=8, gpu=10)
|
||||
):
|
||||
with patch(
|
||||
resource_manager, "get_op_usage", ExecutionResources(cpu=10, gpu=5)
|
||||
):
|
||||
assert_autoscaling_action(
|
||||
delta=5,
|
||||
expected_reason="utilization of 1.5 >= 1.0",
|
||||
)
|
||||
|
||||
# Memory bottleneck: allocation=4 GB, usage=5 GB, 500 MB/actor.
|
||||
# allocation - usage = -1 GB → ceil(1 GB / 500 MB) = 2 actors to remove.
|
||||
# CPU is within budget (allocation.cpu > usage.cpu), so CPU does not trigger.
|
||||
with patch(
|
||||
actor_pool,
|
||||
"per_actor_resource_usage",
|
||||
ExecutionResources(cpu=1, memory=500_000_000),
|
||||
):
|
||||
with patch(
|
||||
resource_manager,
|
||||
"get_allocation",
|
||||
ExecutionResources(cpu=15, memory=4_000_000_000),
|
||||
):
|
||||
with patch(
|
||||
resource_manager,
|
||||
"get_op_usage",
|
||||
ExecutionResources(cpu=10, memory=5_000_000_000),
|
||||
):
|
||||
assert_autoscaling_action(
|
||||
delta=-2,
|
||||
expected_reason="actor pool exceeds resource allocation",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def autoscaler_max_upscaling_delta_setup():
|
||||
resource_manager = MagicMock(
|
||||
spec=ResourceManager,
|
||||
get_budget=MagicMock(return_value=None),
|
||||
get_allocation=MagicMock(return_value=None),
|
||||
)
|
||||
|
||||
actor_pool = MagicMock(
|
||||
spec=_ActorPool,
|
||||
min_size=MagicMock(return_value=5),
|
||||
max_size=MagicMock(return_value=20),
|
||||
current_size=MagicMock(return_value=10),
|
||||
get_current_size=MagicMock(return_value=10),
|
||||
num_pending_actors=MagicMock(return_value=0),
|
||||
num_tasks_in_flight=MagicMock(return_value=40),
|
||||
max_tasks_in_flight_per_actor=MagicMock(return_value=4),
|
||||
get_pool_util=MagicMock(return_value=2.0),
|
||||
)
|
||||
|
||||
op = MagicMock(
|
||||
spec=InternalQueueOperatorMixin,
|
||||
has_completed=MagicMock(return_value=False),
|
||||
_inputs_complete=False,
|
||||
metrics=MagicMock(average_num_inputs_per_task=1, num_inputs_received=1),
|
||||
)
|
||||
op_state = MagicMock(
|
||||
spec=OpState,
|
||||
total_enqueued_input_blocks=MagicMock(return_value=1),
|
||||
)
|
||||
op_state.op = op
|
||||
op_state._scheduling_status = MagicMock(under_resource_limits=True)
|
||||
return resource_manager, actor_pool, op, op_state
|
||||
|
||||
|
||||
def test_actor_pool_scaling_respects_small_max_upscaling_delta(
|
||||
autoscaler_max_upscaling_delta_setup,
|
||||
):
|
||||
resource_manager, actor_pool, op, op_state = autoscaler_max_upscaling_delta_setup
|
||||
autoscaler = DefaultActorAutoscaler(
|
||||
topology=MagicMock(),
|
||||
resource_manager=resource_manager,
|
||||
config=AutoscalingConfig(
|
||||
actor_pool_util_upscaling_threshold=1.0,
|
||||
actor_pool_util_downscaling_threshold=0.5,
|
||||
actor_pool_max_upscaling_delta=3,
|
||||
),
|
||||
)
|
||||
request = autoscaler._derive_target_scaling_config(
|
||||
actor_pool=actor_pool,
|
||||
op=op,
|
||||
op_state=op_state,
|
||||
)
|
||||
# With current_size=10, util=2.0, threshold=1.0:
|
||||
# plan_delta = ceil(10 * (2.0/1.0 - 1)) = ceil(10) = 10
|
||||
# However, delta is limited by max_upscaling_delta=3, so delta = min(10, 3) = 3
|
||||
assert request.delta == 3
|
||||
|
||||
|
||||
def test_actor_pool_scaling_respects_large_max_upscaling_delta(
|
||||
autoscaler_max_upscaling_delta_setup,
|
||||
):
|
||||
resource_manager, actor_pool, op, op_state = autoscaler_max_upscaling_delta_setup
|
||||
autoscaler = DefaultActorAutoscaler(
|
||||
topology=MagicMock(),
|
||||
resource_manager=resource_manager,
|
||||
config=AutoscalingConfig(
|
||||
actor_pool_util_upscaling_threshold=1.0,
|
||||
actor_pool_util_downscaling_threshold=0.5,
|
||||
actor_pool_max_upscaling_delta=100,
|
||||
),
|
||||
)
|
||||
request = autoscaler._derive_target_scaling_config(
|
||||
actor_pool=actor_pool,
|
||||
op=op,
|
||||
op_state=op_state,
|
||||
)
|
||||
# With current_size=10, util=2.0, threshold=1.0:
|
||||
# plan_delta = ceil(10 * (2.0/1.0 - 1)) = ceil(10) = 10
|
||||
# max_upscaling_delta=100 is large enough, but delta is limited by max_size:
|
||||
# max_size(20) - current_size(10) = 10, so delta = min(10, 100, 10) = 10
|
||||
assert request.delta == 10
|
||||
|
||||
|
||||
class BarrierWaiter:
|
||||
def __init__(self, barrier):
|
||||
self._barrier = barrier
|
||||
|
||||
def __call__(self, x):
|
||||
ray.get(self._barrier.wait.remote(), timeout=10)
|
||||
return x
|
||||
|
||||
|
||||
@ray.remote(max_concurrency=10)
|
||||
class Barrier:
|
||||
def __init__(self, n, delay=0):
|
||||
self.n = n
|
||||
self.delay = delay
|
||||
self.max_waiters = 0
|
||||
self.cur_waiters = 0
|
||||
|
||||
def wait(self):
|
||||
self.cur_waiters += 1
|
||||
if self.cur_waiters > self.max_waiters:
|
||||
self.max_waiters = self.cur_waiters
|
||||
self.n -= 1
|
||||
print("wait", self.n)
|
||||
while self.n > 0:
|
||||
time.sleep(0.1)
|
||||
time.sleep(self.delay)
|
||||
print("wait done")
|
||||
self.cur_waiters -= 1
|
||||
|
||||
def get_max_waiters(self):
|
||||
return self.max_waiters
|
||||
|
||||
|
||||
def test_actor_pool_scales_up(ray_start_10_cpus_shared, restore_data_context):
|
||||
# The Ray cluster started by the fixture might not have much object store memory.
|
||||
# To prevent the actor pool from getting backpressured, we decrease the max block
|
||||
# size.
|
||||
ctx = ray.data.DataContext.get_current()
|
||||
ctx.target_max_block_size = 1 * 1024**2
|
||||
|
||||
# The `BarrierWaiter` UDF blocks until there are 2 actors running. If we don't
|
||||
# scale up, the UDF raises a timeout.
|
||||
barrier = Barrier.remote(2)
|
||||
# We produce 3 blocks (1 elem each) such that
|
||||
# - We start wiht actor pool of min_size
|
||||
# - 2 tasks could be submitted to an actor (utilization reaches 200%)
|
||||
# - Autoscaler kicks in and creates another actor
|
||||
# - 3 task is submitted to a new actor (unblocking the barrier)
|
||||
ray.data.range(3, override_num_blocks=3).map(
|
||||
BarrierWaiter,
|
||||
fn_constructor_args=(barrier,),
|
||||
compute=ray.data.ActorPoolStrategy(
|
||||
min_size=1, max_size=2, max_tasks_in_flight_per_actor=2
|
||||
),
|
||||
).take_all()
|
||||
|
||||
|
||||
def test_actor_pool_respects_max_size(ray_start_10_cpus_shared, restore_data_context):
|
||||
# The Ray cluster started by the fixture might not have much object store memory.
|
||||
# To prevent the actor pool from getting backpressured, we decrease the max block
|
||||
# size.
|
||||
ctx = ray.data.DataContext.get_current()
|
||||
ctx.target_max_block_size = 1 * 1024**2
|
||||
|
||||
# The `BarrierWaiter` UDF blocks until there are 3 actors running. Since the max
|
||||
# pool size is 2, the UDF should eventually timeout.
|
||||
barrier = Barrier.remote(3)
|
||||
with pytest.raises(ray.exceptions.RayTaskError):
|
||||
ray.data.range(2, override_num_blocks=2).map(
|
||||
BarrierWaiter,
|
||||
fn_constructor_args=(barrier,),
|
||||
compute=ray.data.ActorPoolStrategy(min_size=1, max_size=2),
|
||||
).take_all()
|
||||
|
||||
|
||||
def test_autoscaling_config_validation_warnings(
|
||||
ray_start_10_cpus_shared, restore_data_context
|
||||
):
|
||||
"""Test that validation warnings are emitted when actor pool config won't allow scaling up."""
|
||||
|
||||
class SimpleMapper:
|
||||
"""Simple callable class for testing autoscaling validation."""
|
||||
|
||||
def __call__(self, row):
|
||||
# Map operates on rows which are dicts
|
||||
return {"value": row["id"] * 2}
|
||||
|
||||
# Test #1: Invalid config (should warn)
|
||||
# - max_tasks_in_flight / max_concurrency == 1
|
||||
# - Default upscaling threshold (200%)
|
||||
with patch(
|
||||
"ray.data._internal.actor_autoscaler.default_actor_autoscaler.logger.warning"
|
||||
) as mock_warning:
|
||||
ds = ray.data.range(2, override_num_blocks=2).map_batches(
|
||||
SimpleMapper,
|
||||
compute=ray.data.ActorPoolStrategy(
|
||||
max_tasks_in_flight_per_actor=1,
|
||||
),
|
||||
max_concurrency=1,
|
||||
)
|
||||
# Take just one item to minimize execution time
|
||||
ds.take_all()
|
||||
|
||||
# Check that warning was called with expected message
|
||||
warn_log_args_str = str(mock_warning.call_args_list)
|
||||
expected_message = (
|
||||
"⚠️ Actor Pool configuration of the "
|
||||
"ActorPoolMapOperator[MapBatches(SimpleMapper)] will not allow it to scale up: "
|
||||
"configured utilization threshold (175.0%) couldn't be reached with "
|
||||
"configured max_concurrency=1 and max_tasks_in_flight_per_actor=1 "
|
||||
"(max utilization will be max_tasks_in_flight_per_actor / max_concurrency = 100%)"
|
||||
)
|
||||
|
||||
assert expected_message in warn_log_args_str
|
||||
|
||||
# Test #2: Provided config is valid (no warnings)
|
||||
# - max_tasks_in_flight / max_concurrency == 2 (default)
|
||||
# - Default upscaling threshold (200%)
|
||||
with patch(
|
||||
"ray.data._internal.actor_autoscaler.default_actor_autoscaler.logger.warning"
|
||||
) as mock_warning:
|
||||
ds = ray.data.range(2, override_num_blocks=2).map_batches(
|
||||
SimpleMapper,
|
||||
compute=ray.data.ActorPoolStrategy(
|
||||
max_tasks_in_flight_per_actor=2,
|
||||
),
|
||||
max_concurrency=1,
|
||||
)
|
||||
ds.take_all()
|
||||
|
||||
# Check that this warning hasn't been emitted
|
||||
warn_log_args_str = str(mock_warning.call_args_list)
|
||||
expected_message = (
|
||||
"⚠️ Actor Pool configuration of the "
|
||||
"ActorPoolMapOperator[MapBatches(SimpleMapper)] will not allow it to scale up: "
|
||||
)
|
||||
|
||||
assert expected_message not in warn_log_args_str
|
||||
|
||||
# Test #3: Default config is valid (no warnings)
|
||||
# - max_tasks_in_flight / max_concurrency == 4 (default)
|
||||
# - Default upscaling threshold (200%)
|
||||
with patch(
|
||||
"ray.data._internal.actor_autoscaler.default_actor_autoscaler.logger.warning"
|
||||
) as mock_warning:
|
||||
ds = ray.data.range(2, override_num_blocks=2).map_batches(
|
||||
SimpleMapper, compute=ray.data.ActorPoolStrategy()
|
||||
)
|
||||
ds.take_all()
|
||||
|
||||
# Check that this warning hasn't been emitted
|
||||
warn_log_args_str = str(mock_warning.call_args_list)
|
||||
expected_message = (
|
||||
"⚠️ Actor Pool configuration of the "
|
||||
"ActorPoolMapOperator[MapBatches(SimpleMapper)] will not allow it to scale up: "
|
||||
)
|
||||
|
||||
assert expected_message not in warn_log_args_str
|
||||
|
||||
# Test #4: Fixed-size pool with invalid config (no warnings)
|
||||
# - max_tasks_in_flight / max_concurrency == 1
|
||||
# - Default upscaling threshold (200%)
|
||||
# - Even though config would normally trigger warning, fixed-size pools
|
||||
# don't scale up by design, so warning should not be emitted
|
||||
with patch(
|
||||
"ray.data._internal.actor_autoscaler.default_actor_autoscaler.logger.warning"
|
||||
) as mock_warning:
|
||||
ds = ray.data.range(2, override_num_blocks=2).map_batches(
|
||||
SimpleMapper,
|
||||
compute=ray.data.ActorPoolStrategy(
|
||||
size=2,
|
||||
max_tasks_in_flight_per_actor=1,
|
||||
),
|
||||
max_concurrency=1,
|
||||
)
|
||||
ds.take_all()
|
||||
|
||||
# Check that this warning hasn't been emitted for fixed-size pool
|
||||
warn_log_args_str = str(mock_warning.call_args_list)
|
||||
expected_message = (
|
||||
"⚠️ Actor Pool configuration of the "
|
||||
"ActorPoolMapOperator[MapBatches(SimpleMapper)] will not allow it to scale up: "
|
||||
)
|
||||
|
||||
assert expected_message not in warn_log_args_str
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def autoscaler_config_mocks():
|
||||
resource_manager = MagicMock(spec=ResourceManager)
|
||||
topology = MagicMock()
|
||||
topology.items = MagicMock(return_value=[])
|
||||
return resource_manager, topology
|
||||
|
||||
|
||||
def test_autoscaling_config_validation_zero_delta(autoscaler_config_mocks):
|
||||
resource_manager, topology = autoscaler_config_mocks
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="actor_pool_max_upscaling_delta must be positive"
|
||||
):
|
||||
DefaultActorAutoscaler(
|
||||
topology=topology,
|
||||
resource_manager=resource_manager,
|
||||
config=AutoscalingConfig(
|
||||
actor_pool_util_upscaling_threshold=1.0,
|
||||
actor_pool_util_downscaling_threshold=0.5,
|
||||
actor_pool_max_upscaling_delta=0,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_autoscaling_config_validation_negative_delta(autoscaler_config_mocks):
|
||||
resource_manager, topology = autoscaler_config_mocks
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="actor_pool_max_upscaling_delta must be positive"
|
||||
):
|
||||
DefaultActorAutoscaler(
|
||||
topology=topology,
|
||||
resource_manager=resource_manager,
|
||||
config=AutoscalingConfig(
|
||||
actor_pool_util_upscaling_threshold=1.0,
|
||||
actor_pool_util_downscaling_threshold=0.5,
|
||||
actor_pool_max_upscaling_delta=-1,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_autoscaling_config_validation_positive_delta(autoscaler_config_mocks):
|
||||
resource_manager, topology = autoscaler_config_mocks
|
||||
|
||||
autoscaler = DefaultActorAutoscaler(
|
||||
topology=topology,
|
||||
resource_manager=resource_manager,
|
||||
config=AutoscalingConfig(
|
||||
actor_pool_util_upscaling_threshold=1.0,
|
||||
actor_pool_util_downscaling_threshold=0.5,
|
||||
actor_pool_max_upscaling_delta=5,
|
||||
),
|
||||
)
|
||||
assert autoscaler._actor_pool_max_upscaling_delta == 5
|
||||
|
||||
|
||||
def test_autoscaling_config_validation_zero_upscaling_threshold(
|
||||
autoscaler_config_mocks,
|
||||
):
|
||||
resource_manager, topology = autoscaler_config_mocks
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="actor_pool_util_upscaling_threshold must be positive"
|
||||
):
|
||||
DefaultActorAutoscaler(
|
||||
topology=topology,
|
||||
resource_manager=resource_manager,
|
||||
config=AutoscalingConfig(
|
||||
actor_pool_util_upscaling_threshold=0,
|
||||
actor_pool_util_downscaling_threshold=0.5,
|
||||
actor_pool_max_upscaling_delta=5,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_autoscaling_config_validation_negative_upscaling_threshold(
|
||||
autoscaler_config_mocks,
|
||||
):
|
||||
resource_manager, topology = autoscaler_config_mocks
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="actor_pool_util_upscaling_threshold must be positive"
|
||||
):
|
||||
DefaultActorAutoscaler(
|
||||
topology=topology,
|
||||
resource_manager=resource_manager,
|
||||
config=AutoscalingConfig(
|
||||
actor_pool_util_upscaling_threshold=-1.0,
|
||||
actor_pool_util_downscaling_threshold=0.5,
|
||||
actor_pool_max_upscaling_delta=5,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def test_get_max_scale_up_tolerates_float_drift():
|
||||
"""Regression test for #64291.
|
||||
|
||||
A budget can carry tiny float drift (e.g. ``gpu=-1e-16``) from chained
|
||||
arithmetic. ``_get_max_scale_up`` reads raw fields (non-negativity assert +
|
||||
``floordiv``), so this must not trip the assert or yield a negative
|
||||
scale-up. ``ExecutionResources`` rounds at construction, so the drift
|
||||
collapses to 0.
|
||||
"""
|
||||
actor_pool = MagicMock()
|
||||
actor_pool.per_actor_resource_usage = MagicMock(
|
||||
return_value=ExecutionResources(cpu=1.0, gpu=0.25, memory=0.0)
|
||||
)
|
||||
# gpu drift rounds to 0 -> 0 actors fit on the gpu dimension -> scale-up 0.
|
||||
budget = ExecutionResources(cpu=4, gpu=-1e-16, memory=0.0)
|
||||
assert _get_max_scale_up(actor_pool, budget) == 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,76 @@
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data.aggregate import (
|
||||
AggregateFn,
|
||||
Max,
|
||||
)
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
RANDOM_SEED = 123
|
||||
|
||||
|
||||
@pytest.mark.parametrize("keys", ["A", ["A", "B"]])
|
||||
def test_agg_inputs(
|
||||
ray_start_regular_shared_2_cpus,
|
||||
keys,
|
||||
configure_shuffle_method,
|
||||
disable_fallback_to_object_extension,
|
||||
):
|
||||
xs = list(range(100))
|
||||
ds = ray.data.from_items([{"A": (x % 3), "B": x, "C": (x % 2)} for x in xs])
|
||||
|
||||
def check_init(k):
|
||||
if len(keys) == 2:
|
||||
assert isinstance(k, tuple), k
|
||||
assert len(k) == 2
|
||||
elif len(keys) == 1:
|
||||
assert isinstance(k, int)
|
||||
return 1
|
||||
|
||||
def check_finalize(v):
|
||||
assert v == 1
|
||||
|
||||
def check_accumulate_merge(a, r):
|
||||
assert a == 1
|
||||
if isinstance(r, int):
|
||||
return 1
|
||||
elif len(r) == 3:
|
||||
assert all(x in r for x in ["A", "B", "C"])
|
||||
else:
|
||||
assert False, r
|
||||
return 1
|
||||
|
||||
output = ds.groupby(keys).aggregate(
|
||||
AggregateFn(
|
||||
init=check_init,
|
||||
accumulate_row=check_accumulate_merge,
|
||||
merge=check_accumulate_merge,
|
||||
finalize=check_finalize,
|
||||
name="foo",
|
||||
)
|
||||
)
|
||||
output.take_all()
|
||||
|
||||
|
||||
def test_agg_errors(
|
||||
ray_start_regular_shared_2_cpus,
|
||||
configure_shuffle_method,
|
||||
disable_fallback_to_object_extension,
|
||||
):
|
||||
|
||||
ds = ray.data.range(100)
|
||||
ds.aggregate(Max("id")) # OK
|
||||
with pytest.raises(ValueError):
|
||||
ds.aggregate(Max())
|
||||
with pytest.raises(ValueError):
|
||||
ds.aggregate(Max(lambda x: x))
|
||||
with pytest.raises(ValueError):
|
||||
ds.aggregate(Max("bad_field"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,565 @@
|
||||
from collections import Counter
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data.aggregate import (
|
||||
ApproximateQuantile,
|
||||
ApproximateTopK,
|
||||
MissingValuePercentage,
|
||||
Unique,
|
||||
ZeroPercentage,
|
||||
)
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
|
||||
class TestMissingValuePercentage:
|
||||
"""Test cases for MissingValuePercentage aggregation."""
|
||||
|
||||
def test_missing_value_percentage_basic(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test basic missing value percentage calculation."""
|
||||
# Create test data with some null values
|
||||
data = [
|
||||
{"id": 1, "value": 10},
|
||||
{"id": 2, "value": None},
|
||||
{"id": 3, "value": 30},
|
||||
{"id": 4, "value": None},
|
||||
{"id": 5, "value": 50},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = ds.aggregate(MissingValuePercentage(on="value"))
|
||||
expected = 40.0 # 2 nulls out of 5 total = 40%
|
||||
|
||||
assert result["missing_pct(value)"] == expected
|
||||
|
||||
def test_missing_value_percentage_no_nulls(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test missing value percentage with no null values."""
|
||||
data = [
|
||||
{"id": 1, "value": 10},
|
||||
{"id": 2, "value": 20},
|
||||
{"id": 3, "value": 30},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = ds.aggregate(MissingValuePercentage(on="value"))
|
||||
expected = 0.0 # 0 nulls out of 3 total = 0%
|
||||
|
||||
assert result["missing_pct(value)"] == expected
|
||||
|
||||
def test_missing_value_percentage_all_nulls(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test missing value percentage with all null values."""
|
||||
data = [
|
||||
{"id": 1, "value": None},
|
||||
{"id": 2, "value": None},
|
||||
{"id": 3, "value": None},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = ds.aggregate(MissingValuePercentage(on="value"))
|
||||
expected = 100.0 # 3 nulls out of 3 total = 100%
|
||||
|
||||
assert result["missing_pct(value)"] == expected
|
||||
|
||||
def test_missing_value_percentage_with_nan(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test missing value percentage with NaN values."""
|
||||
data = [
|
||||
{"id": 1, "value": 10.0},
|
||||
{"id": 2, "value": np.nan},
|
||||
{"id": 3, "value": None},
|
||||
{"id": 4, "value": 40.0},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = ds.aggregate(MissingValuePercentage(on="value"))
|
||||
expected = 50.0 # 2 nulls (NaN + None) out of 4 total = 50%
|
||||
|
||||
assert result["missing_pct(value)"] == expected
|
||||
|
||||
def test_missing_value_percentage_with_string(
|
||||
self, ray_start_regular_shared_2_cpus
|
||||
):
|
||||
"""Test missing value percentage with string values."""
|
||||
data = [
|
||||
{"id": 1, "value": "a"},
|
||||
{"id": 2, "value": None},
|
||||
{"id": 3, "value": None},
|
||||
{"id": 4, "value": "b"},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = ds.aggregate(MissingValuePercentage(on="value"))
|
||||
expected = 50.0 # 2 None out of 4 total = 50%
|
||||
|
||||
assert result["missing_pct(value)"] == expected
|
||||
|
||||
def test_missing_value_percentage_custom_alias(
|
||||
self, ray_start_regular_shared_2_cpus
|
||||
):
|
||||
"""Test missing value percentage with custom alias name."""
|
||||
data = [
|
||||
{"id": 1, "value": 10},
|
||||
{"id": 2, "value": None},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = ds.aggregate(MissingValuePercentage(on="value", alias_name="null_pct"))
|
||||
expected = 50.0 # 1 null out of 2 total = 50%
|
||||
|
||||
assert result["null_pct"] == expected
|
||||
|
||||
def test_missing_value_percentage_large_dataset(
|
||||
self, ray_start_regular_shared_2_cpus
|
||||
):
|
||||
"""Test missing value percentage with larger dataset."""
|
||||
# Create a larger dataset with known null percentage
|
||||
data = []
|
||||
for i in range(1000):
|
||||
value = None if i % 10 == 0 else i # 10% null values
|
||||
data.append({"id": i, "value": value})
|
||||
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = ds.aggregate(MissingValuePercentage(on="value"))
|
||||
expected = 10.0 # 100 nulls out of 1000 total = 10%
|
||||
|
||||
assert abs(result["missing_pct(value)"] - expected) < 0.01
|
||||
|
||||
|
||||
class TestZeroPercentage:
|
||||
"""Test cases for ZeroPercentage aggregation."""
|
||||
|
||||
def test_zero_percentage_basic(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test basic zero percentage calculation."""
|
||||
data = [
|
||||
{"id": 1, "value": 10},
|
||||
{"id": 2, "value": 0},
|
||||
{"id": 3, "value": 30},
|
||||
{"id": 4, "value": 0},
|
||||
{"id": 5, "value": 50},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = ds.aggregate(ZeroPercentage(on="value"))
|
||||
expected = 40.0 # 2 zeros out of 5 total = 40%
|
||||
|
||||
assert result["zero_pct(value)"] == expected
|
||||
|
||||
def test_zero_percentage_no_zeros(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test zero percentage with no zero values."""
|
||||
data = [
|
||||
{"id": 1, "value": 10},
|
||||
{"id": 2, "value": 20},
|
||||
{"id": 3, "value": 30},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = ds.aggregate(ZeroPercentage(on="value"))
|
||||
expected = 0.0 # 0 zeros out of 3 total = 0%
|
||||
|
||||
assert result["zero_pct(value)"] == expected
|
||||
|
||||
def test_zero_percentage_all_zeros(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test zero percentage with all zero values."""
|
||||
data = [
|
||||
{"id": 1, "value": 0},
|
||||
{"id": 2, "value": 0},
|
||||
{"id": 3, "value": 0},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = ds.aggregate(ZeroPercentage(on="value"))
|
||||
expected = 100.0 # 3 zeros out of 3 total = 100%
|
||||
|
||||
assert result["zero_pct(value)"] == expected
|
||||
|
||||
def test_zero_percentage_with_nulls_ignore_nulls_true(
|
||||
self, ray_start_regular_shared_2_cpus
|
||||
):
|
||||
"""Test zero percentage with null values when ignore_nulls=True."""
|
||||
data = [
|
||||
{"id": 1, "value": 10},
|
||||
{"id": 2, "value": 0},
|
||||
{"id": 3, "value": None},
|
||||
{"id": 4, "value": 0},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = ds.aggregate(ZeroPercentage(on="value", ignore_nulls=True))
|
||||
expected = 66.67 # 2 zeros out of 3 non-null values ≈ 66.67%
|
||||
|
||||
assert abs(result["zero_pct(value)"] - expected) < 0.01
|
||||
|
||||
def test_zero_percentage_with_nulls_ignore_nulls_false(
|
||||
self, ray_start_regular_shared_2_cpus
|
||||
):
|
||||
"""Test zero percentage with null values when ignore_nulls=False."""
|
||||
data = [
|
||||
{"id": 1, "value": 10},
|
||||
{"id": 2, "value": 0},
|
||||
{"id": 3, "value": None},
|
||||
{"id": 4, "value": 0},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = ds.aggregate(ZeroPercentage(on="value", ignore_nulls=False))
|
||||
expected = 50.0 # 2 zeros out of 4 total values = 50%
|
||||
|
||||
assert result["zero_pct(value)"] == expected
|
||||
|
||||
def test_zero_percentage_all_nulls(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test zero percentage with all null values."""
|
||||
data = [
|
||||
{"id": 1, "value": None},
|
||||
{"id": 2, "value": None},
|
||||
{"id": 3, "value": None},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = ds.aggregate(ZeroPercentage(on="value", ignore_nulls=True))
|
||||
expected = None # No non-null values to calculate percentage
|
||||
|
||||
assert result["zero_pct(value)"] == expected
|
||||
|
||||
def test_zero_percentage_custom_alias(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test zero percentage with custom alias name."""
|
||||
data = [
|
||||
{"id": 1, "value": 10},
|
||||
{"id": 2, "value": 0},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = ds.aggregate(ZeroPercentage(on="value", alias_name="zero_ratio"))
|
||||
expected = 50.0 # 1 zero out of 2 total = 50%
|
||||
|
||||
assert result["zero_ratio"] == expected
|
||||
|
||||
def test_zero_percentage_large_dataset(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test zero percentage with larger dataset."""
|
||||
# Create a larger dataset with known zero percentage
|
||||
data = []
|
||||
for i in range(1000):
|
||||
value = 0 if i % 5 == 0 else i # 20% zero values
|
||||
data.append({"id": i, "value": value})
|
||||
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = ds.aggregate(ZeroPercentage(on="value"))
|
||||
expected = 20.0 # 200 zeros out of 1000 total = 20%
|
||||
|
||||
assert abs(result["zero_pct(value)"] - expected) < 0.01
|
||||
|
||||
def test_zero_percentage_float_zeros(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test zero percentage with float zero values."""
|
||||
data = [
|
||||
{"id": 1, "value": 10.5},
|
||||
{"id": 2, "value": 0.0},
|
||||
{"id": 3, "value": 30.7},
|
||||
{"id": 4, "value": 0.0},
|
||||
{"id": 5, "value": 50.2},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = ds.aggregate(ZeroPercentage(on="value"))
|
||||
expected = 40.0 # 2 zeros out of 5 total = 40%
|
||||
|
||||
assert result["zero_pct(value)"] == expected
|
||||
|
||||
def test_zero_percentage_negative_values(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test zero percentage with negative values (zeros should still be counted)."""
|
||||
data = [
|
||||
{"id": 1, "value": -10},
|
||||
{"id": 2, "value": 0},
|
||||
{"id": 3, "value": 30},
|
||||
{"id": 4, "value": -5},
|
||||
{"id": 5, "value": 0},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = ds.aggregate(ZeroPercentage(on="value"))
|
||||
expected = 40.0 # 2 zeros out of 5 total = 40%
|
||||
|
||||
assert result["zero_pct(value)"] == expected
|
||||
|
||||
|
||||
class TestApproximateQuantile:
|
||||
"""Test cases for ApproximateQuantile aggregation."""
|
||||
|
||||
def test_approximate_quantile_basic(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test basic approximate quantile calculation."""
|
||||
data = [
|
||||
{
|
||||
"id": 1,
|
||||
"value": 10,
|
||||
},
|
||||
{"id": 2, "value": 0},
|
||||
{"id": 3, "value": 30},
|
||||
{"id": 4, "value": 0},
|
||||
{"id": 5, "value": 50},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = ds.aggregate(
|
||||
ApproximateQuantile(on="value", quantiles=[0.1, 0.5, 0.9])
|
||||
)
|
||||
expected = [0.0, 10.0, 50.0]
|
||||
assert result["approx_quantile(value)"] == expected
|
||||
|
||||
def test_approximate_quantile_ignores_nulls(self, ray_start_regular_shared_2_cpus):
|
||||
data = [
|
||||
{"id": 1, "value": 5.0},
|
||||
{"id": 2, "value": None},
|
||||
{"id": 3, "value": 15.0},
|
||||
{"id": 4, "value": None},
|
||||
{"id": 5, "value": 25.0},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = ds.aggregate(ApproximateQuantile(on="value", quantiles=[0.5]))
|
||||
assert result["approx_quantile(value)"] == [15.0]
|
||||
|
||||
def test_approximate_quantile_custom_alias(self, ray_start_regular_shared_2_cpus):
|
||||
data = [
|
||||
{"id": 1, "value": 1.0},
|
||||
{"id": 2, "value": 3.0},
|
||||
{"id": 3, "value": 5.0},
|
||||
{"id": 4, "value": 7.0},
|
||||
{"id": 5, "value": 9.0},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
quantiles = [0.0, 1.0]
|
||||
result = ds.aggregate(
|
||||
ApproximateQuantile(
|
||||
on="value", quantiles=quantiles, alias_name="value_range"
|
||||
)
|
||||
)
|
||||
|
||||
assert result["value_range"] == [1.0, 9.0]
|
||||
assert len(result["value_range"]) == len(quantiles)
|
||||
|
||||
def test_approximate_quantile_groupby(self, ray_start_regular_shared_2_cpus):
|
||||
data = [
|
||||
{"group": "A", "value": 1.0},
|
||||
{"group": "A", "value": 2.0},
|
||||
{"group": "A", "value": 3.0},
|
||||
{"group": "B", "value": 10.0},
|
||||
{"group": "B", "value": 20.0},
|
||||
{"group": "B", "value": 30.0},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = (
|
||||
ds.groupby("group")
|
||||
.aggregate(ApproximateQuantile(on="value", quantiles=[0.5]))
|
||||
.take_all()
|
||||
)
|
||||
|
||||
result_by_group = {
|
||||
row["group"]: row["approx_quantile(value)"] for row in result
|
||||
}
|
||||
|
||||
assert result_by_group["A"] == [2.0]
|
||||
assert result_by_group["B"] == [20.0]
|
||||
|
||||
|
||||
class TestApproximateTopK:
|
||||
"""Test cases for ApproximateTopK aggregation."""
|
||||
|
||||
def test_approximate_topk_ignores_nulls(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test that null values are ignored."""
|
||||
data = [
|
||||
*[{"word": "apple"} for _ in range(5)],
|
||||
*[{"word": None} for _ in range(10)],
|
||||
*[{"word": "banana"} for _ in range(3)],
|
||||
*[{"word": "cherry"} for _ in range(2)],
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
result = ds.aggregate(ApproximateTopK(on="word", k=2))
|
||||
assert result["approx_topk(word)"] == [
|
||||
{"word": "apple", "count": 5},
|
||||
{"word": "banana", "count": 3},
|
||||
]
|
||||
|
||||
def test_approximate_topk_custom_alias(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test approximate top k with custom alias."""
|
||||
data = [
|
||||
*[{"item": "x"} for _ in range(3)],
|
||||
*[{"item": "y"} for _ in range(2)],
|
||||
*[{"item": "z"} for _ in range(1)],
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
result = ds.aggregate(ApproximateTopK(on="item", k=2, alias_name="top_items"))
|
||||
assert "top_items" in result
|
||||
assert result["top_items"] == [
|
||||
{"item": "x", "count": 3},
|
||||
{"item": "y", "count": 2},
|
||||
]
|
||||
|
||||
def test_approximate_topk_groupby(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test approximate top k with groupby."""
|
||||
data = [
|
||||
*[{"category": "A", "item": "apple"} for _ in range(5)],
|
||||
*[{"category": "A", "item": "banana"} for _ in range(3)],
|
||||
*[{"category": "B", "item": "cherry"} for _ in range(4)],
|
||||
*[{"category": "B", "item": "date"} for _ in range(2)],
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
result = (
|
||||
ds.groupby("category").aggregate(ApproximateTopK(on="item", k=1)).take_all()
|
||||
)
|
||||
|
||||
result_by_category = {
|
||||
row["category"]: row["approx_topk(item)"] for row in result
|
||||
}
|
||||
|
||||
assert result_by_category["A"] == [{"item": "apple", "count": 5}]
|
||||
assert result_by_category["B"] == [{"item": "cherry", "count": 4}]
|
||||
|
||||
def test_approximate_topk_all_unique(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test approximate top k when all items are unique."""
|
||||
data = [{"id": f"item_{i}"} for i in range(10)]
|
||||
ds = ray.data.from_items(data)
|
||||
result = ds.aggregate(ApproximateTopK(on="id", k=3))
|
||||
|
||||
# All items have count 1, so we should get exactly 3 items
|
||||
assert len(result["approx_topk(id)"]) == 3
|
||||
for item in result["approx_topk(id)"]:
|
||||
assert item["count"] == 1
|
||||
|
||||
def test_approximate_topk_fewer_items_than_k(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test approximate top k when dataset has fewer unique items than k."""
|
||||
data = [
|
||||
{"id": "a"},
|
||||
{"id": "b"},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
result = ds.aggregate(ApproximateTopK(on="id", k=5))
|
||||
|
||||
# Should only return 2 items since that's all we have
|
||||
assert len(result["approx_topk(id)"]) == 2
|
||||
|
||||
def test_approximate_topk_different_log_capacity(
|
||||
self, ray_start_regular_shared_2_cpus
|
||||
):
|
||||
"""Test that different log_capacity values still produce correct top k."""
|
||||
data = [
|
||||
*[{"id": "frequent"} for _ in range(100)],
|
||||
*[{"id": "common"} for _ in range(50)],
|
||||
*[{"id": f"rare_{i}"} for i in range(50)], # 50 unique rare items
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
# Test with smaller log_capacity
|
||||
result_small = ds.aggregate(ApproximateTopK(on="id", k=2, log_capacity=10))
|
||||
# Test with larger log_capacity
|
||||
result_large = ds.aggregate(ApproximateTopK(on="id", k=2, log_capacity=15))
|
||||
|
||||
# Both should correctly identify the top 2
|
||||
for result in [result_small, result_large]:
|
||||
assert result["approx_topk(id)"][0] == {"id": "frequent", "count": 100}
|
||||
assert result["approx_topk(id)"][1] == {"id": "common", "count": 50}
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("data", "expected1", "expected2"),
|
||||
[
|
||||
(
|
||||
[{"id": 1}, {"id": 1}, {"id": 2}],
|
||||
{"id": 1, "count": 2},
|
||||
{"id": 2, "count": 1},
|
||||
),
|
||||
(
|
||||
[{"id": [1, 2, 3]}, {"id": [1, 2, 3]}, {"id": [1, 2]}],
|
||||
{"id": [1, 2, 3], "count": 2},
|
||||
{"id": [1, 2], "count": 1},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_approximate_topk_non_string_datatype(
|
||||
self, data, expected1, expected2, ray_start_regular_shared_2_cpus
|
||||
):
|
||||
"""Test that ApproximateTopK works with non-string type elements."""
|
||||
ds = ray.data.from_items(data)
|
||||
|
||||
result = ds.aggregate(ApproximateTopK(on="id", k=2, log_capacity=3))
|
||||
assert result["approx_topk(id)"][0] == expected1
|
||||
assert result["approx_topk(id)"][1] == expected2
|
||||
|
||||
def test_approximate_topk_encode_lists(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test ApproximateTopK list encode feature."""
|
||||
data = [{"id": [1, 1, 1]}, {"id": [2, 2]}, {"id": [3]}]
|
||||
ds = ray.data.from_items(data)
|
||||
result = ds.aggregate(
|
||||
ApproximateTopK(on="id", k=4, log_capacity=10, encode_lists=True)
|
||||
)
|
||||
assert result["approx_topk(id)"][0] == {"id": 1, "count": 3}
|
||||
assert result["approx_topk(id)"][1] == {"id": 2, "count": 2}
|
||||
assert result["approx_topk(id)"][2] == {"id": 3, "count": 1}
|
||||
|
||||
|
||||
class TestUnique:
|
||||
"""Test cases for Unique aggregation."""
|
||||
|
||||
def test_unique_basic(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test basic Unique aggregation."""
|
||||
data = [{"id": "a"}, {"id": "b"}, {"id": "b"}, {"id": None}]
|
||||
ds = ray.data.from_items(data)
|
||||
result = ds.aggregate(Unique(on="id", ignore_nulls=False))
|
||||
|
||||
assert Counter(result["unique(id)"]) == Counter(["a", "b", None])
|
||||
|
||||
def test_unique_ignores_nulls(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test Unique properly ignores nulls."""
|
||||
data = [{"id": "a"}, {"id": None}, {"id": "b"}, {"id": "b"}, {"id": None}]
|
||||
ds = ray.data.from_items(data)
|
||||
result = ds.aggregate(Unique(on="id", ignore_nulls=True))
|
||||
|
||||
assert Counter(result["unique(id)"]) == Counter(["a", "b"])
|
||||
|
||||
def test_unique_custom_alias(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test Unique with custom alias."""
|
||||
data = [{"id": "a"}, {"id": "b"}, {"id": "b"}]
|
||||
ds = ray.data.from_items(data)
|
||||
result = ds.aggregate(Unique(on="id", alias_name="custom"))
|
||||
|
||||
assert sorted(result["custom"]) == ["a", "b"]
|
||||
|
||||
def test_unique_list_datatype(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test Unique works with non-hashable types like list."""
|
||||
data = [
|
||||
{"id": ["a", "b", "c"]},
|
||||
{"id": ["a", "b", "c"]},
|
||||
{"id": ["a", "b", "c"]},
|
||||
]
|
||||
ds = ray.data.from_items(data)
|
||||
result = ds.aggregate(Unique(on="id"))
|
||||
|
||||
assert result["unique(id)"][0] == ["a", "b", "c"]
|
||||
|
||||
def test_unique_encode_lists(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test Unique works when encode_lists is True."""
|
||||
data = [{"id": ["a", "b", "c"]}, {"id": ["a", "a", "a", "b", None]}]
|
||||
ds = ray.data.from_items(data)
|
||||
result = ds.aggregate(Unique(on="id", encode_lists=True, ignore_nulls=False))
|
||||
|
||||
answer = ["a", "b", "c", None]
|
||||
|
||||
assert Counter(result["unique(id)"]) == Counter(answer)
|
||||
|
||||
def test_unique_encode_lists_ignores_nulls(self, ray_start_regular_shared_2_cpus):
|
||||
"""Test Unique will drop null values when encode_lists is True."""
|
||||
data = [{"id": ["a", "b", "c"]}, {"id": ["a", "a", "a", "b", None]}]
|
||||
ds = ray.data.from_items(data)
|
||||
result = ds.aggregate(Unique(on="id", encode_lists=True, ignore_nulls=True))
|
||||
|
||||
answer = ["a", "b", "c"]
|
||||
|
||||
assert Counter(result["unique(id)"]) == Counter(answer)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,351 @@
|
||||
import base64
|
||||
import os
|
||||
import sys
|
||||
import types
|
||||
from decimal import Decimal
|
||||
from tempfile import TemporaryDirectory
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
from pyarrow import ArrowInvalid
|
||||
|
||||
import ray
|
||||
from ray._common.test_utils import run_string_as_driver
|
||||
from ray.data._internal.arrow_block import (
|
||||
ArrowBlockAccessor,
|
||||
ArrowBlockBuilder,
|
||||
)
|
||||
from ray.data._internal.arrow_ops.transform_pyarrow import combine_chunked_array
|
||||
from ray.data._internal.util import GiB, MiB
|
||||
from ray.data.block import BlockAccessor
|
||||
from ray.data.context import DataContext
|
||||
|
||||
|
||||
def test_combine_chunked_fixed_width_array_large():
|
||||
"""Verifies `combine_chunked_array` on fixed-width arrays > 2 GiB, produces
|
||||
single contiguous PA Array"""
|
||||
|
||||
# 144 MiB
|
||||
ones_1gb = np.ones(shape=(550, 128, 128, 4), dtype=np.int32()).ravel()
|
||||
|
||||
# Total ~2.15 GiB
|
||||
input_ = pa.chunked_array(
|
||||
[
|
||||
pa.array(ones_1gb),
|
||||
]
|
||||
* 16
|
||||
)
|
||||
|
||||
assert round(input_.nbytes / GiB, 2) == 2.15
|
||||
|
||||
result = combine_chunked_array(input_)
|
||||
|
||||
assert isinstance(result, pa.Int32Array)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"array_type,input_factory",
|
||||
[
|
||||
(
|
||||
pa.binary(),
|
||||
lambda num_bytes: np.arange(num_bytes, dtype=np.uint8).tobytes(),
|
||||
),
|
||||
(
|
||||
pa.string(),
|
||||
lambda num_bytes: base64.encodebytes(
|
||||
np.arange(num_bytes, dtype=np.int8).tobytes()
|
||||
).decode("ascii"),
|
||||
),
|
||||
(pa.list_(pa.uint8()), lambda num_bytes: np.arange(num_bytes, dtype=np.uint8)),
|
||||
],
|
||||
)
|
||||
def test_combine_chunked_variable_width_array_large(array_type, input_factory):
|
||||
"""Verifies `combine_chunked_array` on variable-width arrays > 2 GiB,
|
||||
safely produces new ChunkedArray with provided chunks recombined into
|
||||
larger ones up to INT32_MAX in size"""
|
||||
|
||||
one_half_gb_arr = pa.array([input_factory(GiB / 2)], type=array_type)
|
||||
chunked_arr = pa.chunked_array(
|
||||
[one_half_gb_arr, one_half_gb_arr, one_half_gb_arr, one_half_gb_arr]
|
||||
)
|
||||
|
||||
# 2 GiB + offsets (4 x int32)
|
||||
num_bytes = chunked_arr.nbytes
|
||||
expected_num_bytes = 4 * one_half_gb_arr.nbytes
|
||||
|
||||
num_chunks = len(chunked_arr.chunks)
|
||||
assert num_chunks == 4
|
||||
assert num_bytes == expected_num_bytes
|
||||
|
||||
# Assert attempt to combine directly fails
|
||||
with pytest.raises(ArrowInvalid):
|
||||
chunked_arr.combine_chunks()
|
||||
|
||||
# Safe combination succeeds by avoiding overflowing combination
|
||||
combined = combine_chunked_array(chunked_arr)
|
||||
|
||||
num_bytes = combined.nbytes
|
||||
|
||||
num_chunks = len(combined.chunks)
|
||||
assert num_chunks == 2
|
||||
assert num_bytes == expected_num_bytes
|
||||
|
||||
|
||||
def test_add_rows_with_different_column_names(ray_start_regular_shared):
|
||||
builder = ArrowBlockBuilder()
|
||||
|
||||
builder.add({"col1": "spam"})
|
||||
builder.add({"col2": "foo"})
|
||||
block = builder.build()
|
||||
|
||||
expected_table = pa.Table.from_pydict(
|
||||
{"col1": ["spam", None], "col2": [None, "foo"]}
|
||||
)
|
||||
assert block.equals(expected_table)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def binary_dataset_single_file_gt_2gb():
|
||||
total_size = int(2.1 * GiB)
|
||||
chunk_size = 256 * MiB
|
||||
num_chunks = total_size // chunk_size
|
||||
remainder = total_size % chunk_size
|
||||
|
||||
with TemporaryDirectory() as tmp_dir:
|
||||
dataset_path = f"{tmp_dir}/binary_dataset_gt_2gb_single_file"
|
||||
|
||||
# Create directory
|
||||
os.mkdir(dataset_path)
|
||||
|
||||
with open(f"{dataset_path}/chunk.bin", "wb") as f:
|
||||
for i in range(num_chunks):
|
||||
f.write(b"a" * chunk_size)
|
||||
|
||||
print(f">>> Written chunk #{i}")
|
||||
|
||||
if remainder:
|
||||
f.write(b"a" * remainder)
|
||||
|
||||
print(f">>> Wrote chunked dataset at: {dataset_path}")
|
||||
|
||||
yield dataset_path, total_size
|
||||
|
||||
print(f">>> Cleaning up dataset: {dataset_path}")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"col_name",
|
||||
[
|
||||
"bytes",
|
||||
# TODO fix numpy conversion
|
||||
# "text",
|
||||
],
|
||||
)
|
||||
def test_single_row_gt_2gb(
|
||||
ray_start_regular_shared,
|
||||
restore_data_context,
|
||||
binary_dataset_single_file_gt_2gb,
|
||||
col_name,
|
||||
):
|
||||
# Disable (automatic) fallback to `ArrowPythonObjectType` extension type
|
||||
DataContext.get_current().enable_fallback_to_arrow_object_ext_type = False
|
||||
|
||||
dataset_path, target_binary_size = binary_dataset_single_file_gt_2gb
|
||||
|
||||
def _id(row):
|
||||
bs = row[col_name]
|
||||
assert round(len(bs) / GiB, 1) == round(target_binary_size / GiB, 1)
|
||||
return row
|
||||
|
||||
if col_name == "text":
|
||||
ds = ray.data.read_text(dataset_path)
|
||||
elif col_name == "bytes":
|
||||
ds = ray.data.read_binary_files(dataset_path)
|
||||
|
||||
total = ds.map(_id).count()
|
||||
|
||||
assert total == 1
|
||||
|
||||
|
||||
def test_random_shuffle(ray_start_regular_shared):
|
||||
TOTAL_ROWS = 10000
|
||||
table = pa.table({"id": pa.array(range(TOTAL_ROWS))})
|
||||
block_accessor = ArrowBlockAccessor(table)
|
||||
|
||||
# Perform the random shuffle
|
||||
shuffled_table = block_accessor.random_shuffle(random_seed=None)
|
||||
assert shuffled_table.num_rows == TOTAL_ROWS
|
||||
|
||||
# Access the shuffled data
|
||||
block_accessor = ArrowBlockAccessor(shuffled_table)
|
||||
shuffled_data = block_accessor.to_pandas()["id"].tolist()
|
||||
original_data = list(range(TOTAL_ROWS))
|
||||
|
||||
# Ensure the shuffled data is not identical to the original
|
||||
assert (
|
||||
shuffled_data != original_data
|
||||
), "Shuffling should result in a different order"
|
||||
|
||||
# Ensure the entire set of original values is still in the shuffled dataset
|
||||
assert (
|
||||
sorted(shuffled_data) == original_data
|
||||
), "The shuffled data should contain all the original values"
|
||||
|
||||
|
||||
def test_register_arrow_types(ray_start_regular_shared, tmp_path):
|
||||
# Test that our custom arrow extension types are registered on initialization.
|
||||
ds = ray.data.from_items(np.zeros((8, 8, 8), dtype=np.int64))
|
||||
tmp_file = f"{tmp_path}/test.parquet"
|
||||
ds.write_parquet(tmp_file)
|
||||
|
||||
ds = ray.data.read_parquet(tmp_file)
|
||||
schema = "Column Type\n------ ----\nitem ArrowTensorTypeV2(shape=(8, 8), dtype=int64)"
|
||||
assert str(ds.schema()) == schema
|
||||
|
||||
# Also run in driver script to eliminate existing imports.
|
||||
driver_script = """import ray
|
||||
ds = ray.data.read_parquet("{0}")
|
||||
schema = ds.schema()
|
||||
assert str(schema) == \"\"\"{1}\"\"\"
|
||||
""".format(
|
||||
tmp_file, schema
|
||||
)
|
||||
run_string_as_driver(driver_script)
|
||||
|
||||
|
||||
def test_dict_doesnt_fallback_to_pandas_block(ray_start_regular_shared):
|
||||
# If the UDF returns a column with dict, previously, we would
|
||||
# fall back to pandas, because we couldn't convert it to
|
||||
# an Arrow block. This test checks that the block
|
||||
# construction now correctly goes to Arrow.
|
||||
def fn(batch):
|
||||
batch["data_dict"] = [{"data": 0} for _ in range(len(batch["id"]))]
|
||||
batch["data_objects"] = [
|
||||
types.SimpleNamespace(a=1, b="test") for _ in range(len(batch["id"]))
|
||||
]
|
||||
return batch
|
||||
|
||||
ds = ray.data.range(10).map_batches(fn)
|
||||
ds = ds.materialize()
|
||||
block = ray.get(ds.get_internal_block_refs()[0])
|
||||
assert isinstance(block, pa.Table), type(block)
|
||||
df_from_block = block.to_pandas()
|
||||
assert df_from_block["data_dict"].iloc[0] == {"data": 0}
|
||||
assert df_from_block["data_objects"].iloc[0] == types.SimpleNamespace(a=1, b="test")
|
||||
|
||||
def fn2(batch):
|
||||
batch["data_none"] = [None for _ in range(len(batch["id"]))]
|
||||
return batch
|
||||
|
||||
ds2 = ray.data.range(10).map_batches(fn2)
|
||||
ds2 = ds2.materialize()
|
||||
block = ray.get(ds2.get_internal_block_refs()[0])
|
||||
assert isinstance(block, pa.Table), type(block)
|
||||
df_from_block = block.to_pandas()
|
||||
assert df_from_block["data_none"].iloc[0] is None
|
||||
|
||||
|
||||
# Test for https://github.com/ray-project/ray/issues/49338.
|
||||
def test_build_block_with_null_column(ray_start_regular_shared, restore_data_context):
|
||||
ctx = DataContext.get_current()
|
||||
ctx.execution_options.preserve_order = True
|
||||
|
||||
# The blocks need to contain a tensor column to trigger the bug.
|
||||
block1 = BlockAccessor.batch_to_block(
|
||||
{"string": [None], "array": np.zeros((1, 2, 2))}
|
||||
)
|
||||
block2 = BlockAccessor.batch_to_block(
|
||||
{"string": ["spam"], "array": np.zeros((1, 2, 2))}
|
||||
)
|
||||
|
||||
builder = ArrowBlockBuilder()
|
||||
builder.add_block(block1)
|
||||
builder.add_block(block2)
|
||||
block = builder.build()
|
||||
|
||||
rows = list(BlockAccessor.for_block(block).iter_rows(True))
|
||||
assert len(rows) == 2
|
||||
|
||||
assert rows[0]["string"] is None
|
||||
assert rows[1]["string"] == "spam"
|
||||
assert np.array_equal(rows[0]["array"], np.zeros((2, 2)))
|
||||
assert np.array_equal(rows[1]["array"], np.zeros((2, 2)))
|
||||
|
||||
|
||||
def test_arrow_block_timestamp_ns(ray_start_regular_shared):
|
||||
# Input data with nanosecond precision timestamps
|
||||
data_rows = [
|
||||
{"col1": 1, "col2": pd.Timestamp("2023-01-01T00:00:00.123456789")},
|
||||
{"col1": 2, "col2": pd.Timestamp("2023-01-01T01:15:30.987654321")},
|
||||
{"col1": 3, "col2": pd.Timestamp("2023-01-01T02:30:15.111111111")},
|
||||
{"col1": 4, "col2": pd.Timestamp("2023-01-01T03:45:45.222222222")},
|
||||
{"col1": 5, "col2": pd.Timestamp("2023-01-01T05:00:00.333333333")},
|
||||
]
|
||||
|
||||
# Initialize ArrowBlockBuilder
|
||||
arrow_builder = ArrowBlockBuilder()
|
||||
for row in data_rows:
|
||||
arrow_builder.add(row)
|
||||
arrow_block = arrow_builder.build()
|
||||
|
||||
assert arrow_block.schema.field("col2").type == pa.timestamp("ns")
|
||||
for i, row in enumerate(data_rows):
|
||||
result_timestamp = arrow_block["col2"][i].as_py()
|
||||
# Convert both values to pandas Timestamp to preserve nanosecond precision for
|
||||
# comparison.
|
||||
assert pd.Timestamp(row["col2"]) == pd.Timestamp(
|
||||
result_timestamp
|
||||
), f"Timestamp mismatch at row {i} in ArrowBlockBuilder output"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"input_array,transform,expected_type,expected_values",
|
||||
[
|
||||
(
|
||||
pa.array([None, None], type=pa.string()),
|
||||
None,
|
||||
pa.string(),
|
||||
[None, None],
|
||||
),
|
||||
(
|
||||
pa.array([None, None], type=pa.list_(pa.string())),
|
||||
None,
|
||||
pa.list_(pa.string()),
|
||||
[None, None],
|
||||
),
|
||||
(
|
||||
pa.array([None, None], type=pa.decimal128(10, 2)),
|
||||
lambda df: df.fillna({"x": 0}),
|
||||
pa.decimal128(10, 2),
|
||||
[Decimal("0.00"), Decimal("0.00")],
|
||||
),
|
||||
(
|
||||
pa.array([["a", "b"], None], type=pa.list_(pa.string())),
|
||||
None,
|
||||
pa.list_(pa.string()),
|
||||
[["a", "b"], None],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_arrow_block_to_pandas_preserves_arrow_types_through_roundtrip(
|
||||
input_array, transform, expected_type, expected_values
|
||||
):
|
||||
table = pa.table({"x": input_array})
|
||||
|
||||
df = ArrowBlockAccessor(table).to_pandas()
|
||||
assert isinstance(df.dtypes["x"], pd.ArrowDtype)
|
||||
assert df.dtypes["x"].pyarrow_dtype == expected_type
|
||||
|
||||
if transform is not None:
|
||||
df = transform(df)
|
||||
|
||||
roundtripped = BlockAccessor.for_block(df).to_arrow()
|
||||
|
||||
assert roundtripped.schema.field("x").type == expected_type
|
||||
assert roundtripped.to_pydict() == {"x": expected_values}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,514 @@
|
||||
import logging
|
||||
import os
|
||||
import types
|
||||
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
import pyarrow.parquet as pq
|
||||
import pytest
|
||||
from packaging.version import parse as parse_version
|
||||
from pytest_lazy_fixtures import lf as lazy_fixture
|
||||
|
||||
import ray
|
||||
import ray.cloudpickle as pickle
|
||||
import ray.data
|
||||
import ray.train
|
||||
from ray.data._internal.utils.arrow_utils import get_pyarrow_version
|
||||
from ray.data.extensions.object_extension import (
|
||||
ArrowPythonObjectArray,
|
||||
)
|
||||
from ray.data.extensions.tensor_extension import (
|
||||
ArrowTensorArray,
|
||||
ArrowVariableShapedTensorArray,
|
||||
)
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def null_array():
|
||||
return pa.array([])
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def int_array():
|
||||
return pa.array(list(range(1000)))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def int_array_with_nulls():
|
||||
return pa.array((list(range(9)) + [None]) * 100)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def float_array():
|
||||
return pa.array([float(i) for i in range(1000)])
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def boolean_array():
|
||||
return pa.array([True, False] * 500)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def string_array():
|
||||
return pa.array(["foo", "bar", "bz", None, "quux"] * 200)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def large_string_array():
|
||||
return pa.array(["foo", "bar", "bz", None, "quux"] * 200, type=pa.large_string())
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def binary_array():
|
||||
return pa.array([b"foo", b"bar", b"bz", None, b"quux"] * 200)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fixed_size_binary_array():
|
||||
return pa.array([b"foo", b"bar", b"baz", None, b"qux"] * 200, type=pa.binary(3))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def large_binary_array():
|
||||
return pa.array(
|
||||
[b"foo", b"bar", b"bz", None, b"quux"] * 200, type=pa.large_binary()
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def list_array():
|
||||
return pa.array(([None] + [list(range(9)) + [None]] * 9) * 100)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def large_list_array():
|
||||
# Large list array with nulls
|
||||
return pa.array(
|
||||
([None] + [list(range(9)) + [None]] * 9) * 100,
|
||||
type=pa.large_list(pa.int64()),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fixed_size_list_array():
|
||||
# Fixed size list array
|
||||
return pa.FixedSizeListArray.from_arrays(
|
||||
pa.array((list(range(9)) + [None]) * 1000), 10
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def map_array():
|
||||
return pa.array(
|
||||
[list(zip("abcdefghij", range(10))) for _ in range(1000)],
|
||||
type=pa.map_(pa.string(), pa.int64()),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def struct_array():
|
||||
# Struct array
|
||||
return pa.array({"a": i} for i in range(1000))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sparse_union_array():
|
||||
return pa.UnionArray.from_sparse(
|
||||
pa.array([0, 1] * 500, type=pa.int8()),
|
||||
[pa.array(list(range(1000))), pa.array([True, False] * 500)],
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dense_union_array():
|
||||
return pa.UnionArray.from_dense(
|
||||
pa.array([0, 1] * 500, type=pa.int8()),
|
||||
pa.array(
|
||||
[i if i % 2 == 0 else (i % 3) % 2 for i in range(1000)], type=pa.int32()
|
||||
),
|
||||
[pa.array(list(range(1000))), pa.array([True, False])],
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def dictionary_array():
|
||||
return pa.DictionaryArray.from_arrays(
|
||||
pa.array((list(range(9)) + [None]) * 100),
|
||||
pa.array(["a", "b", "c", "d", "e", "f", "g", "h", "i"]),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tensor_array():
|
||||
return ArrowTensorArray.from_numpy(np.arange(1000 * 4 * 4).reshape((1000, 4, 4)))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def boolean_tensor_array():
|
||||
return ArrowTensorArray.from_numpy(
|
||||
np.array(
|
||||
[True, False, False, True, False, False, True, True] * 2 * 1000
|
||||
).reshape((1000, 4, 4))
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def variable_shaped_tensor_array():
|
||||
return ArrowVariableShapedTensorArray.from_numpy(
|
||||
np.array(
|
||||
[
|
||||
np.arange(4).reshape((2, 2)),
|
||||
np.arange(4, 13).reshape((3, 3)),
|
||||
]
|
||||
* 500,
|
||||
dtype=object,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def boolean_variable_shaped_tensor_array():
|
||||
return ArrowVariableShapedTensorArray.from_numpy(
|
||||
np.array(
|
||||
[
|
||||
np.array([[True, False], [False, True]]),
|
||||
np.array(
|
||||
[
|
||||
[False, True, False],
|
||||
[True, True, False],
|
||||
[False, False, False],
|
||||
],
|
||||
),
|
||||
]
|
||||
* 500,
|
||||
dtype=object,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def list_of_struct_array():
|
||||
return pa.array([{"a": i}, {"a": -i}] for i in range(1000))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def list_of_empty_struct_array():
|
||||
return pa.array([{}, {}] for i in range(1000))
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def complex_nested_array():
|
||||
return pa.UnionArray.from_sparse(
|
||||
pa.array([0, 1] * 500, type=pa.int8()),
|
||||
[
|
||||
pa.array(
|
||||
[
|
||||
{
|
||||
"a": i % 2 == 0,
|
||||
"b": i,
|
||||
"c": "bar",
|
||||
}
|
||||
for i in range(1000)
|
||||
]
|
||||
),
|
||||
pa.array(
|
||||
[list(zip("abcdefghij", range(10))) for _ in range(1000)],
|
||||
type=pa.map_(pa.string(), pa.int64()),
|
||||
),
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def pickled_objects_array():
|
||||
elements = ["test", 20, False, {"some": "value"}, None, np.zeros((10, 10))]
|
||||
elements *= 1 + 1000 // len(elements)
|
||||
elements = elements[:1000]
|
||||
|
||||
arr = np.array(elements, dtype=object)
|
||||
return ArrowPythonObjectArray.from_objects(arr)
|
||||
|
||||
|
||||
pytest_custom_serialization_arrays = [
|
||||
# Null array
|
||||
(lazy_fixture("null_array"), 1.0),
|
||||
# Int array
|
||||
(lazy_fixture("int_array"), 0.1),
|
||||
# Array with nulls
|
||||
(lazy_fixture("int_array_with_nulls"), 0.1),
|
||||
# Float array
|
||||
(lazy_fixture("float_array"), 0.1),
|
||||
# Boolean array
|
||||
# Due to bit-packing, most of the pickle bytes are metadata.
|
||||
(lazy_fixture("boolean_array"), 0.8),
|
||||
# String array
|
||||
(lazy_fixture("string_array"), 0.1),
|
||||
# Large string array
|
||||
(lazy_fixture("large_string_array"), 0.1),
|
||||
# Binary array
|
||||
(lazy_fixture("binary_array"), 0.1),
|
||||
# Fixed size binary array
|
||||
(lazy_fixture("fixed_size_binary_array"), 0.1),
|
||||
# Large binary array
|
||||
(lazy_fixture("large_binary_array"), 0.1),
|
||||
# List array with nulls
|
||||
(lazy_fixture("list_array"), 0.1),
|
||||
# Large list array with nulls
|
||||
(lazy_fixture("large_list_array"), 0.1),
|
||||
# Fixed size list array
|
||||
(lazy_fixture("fixed_size_list_array"), 0.1),
|
||||
# Map array
|
||||
(lazy_fixture("map_array"), 0.1),
|
||||
# Struct array
|
||||
(lazy_fixture("struct_array"), 0.1),
|
||||
# Union array (sparse)
|
||||
(lazy_fixture("sparse_union_array"), 0.1),
|
||||
# Union array (dense)
|
||||
(lazy_fixture("dense_union_array"), 0.1),
|
||||
# Dictionary array
|
||||
(lazy_fixture("dictionary_array"), 0.1),
|
||||
# Tensor extension array
|
||||
(lazy_fixture("tensor_array"), 0.1),
|
||||
# Boolean tensor extension array
|
||||
(lazy_fixture("boolean_tensor_array"), 0.25),
|
||||
# Variable-shaped tensor extension array
|
||||
(lazy_fixture("variable_shaped_tensor_array"), 0.1),
|
||||
# Boolean variable-shaped tensor extension array
|
||||
(lazy_fixture("boolean_variable_shaped_tensor_array"), 0.25),
|
||||
# List of struct array
|
||||
(lazy_fixture("list_of_struct_array"), 0.1),
|
||||
# List of empty struct array
|
||||
(lazy_fixture("list_of_empty_struct_array"), 0.1),
|
||||
# Complex nested array
|
||||
(lazy_fixture("complex_nested_array"), 0.1),
|
||||
# Array of pickled objects
|
||||
(lazy_fixture("pickled_objects_array"), 0.1),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("data,cap_mult", pytest_custom_serialization_arrays)
|
||||
def test_custom_arrow_data_serializer(ray_start_regular_shared, data, cap_mult):
|
||||
if len(data) == 0:
|
||||
data = pa.table({"a": []})
|
||||
else:
|
||||
data = pa.Table.from_arrays(
|
||||
[data, data, pa.array(range(1000), type=pa.int32())],
|
||||
schema=pa.schema(
|
||||
[
|
||||
pa.field("arr1", data.type),
|
||||
pa.field("arr2", data.type),
|
||||
pa.field("arr3", pa.int32()),
|
||||
],
|
||||
metadata={b"foo": b"bar"},
|
||||
),
|
||||
)
|
||||
ray._private.worker.global_worker.get_serialization_context()
|
||||
data.validate()
|
||||
pyarrow_version = get_pyarrow_version()
|
||||
if pyarrow_version >= parse_version("7.0.0"):
|
||||
# get_total_buffer_size API was added in Arrow 7.0.0.
|
||||
buf_size = data.get_total_buffer_size()
|
||||
# Create a zero-copy slice view of data.
|
||||
view = data.slice(10, 10)
|
||||
s_arr = pickle.dumps(data)
|
||||
s_view = pickle.dumps(view)
|
||||
post_slice = pickle.loads(s_view)
|
||||
post_slice.validate()
|
||||
# Check for round-trip equality.
|
||||
assert view.equals(post_slice), post_slice
|
||||
# Check that the slice view was truncated upon serialization.
|
||||
assert len(s_view) <= cap_mult * len(s_arr)
|
||||
for column, pre_column in zip(post_slice.columns, view.columns):
|
||||
# Check that offset was reset on slice.
|
||||
if column.num_chunks > 0:
|
||||
assert column.chunk(0).offset == 0
|
||||
# Check that null count was either properly cached or recomputed.
|
||||
assert column.null_count == pre_column.null_count
|
||||
if pyarrow_version >= parse_version("7.0.0"):
|
||||
# Check that slice buffer only contains slice data.
|
||||
slice_buf_size = post_slice.get_total_buffer_size()
|
||||
if buf_size > 0:
|
||||
assert buf_size / slice_buf_size - len(data) / len(post_slice) < 100
|
||||
|
||||
|
||||
def test_custom_arrow_data_serializer_fallback(
|
||||
ray_start_regular_shared, propagate_logs, caplog
|
||||
):
|
||||
# Reset serialization fallback set so warning is logged.
|
||||
import ray._private.arrow_serialization as arrow_ser_module
|
||||
|
||||
arrow_ser_module._serialization_fallback_set = set()
|
||||
|
||||
data = pa.table(
|
||||
{
|
||||
"a": pa.UnionArray.from_dense(
|
||||
pa.array([0, 1] * 500, type=pa.int8()),
|
||||
pa.array(
|
||||
[i if i % 2 == 0 else (i % 3) % 2 for i in range(1000)],
|
||||
type=pa.int32(),
|
||||
),
|
||||
[pa.array(list(range(1000))), pa.array([True, False])],
|
||||
)
|
||||
}
|
||||
)
|
||||
cap_mult = 0.1
|
||||
ray._private.worker.global_worker.get_serialization_context()
|
||||
data.validate()
|
||||
pyarrow_version = get_pyarrow_version()
|
||||
if pyarrow_version >= parse_version("7.0.0"):
|
||||
# get_total_buffer_size API was added in Arrow 7.0.0.
|
||||
buf_size = data.get_total_buffer_size()
|
||||
# Create a zero-copy slice view of data.
|
||||
view = data.slice(10, 10)
|
||||
# Confirm that (1) fallback works, and (2) warning is logged.
|
||||
with caplog.at_level(
|
||||
logging.WARNING,
|
||||
logger="ray.data._internal.arrow_serialization",
|
||||
):
|
||||
s_arr = pickle.dumps(data)
|
||||
assert "Failed to complete optimized serialization" in caplog.text
|
||||
caplog.clear()
|
||||
# Confirm that we only warn once per process.
|
||||
with caplog.at_level(
|
||||
logging.WARNING,
|
||||
logger="ray.data._internal.arrow_serialization",
|
||||
):
|
||||
s_view = pickle.dumps(view)
|
||||
assert "Failed to complete optimized serialization" not in caplog.text
|
||||
post_slice = pickle.loads(s_view)
|
||||
post_slice.validate()
|
||||
# Check for round-trip equality.
|
||||
assert view.equals(post_slice), post_slice
|
||||
# Check that the slice view was truncated upon serialization.
|
||||
assert len(s_view) <= cap_mult * len(s_arr)
|
||||
for column, pre_column in zip(post_slice.columns, view.columns):
|
||||
# Check that offset was reset on slice.
|
||||
if column.num_chunks > 0:
|
||||
assert column.chunk(0).offset == 0
|
||||
# Check that null count was either properly cached or recomputed.
|
||||
assert column.null_count == pre_column.null_count
|
||||
if pyarrow_version >= parse_version("7.0.0"):
|
||||
# Check that slice buffer only contains slice data.
|
||||
slice_buf_size = post_slice.get_total_buffer_size()
|
||||
if buf_size > 0:
|
||||
assert buf_size / slice_buf_size - len(data) / len(post_slice) < 100
|
||||
|
||||
|
||||
def test_arrow_scalar_conversion(ray_start_regular_shared):
|
||||
ds = ray.data.from_items([1])
|
||||
|
||||
def fn(batch: list):
|
||||
return {"id": np.array([1])}
|
||||
|
||||
ds = ds.map_batches(fn)
|
||||
res = ds.take()
|
||||
assert res == [{"id": 1}], res
|
||||
|
||||
|
||||
def test_arrow_object_and_array_support(ray_start_regular_shared):
|
||||
obj = types.SimpleNamespace(some_attribute="test")
|
||||
|
||||
def f(batch):
|
||||
batch_size = len(batch["id"])
|
||||
return {
|
||||
"array": np.zeros((batch_size, 32, 32, 3)),
|
||||
"unsupported": [obj] * batch_size,
|
||||
}
|
||||
|
||||
res = ray.data.range(5).map_batches(f, batch_size=None).take(1)
|
||||
assert res[0]["array"].shape == (32, 32, 3)
|
||||
assert np.all(res[0]["array"] == 0)
|
||||
assert res[0]["unsupported"] == obj
|
||||
|
||||
|
||||
def test_custom_arrow_data_serializer_parquet_roundtrip(
|
||||
ray_start_regular_shared, tmp_path
|
||||
):
|
||||
ray._private.worker.global_worker.get_serialization_context()
|
||||
t = pa.table({"a": list(range(10000000))})
|
||||
pq.write_table(t, f"{tmp_path}/test.parquet")
|
||||
t2 = pq.read_table(f"{tmp_path}/test.parquet")
|
||||
s_t = pickle.dumps(t)
|
||||
s_t2 = pickle.dumps(t2)
|
||||
# Check that the post-Parquet slice view chunks don't cause a serialization blow-up.
|
||||
assert len(s_t2) < 1.1 * len(s_t)
|
||||
# Check for round-trip equality.
|
||||
assert t2.equals(pickle.loads(s_t2))
|
||||
|
||||
|
||||
def test_arrow_schema_ipc_serialization(ray_start_regular_shared):
|
||||
"""Test that Arrow Schema uses IPC serialization for performance."""
|
||||
from ray._private.arrow_serialization import (
|
||||
_arrow_schema_reduce,
|
||||
_restore_schema_from_ipc,
|
||||
)
|
||||
|
||||
# Verify the reducer is registered
|
||||
ray._private.worker.global_worker.get_serialization_context()
|
||||
assert pa.Schema in pickle.CloudPickler.dispatch
|
||||
assert pickle.CloudPickler.dispatch[pa.Schema] == _arrow_schema_reduce
|
||||
|
||||
# Create a complex schema with various types
|
||||
schema = pa.schema(
|
||||
[
|
||||
pa.field("id", pa.int64()),
|
||||
pa.field("name", pa.string()),
|
||||
pa.field("timestamp", pa.timestamp("us", tz="UTC")),
|
||||
pa.field("tags", pa.list_(pa.string())),
|
||||
pa.field("metadata", pa.map_(pa.string(), pa.string())),
|
||||
pa.field(
|
||||
"nested",
|
||||
pa.struct(
|
||||
[
|
||||
pa.field("x", pa.float64()),
|
||||
pa.field("y", pa.float64()),
|
||||
]
|
||||
),
|
||||
),
|
||||
pa.field("category", pa.dictionary(pa.int8(), pa.string())),
|
||||
pa.field("decimal_val", pa.decimal128(18, 6)),
|
||||
],
|
||||
metadata={b"foo": b"bar"},
|
||||
)
|
||||
|
||||
# Test roundtrip serialization
|
||||
serialized = pickle.dumps(schema)
|
||||
deserialized = pickle.loads(serialized)
|
||||
assert schema.equals(deserialized)
|
||||
assert schema.metadata == deserialized.metadata
|
||||
|
||||
# Verify the reducer uses IPC format (check via direct call)
|
||||
restore_func, (ipc_bytes,) = _arrow_schema_reduce(schema)
|
||||
assert restore_func == _restore_schema_from_ipc
|
||||
# IPC bytes should match what schema.serialize() produces
|
||||
assert ipc_bytes == schema.serialize().to_pybytes()
|
||||
# Verify restore works
|
||||
restored = restore_func(ipc_bytes)
|
||||
assert schema.equals(restored)
|
||||
|
||||
|
||||
def test_custom_arrow_data_serializer_disable(shutdown_only):
|
||||
ray.shutdown()
|
||||
ray.worker._post_init_hooks = []
|
||||
context = ray.worker.global_worker.get_serialization_context()
|
||||
context._unregister_cloudpickle_reducer(pa.Table)
|
||||
# Disable custom Arrow array serialization.
|
||||
os.environ["RAY_DISABLE_CUSTOM_ARROW_ARRAY_SERIALIZATION"] = "1"
|
||||
ray.init()
|
||||
# Create a zero-copy slice view of table.
|
||||
t = pa.table({"a": list(range(10000000))})
|
||||
view = t.slice(10, 10)
|
||||
s_t = pickle.dumps(t)
|
||||
s_view = pickle.dumps(view)
|
||||
# Check that the slice view contains the full buffer of the underlying array.
|
||||
d_view = pickle.loads(s_view)
|
||||
assert d_view["a"].chunk(0).buffers()[1].size == t["a"].chunk(0).buffers()[1].size
|
||||
# Check that the serialized slice view is large
|
||||
assert len(s_view) > 0.8 * len(s_t)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,20 @@
|
||||
"""Tests for batch_size="auto" in BatchMapTransformFn and map_batches."""
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
|
||||
|
||||
def test_map_batches_auto_correctness(ray_start_regular_shared):
|
||||
"""batch_size='auto' preserves all rows and values in an end-to-end pipeline."""
|
||||
n_rows = 100
|
||||
rows = (
|
||||
ray.data.range(n_rows)
|
||||
.map_batches(lambda batch: batch, batch_size="auto")
|
||||
.take_all()
|
||||
)
|
||||
assert len(rows) == n_rows
|
||||
assert sorted(r["id"] for r in rows) == list(range(n_rows))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -0,0 +1,184 @@
|
||||
from dataclasses import astuple, dataclass
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data._internal.util import _autodetect_parallelism
|
||||
from ray.data.context import DataContext
|
||||
from ray.tests.conftest import * # noqa
|
||||
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
||||
|
||||
|
||||
@dataclass
|
||||
class TestCase:
|
||||
avail_cpus: int
|
||||
target_max_block_size: int
|
||||
data_size: int
|
||||
expected_parallelism: int
|
||||
|
||||
|
||||
MiB = 1024 * 1024
|
||||
GiB = 1024 * MiB
|
||||
|
||||
TEST_CASES = [
|
||||
TestCase(
|
||||
avail_cpus=4,
|
||||
target_max_block_size=DataContext.get_current().target_max_block_size,
|
||||
data_size=1024,
|
||||
expected_parallelism=8, # avail_cpus has precedence
|
||||
),
|
||||
TestCase(
|
||||
avail_cpus=4,
|
||||
target_max_block_size=DataContext.get_current().target_max_block_size,
|
||||
data_size=10 * MiB,
|
||||
expected_parallelism=10, # MIN_BLOCK_SIZE has precedence
|
||||
),
|
||||
TestCase(
|
||||
avail_cpus=4,
|
||||
target_max_block_size=DataContext.get_current().target_max_block_size,
|
||||
data_size=20 * MiB,
|
||||
expected_parallelism=20, # MIN_BLOCK_SIZE has precedence
|
||||
),
|
||||
TestCase(
|
||||
avail_cpus=4,
|
||||
target_max_block_size=DataContext.get_current().target_max_block_size,
|
||||
data_size=100 * MiB,
|
||||
expected_parallelism=100, # MIN_BLOCK_SIZE has precedence
|
||||
),
|
||||
TestCase(
|
||||
avail_cpus=4,
|
||||
target_max_block_size=DataContext.get_current().target_max_block_size,
|
||||
data_size=1 * GiB,
|
||||
expected_parallelism=200, # MIN_PARALLELISM has precedence
|
||||
),
|
||||
TestCase(
|
||||
avail_cpus=4,
|
||||
target_max_block_size=DataContext.get_current().target_max_block_size,
|
||||
data_size=10 * GiB,
|
||||
expected_parallelism=200, # MIN_PARALLELISM has precedence
|
||||
),
|
||||
TestCase(
|
||||
avail_cpus=150,
|
||||
target_max_block_size=DataContext.get_current().target_max_block_size,
|
||||
data_size=10 * GiB,
|
||||
expected_parallelism=300, # avail_cpus has precedence
|
||||
),
|
||||
TestCase(
|
||||
avail_cpus=400,
|
||||
target_max_block_size=DataContext.get_current().target_max_block_size,
|
||||
data_size=10 * GiB,
|
||||
expected_parallelism=800, # avail_cpus has precedence
|
||||
),
|
||||
TestCase(
|
||||
avail_cpus=400,
|
||||
target_max_block_size=DataContext.get_current().target_max_block_size,
|
||||
data_size=1 * MiB,
|
||||
expected_parallelism=800, # avail_cpus has precedence
|
||||
),
|
||||
TestCase(
|
||||
avail_cpus=4,
|
||||
target_max_block_size=DataContext.get_current().target_max_block_size,
|
||||
data_size=1000 * GiB,
|
||||
expected_parallelism=8000, # MAX_BLOCK_SIZE has precedence
|
||||
),
|
||||
TestCase(
|
||||
avail_cpus=4,
|
||||
target_max_block_size=DataContext.get_current().target_max_block_size,
|
||||
data_size=10000 * GiB,
|
||||
expected_parallelism=80000, # MAX_BLOCK_SIZE has precedence
|
||||
),
|
||||
TestCase(
|
||||
avail_cpus=4,
|
||||
target_max_block_size=512 * MiB,
|
||||
data_size=1000 * GiB,
|
||||
expected_parallelism=2000, # passed max_block_size has precedence
|
||||
),
|
||||
TestCase(
|
||||
avail_cpus=4,
|
||||
target_max_block_size=512 * MiB,
|
||||
data_size=10000 * GiB,
|
||||
expected_parallelism=20000, # passed max_block_size has precedence
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"avail_cpus,target_max_block_size,data_size,expected",
|
||||
[astuple(test) for test in TEST_CASES],
|
||||
)
|
||||
def test_autodetect_parallelism(
|
||||
shutdown_only, avail_cpus, target_max_block_size, data_size, expected
|
||||
):
|
||||
class MockReader:
|
||||
def estimate_inmemory_data_size(self):
|
||||
return data_size
|
||||
|
||||
result, _, _ = _autodetect_parallelism(
|
||||
parallelism=-1,
|
||||
target_max_block_size=target_max_block_size,
|
||||
ctx=DataContext.get_current(),
|
||||
datasource_or_legacy_reader=MockReader(),
|
||||
avail_cpus=avail_cpus,
|
||||
)
|
||||
assert result == expected, (result, expected)
|
||||
|
||||
|
||||
def test_auto_parallelism_basic(shutdown_only):
|
||||
ray.init(num_cpus=8)
|
||||
context = DataContext.get_current()
|
||||
context.read_op_min_num_blocks = 1
|
||||
# Datasource bound.
|
||||
ds = ray.data.range_tensor(5, shape=(100,), override_num_blocks=-1)
|
||||
assert ds._logical_plan.initial_num_blocks() == 5, ds
|
||||
# CPU bound. TODO(ekl) we should fix range datasource to respect parallelism more
|
||||
# properly, currently it can go a little over.
|
||||
ds = ray.data.range_tensor(10000, shape=(100,), override_num_blocks=-1)
|
||||
assert ds._logical_plan.initial_num_blocks() == 16, ds
|
||||
# Block size bound.
|
||||
ds = ray.data.range_tensor(100000000, shape=(100,), override_num_blocks=-1)
|
||||
assert ds._logical_plan.initial_num_blocks() >= 590, ds
|
||||
assert ds._logical_plan.initial_num_blocks() <= 600, ds
|
||||
|
||||
|
||||
def test_auto_parallelism_placement_group(shutdown_only):
|
||||
ray.init(num_cpus=16, num_gpus=8)
|
||||
|
||||
@ray.remote
|
||||
def run():
|
||||
context = DataContext.get_current()
|
||||
context.min_parallelism = 1
|
||||
ds = ray.data.range_tensor(2000, shape=(100,), override_num_blocks=-1)
|
||||
return ds._logical_plan.initial_num_blocks()
|
||||
|
||||
# 1/16 * 4 * 16 = 4
|
||||
pg = ray.util.placement_group([{"CPU": 1}])
|
||||
num_blocks = ray.get(
|
||||
run.options(
|
||||
scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=pg)
|
||||
).remote()
|
||||
)
|
||||
assert num_blocks == 4, num_blocks
|
||||
|
||||
# 2/16 * 4 * 16 = 8
|
||||
pg = ray.util.placement_group([{"CPU": 2}])
|
||||
num_blocks = ray.get(
|
||||
run.options(
|
||||
scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=pg)
|
||||
).remote()
|
||||
)
|
||||
assert num_blocks == 8, num_blocks
|
||||
|
||||
# 1/8 * 4 * 16 = 8
|
||||
pg = ray.util.placement_group([{"CPU": 1, "GPU": 1}])
|
||||
num_blocks = ray.get(
|
||||
run.options(
|
||||
scheduling_strategy=PlacementGroupSchedulingStrategy(placement_group=pg)
|
||||
).remote()
|
||||
)
|
||||
assert num_blocks == 8, num_blocks
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,971 @@
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.cluster_utils import Cluster
|
||||
from ray.data._internal.cluster_autoscaler.default_autoscaling_coordinator import (
|
||||
HEAD_NODE_RESOURCE_LABEL,
|
||||
DefaultAutoscalingCoordinator,
|
||||
_AutoscalingCoordinatorActor,
|
||||
_format_resources_for_log,
|
||||
get_or_create_autoscaling_coordinator,
|
||||
)
|
||||
from ray.data._internal.util import GiB
|
||||
from ray.tests.conftest import wait_for_condition
|
||||
|
||||
CLUSTER_NODES_WITH_HEAD = [
|
||||
# Head node should be included if it has non-zero CPUs or GPUs.
|
||||
{
|
||||
"Resources": {
|
||||
"CPU": 10,
|
||||
"GPU": 5,
|
||||
"object_store_memory": 1000,
|
||||
HEAD_NODE_RESOURCE_LABEL: 1,
|
||||
},
|
||||
"Alive": True,
|
||||
},
|
||||
# Dead node should be excluded.
|
||||
{
|
||||
"Resources": {
|
||||
"CPU": 10,
|
||||
"GPU": 5,
|
||||
"object_store_memory": 1000,
|
||||
},
|
||||
"Alive": False,
|
||||
},
|
||||
]
|
||||
|
||||
CLUSTER_NODES_WITHOUT_HEAD = [
|
||||
{
|
||||
"Resources": {"CPU": 10, "GPU": 5, "object_store_memory": 1000},
|
||||
"Alive": True,
|
||||
},
|
||||
# Head node should be excluded if CPUs and GPUs are both 0.
|
||||
{
|
||||
"Resources": {
|
||||
"CPU": 0,
|
||||
"GPU": 0,
|
||||
"object_store_memory": 1000,
|
||||
HEAD_NODE_RESOURCE_LABEL: 1,
|
||||
},
|
||||
"Alive": True,
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"cluster_nodes",
|
||||
[
|
||||
CLUSTER_NODES_WITH_HEAD,
|
||||
CLUSTER_NODES_WITHOUT_HEAD,
|
||||
],
|
||||
)
|
||||
def test_basic(cluster_nodes):
|
||||
mocked_time = 0
|
||||
|
||||
mock_request_resources = Mock()
|
||||
as_coordinator = _AutoscalingCoordinatorActor(
|
||||
get_current_time=lambda: mocked_time,
|
||||
send_resources_request=mock_request_resources,
|
||||
get_cluster_nodes=lambda: cluster_nodes,
|
||||
)
|
||||
|
||||
req1 = [{"CPU": 3, "GPU": 1, "object_store_memory": 100}]
|
||||
req1_timeout = 2
|
||||
as_coordinator.request_resources(
|
||||
requester_id="requester1",
|
||||
resources=req1,
|
||||
expire_after_s=req1_timeout,
|
||||
)
|
||||
mock_request_resources.assert_called_once_with(req1)
|
||||
res1 = as_coordinator.get_allocated_resources("requester1")
|
||||
|
||||
def _remove_head_node_resources(res):
|
||||
for r in res:
|
||||
if HEAD_NODE_RESOURCE_LABEL in r:
|
||||
del r[HEAD_NODE_RESOURCE_LABEL]
|
||||
|
||||
_remove_head_node_resources(res1)
|
||||
assert res1 == req1
|
||||
|
||||
# Send the same request again. `mock_request_resources` won't be called
|
||||
# since the request is not updated.
|
||||
as_coordinator.request_resources(
|
||||
requester_id="requester1",
|
||||
resources=req1,
|
||||
expire_after_s=req1_timeout,
|
||||
)
|
||||
assert mock_request_resources.call_count == 1
|
||||
|
||||
# Send a request from requester2, with request_remaining=True.
|
||||
# requester2 should get the requested + the remaining resources.
|
||||
req2 = [{"CPU": 2, "GPU": 1, "object_store_memory": 100}]
|
||||
req2_timeout = 20
|
||||
as_coordinator.request_resources(
|
||||
requester_id="requester2",
|
||||
resources=req2,
|
||||
expire_after_s=req2_timeout,
|
||||
request_remaining=True,
|
||||
)
|
||||
mock_request_resources.assert_called_with(req1 + req2)
|
||||
res2 = as_coordinator.get_allocated_resources("requester2")
|
||||
_remove_head_node_resources(res2)
|
||||
assert res2 == req2 + [{"CPU": 5, "GPU": 3, "object_store_memory": 800}]
|
||||
|
||||
# Test updating req1
|
||||
req1_updated = [{"CPU": 4, "GPU": 2, "object_store_memory": 300}]
|
||||
as_coordinator.request_resources(
|
||||
requester_id="requester1",
|
||||
resources=req1_updated,
|
||||
expire_after_s=req1_timeout,
|
||||
)
|
||||
mock_request_resources.assert_called_with(req1_updated + req2)
|
||||
res1 = as_coordinator.get_allocated_resources("requester1")
|
||||
_remove_head_node_resources(res1)
|
||||
assert res1 == req1_updated
|
||||
res2 = as_coordinator.get_allocated_resources("requester2")
|
||||
_remove_head_node_resources(res2)
|
||||
assert res2 == req2 + [{"CPU": 4, "GPU": 2, "object_store_memory": 600}]
|
||||
|
||||
# After req1_timeout, req1 should be expired.
|
||||
mocked_time = req1_timeout + 0.1
|
||||
as_coordinator._tick()
|
||||
mock_request_resources.assert_called_with(req2)
|
||||
res1 = as_coordinator.get_allocated_resources("requester1")
|
||||
res2 = as_coordinator.get_allocated_resources("requester2")
|
||||
_remove_head_node_resources(res1)
|
||||
_remove_head_node_resources(res2)
|
||||
assert res1 == []
|
||||
assert res2 == req2 + [{"CPU": 8, "GPU": 4, "object_store_memory": 900}]
|
||||
|
||||
# After req2_timeout, req2 should be expired.
|
||||
mocked_time = req2_timeout + 0.1
|
||||
as_coordinator._tick()
|
||||
mock_request_resources.assert_called_with([])
|
||||
res1 = as_coordinator.get_allocated_resources("requester1")
|
||||
res2 = as_coordinator.get_allocated_resources("requester2")
|
||||
_remove_head_node_resources(res1)
|
||||
_remove_head_node_resources(res2)
|
||||
assert res1 == []
|
||||
assert res2 == []
|
||||
|
||||
# Test canceling a request
|
||||
as_coordinator.cancel_request("requester2")
|
||||
res2 = as_coordinator.get_allocated_resources("requester2")
|
||||
_remove_head_node_resources(res2)
|
||||
assert res2 == []
|
||||
|
||||
|
||||
def test_double_allocation_with_multiple_request_remaining():
|
||||
"""Test fair allocation when multiple requesters have request_remaining=True."""
|
||||
cluster_nodes = [
|
||||
{
|
||||
"Resources": {
|
||||
"CPU": 10,
|
||||
"GPU": 5,
|
||||
"object_store_memory": 1000,
|
||||
},
|
||||
"Alive": True,
|
||||
}
|
||||
]
|
||||
|
||||
mocked_time = 0
|
||||
mock_request_resources = Mock()
|
||||
coordinator = _AutoscalingCoordinatorActor(
|
||||
get_current_time=lambda: mocked_time,
|
||||
send_resources_request=mock_request_resources,
|
||||
get_cluster_nodes=lambda: cluster_nodes,
|
||||
)
|
||||
|
||||
# Requester1: asks for CPU=2, GPU=1 with request_remaining=True
|
||||
req1 = [{"CPU": 2, "GPU": 1, "object_store_memory": 100}]
|
||||
coordinator.request_resources(
|
||||
requester_id="requester1",
|
||||
resources=req1,
|
||||
expire_after_s=100,
|
||||
request_remaining=True,
|
||||
)
|
||||
|
||||
# Requester2: asks for CPU=3, GPU=1 with request_remaining=True
|
||||
req2 = [{"CPU": 3, "GPU": 1, "object_store_memory": 200}]
|
||||
coordinator.request_resources(
|
||||
requester_id="requester2",
|
||||
resources=req2,
|
||||
expire_after_s=100,
|
||||
request_remaining=True,
|
||||
)
|
||||
|
||||
# Get allocated resources
|
||||
res1 = coordinator.get_allocated_resources("requester1")
|
||||
res2 = coordinator.get_allocated_resources("requester2")
|
||||
|
||||
# After allocating specific requests (req1 and req2):
|
||||
# Remaining = CPU: 10-2-3=5, GPU: 5-1-1=3, memory: 1000-100-200=700
|
||||
# With fair allocation, each requester gets 1/2 of remaining resources
|
||||
expected_remaining_per_requester = {
|
||||
"CPU": 5 // 2, # = 2
|
||||
"GPU": 3 // 2, # = 1
|
||||
"object_store_memory": 700 // 2, # = 350
|
||||
}
|
||||
|
||||
# Both requesters should get their specific requests + fair share of remaining
|
||||
assert res1 == req1 + [expected_remaining_per_requester]
|
||||
assert res2 == req2 + [expected_remaining_per_requester]
|
||||
|
||||
|
||||
def test_format_resources_for_log():
|
||||
# Two bundles that are identical except for sub-precision differences in
|
||||
# object_store_memory (8.96 vs 9.04 GiB), plus custom labels and a GPU: 0
|
||||
# entry. They should aggregate into a single entry, with custom/zero
|
||||
# resources dropped.
|
||||
resources = [
|
||||
{
|
||||
"CPU": 8,
|
||||
"GPU": 0,
|
||||
"memory": 32 * GiB,
|
||||
"object_store_memory": int(8.96 * GiB),
|
||||
"anyscale/cpu_only:true": 1.0,
|
||||
"anyscale/region:us-west-2": 1.0,
|
||||
"node:10.0.193.159": 1.0,
|
||||
},
|
||||
{
|
||||
"CPU": 8,
|
||||
"GPU": 0,
|
||||
"memory": 32 * GiB,
|
||||
"object_store_memory": int(9.04 * GiB),
|
||||
"anyscale/cpu_only:true": 1.0,
|
||||
"anyscale/region:us-west-2": 1.0,
|
||||
"node:10.0.241.173": 1.0,
|
||||
},
|
||||
{
|
||||
"CPU": 0,
|
||||
"GPU": 0,
|
||||
"memory": 0,
|
||||
"object_store_memory": 0,
|
||||
"anyscale/cpu_only:true": 1.0,
|
||||
"node:10.0.252.14": 1.0,
|
||||
},
|
||||
]
|
||||
|
||||
log_message = _format_resources_for_log(resources)
|
||||
assert log_message == (
|
||||
"[2 x {CPU: 8, memory: 32.0GiB, object_store_memory: 9.0GiB}]"
|
||||
)
|
||||
assert "anyscale/" not in log_message
|
||||
assert "node:" not in log_message
|
||||
assert "GPU" not in log_message
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cluster():
|
||||
"""Initialize a Ray cluster with a 0 CPU head node and no workers."""
|
||||
cluster = Cluster()
|
||||
cluster.add_node(num_cpus=0)
|
||||
cluster.wait_for_nodes()
|
||||
cluster.connect()
|
||||
yield cluster
|
||||
ray.shutdown()
|
||||
cluster.shutdown()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("gpu_tasks_include_cpu", [True, False])
|
||||
def test_autoscaling_coordinator_e2e(cluster, gpu_tasks_include_cpu):
|
||||
"""Integration test for AutoscalingCoordinator.
|
||||
|
||||
This test creates 2 dummy components that request resources from
|
||||
AutoscalingCoordinator, and checks allocated resources are correct.
|
||||
"""
|
||||
object_store_memory = 100 * 1024**2
|
||||
num_cpu_nodes = 4
|
||||
cpu_node_spec = {"num_cpus": 8, "object_store_memory": object_store_memory}
|
||||
num_gpu_nodes = 2
|
||||
gpu_node_spec = {
|
||||
"num_cpus": 4,
|
||||
"num_gpus": 1,
|
||||
"object_store_memory": object_store_memory,
|
||||
}
|
||||
|
||||
for _ in range(num_cpu_nodes):
|
||||
cluster.add_node(**cpu_node_spec)
|
||||
for _ in range(num_gpu_nodes):
|
||||
cluster.add_node(**gpu_node_spec)
|
||||
|
||||
cluster.wait_for_nodes()
|
||||
|
||||
@ray.remote
|
||||
def request_and_check_resources(
|
||||
requester_id, resources, expected, request_remaining
|
||||
):
|
||||
as_coordinator = get_or_create_autoscaling_coordinator()
|
||||
ray.get(
|
||||
as_coordinator.request_resources.remote(
|
||||
requester_id=requester_id,
|
||||
resources=resources,
|
||||
expire_after_s=100,
|
||||
request_remaining=request_remaining,
|
||||
)
|
||||
)
|
||||
|
||||
def check_allocated_resources():
|
||||
allocated = ray.get(
|
||||
as_coordinator.get_allocated_resources.remote(requester_id)
|
||||
)
|
||||
allocated = [
|
||||
{
|
||||
k: int(v)
|
||||
for k, v in r.items()
|
||||
if k in ["CPU", "GPU", "object_store_memory"] and v != 0
|
||||
}
|
||||
for r in allocated
|
||||
if "node:__internal_head__" not in r
|
||||
]
|
||||
allocated = [r for r in allocated if len(r) > 0]
|
||||
if allocated != expected:
|
||||
print(
|
||||
f"{requester_id}: Allocated resources: {allocated}, "
|
||||
f"expected: {expected}. Retrying."
|
||||
)
|
||||
return False
|
||||
else:
|
||||
return True
|
||||
|
||||
wait_for_condition(
|
||||
check_allocated_resources,
|
||||
retry_interval_ms=1000,
|
||||
timeout=5,
|
||||
)
|
||||
return "ok"
|
||||
|
||||
res1_resources = [
|
||||
{
|
||||
"CPU": cpu_node_spec["num_cpus"],
|
||||
"object_store_memory": object_store_memory,
|
||||
}
|
||||
] * num_cpu_nodes
|
||||
req2_resources = [
|
||||
{
|
||||
"GPU": gpu_node_spec["num_gpus"],
|
||||
}
|
||||
] * num_gpu_nodes
|
||||
if gpu_tasks_include_cpu:
|
||||
for r in req2_resources:
|
||||
r["CPU"] = 1
|
||||
remaining = [
|
||||
{
|
||||
"CPU": gpu_node_spec["num_cpus"] - (1 if gpu_tasks_include_cpu else 0),
|
||||
"object_store_memory": object_store_memory,
|
||||
}
|
||||
] * num_gpu_nodes
|
||||
|
||||
res1 = request_and_check_resources.remote(
|
||||
requester_id="requester1",
|
||||
resources=res1_resources,
|
||||
expected=res1_resources + remaining,
|
||||
request_remaining=True,
|
||||
)
|
||||
res2 = request_and_check_resources.remote(
|
||||
requester_id="requester2",
|
||||
resources=req2_resources,
|
||||
expected=req2_resources,
|
||||
request_remaining=False,
|
||||
)
|
||||
|
||||
assert ray.get([res1, res2]) == ["ok"] * 2
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def autoscaling_coordinator_actor(ray_start_regular_shared):
|
||||
actor_cls = ray.remote(num_cpus=0)(_AutoscalingCoordinatorActor)
|
||||
actor = actor_cls.remote(
|
||||
send_resources_request=lambda b: None,
|
||||
get_cluster_nodes=lambda: [
|
||||
{"Alive": True, "Resources": {"CPU": 4}, "NodeID": "n1"}
|
||||
],
|
||||
)
|
||||
yield actor
|
||||
ray.kill(actor)
|
||||
|
||||
|
||||
def test_get_allocated_resources_eventually_consistent(autoscaling_coordinator_actor):
|
||||
"""get_allocated_resources eventually reflects a submitted request_resources call."""
|
||||
coordinator = DefaultAutoscalingCoordinator(
|
||||
"test", autoscaling_coordinator_actor=autoscaling_coordinator_actor
|
||||
)
|
||||
|
||||
coordinator.request_resources(resources=[{"CPU": 1}], expire_after_s=60)
|
||||
|
||||
wait_for_condition(
|
||||
lambda: coordinator.get_allocated_resources() == [{"CPU": 1}],
|
||||
retry_interval_ms=100,
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
|
||||
def test_get_allocated_resources_returns_cached_while_pending(
|
||||
autoscaling_coordinator_actor, monkeypatch
|
||||
):
|
||||
"""Returns the last cached value without blocking when a ref is still in-flight."""
|
||||
coordinator = DefaultAutoscalingCoordinator(
|
||||
"test", autoscaling_coordinator_actor=autoscaling_coordinator_actor
|
||||
)
|
||||
|
||||
coordinator.request_resources(resources=[{"CPU": 1}], expire_after_s=60)
|
||||
wait_for_condition(
|
||||
lambda: coordinator.get_allocated_resources() == [{"CPU": 1}],
|
||||
retry_interval_ms=100,
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
# Make ray.wait report all refs as still pending.
|
||||
def fake_wait(refs, *args, **kwargs):
|
||||
return [], refs
|
||||
|
||||
monkeypatch.setattr(ray, "wait", fake_wait)
|
||||
|
||||
coordinator.request_resources(resources=[{"CPU": 2}], expire_after_s=60)
|
||||
# Should return the stale cached value, not block.
|
||||
assert coordinator.get_allocated_resources() == [{"CPU": 1}]
|
||||
|
||||
|
||||
def test_get_allocated_resources_returns_cached_on_actor_error(
|
||||
autoscaling_coordinator_actor, monkeypatch
|
||||
):
|
||||
"""Actor errors fall back to the cached value, log a warning, and never raise.
|
||||
|
||||
Recovery is automatic: a fresh request is submitted on the next call.
|
||||
"""
|
||||
coordinator = DefaultAutoscalingCoordinator(
|
||||
"test", autoscaling_coordinator_actor=autoscaling_coordinator_actor
|
||||
)
|
||||
|
||||
coordinator.request_resources(resources=[{"CPU": 1}], expire_after_s=60)
|
||||
wait_for_condition(
|
||||
lambda: coordinator.get_allocated_resources() == [{"CPU": 1}],
|
||||
retry_interval_ms=100,
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
def fake_wait(refs, *args, **kwargs):
|
||||
# Report the ref as ready so ray.get is attempted.
|
||||
return refs, []
|
||||
|
||||
monkeypatch.setattr(ray, "wait", fake_wait)
|
||||
monkeypatch.setattr(ray, "get", Mock(side_effect=ray.exceptions.RayActorError()))
|
||||
|
||||
# Must return the last cached value, not raise.
|
||||
assert coordinator.get_allocated_resources() == [{"CPU": 1}]
|
||||
|
||||
# Recovery: submit a new request after the error and verify it eventually
|
||||
# resolves, proving the coordinator can communicate with the actor again.
|
||||
monkeypatch.undo()
|
||||
coordinator.request_resources(resources=[{"CPU": 2}], expire_after_s=60)
|
||||
wait_for_condition(
|
||||
lambda: coordinator.get_allocated_resources() == [{"CPU": 2}],
|
||||
retry_interval_ms=100,
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
|
||||
def test_cancel_request_makes_get_return_empty(autoscaling_coordinator_actor):
|
||||
"""After cancel_request, get_allocated_resources eventually returns []."""
|
||||
coordinator = DefaultAutoscalingCoordinator(
|
||||
"test", autoscaling_coordinator_actor=autoscaling_coordinator_actor
|
||||
)
|
||||
|
||||
coordinator.request_resources(resources=[{"CPU": 1}], expire_after_s=60)
|
||||
wait_for_condition(
|
||||
lambda: coordinator.get_allocated_resources() == [{"CPU": 1}],
|
||||
retry_interval_ms=100,
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
coordinator.cancel_request()
|
||||
wait_for_condition(
|
||||
lambda: coordinator.get_allocated_resources() == [],
|
||||
retry_interval_ms=100,
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
|
||||
def test_non_ray_errors_propagate(autoscaling_coordinator_actor, monkeypatch):
|
||||
"""Non-Ray errors during result consumption propagate rather than being swallowed.
|
||||
|
||||
Guards against accidentally broadening the catch from RayError to Exception.
|
||||
"""
|
||||
coordinator = DefaultAutoscalingCoordinator(
|
||||
"test", autoscaling_coordinator_actor=autoscaling_coordinator_actor
|
||||
)
|
||||
|
||||
coordinator.request_resources(resources=[{"CPU": 1}], expire_after_s=60)
|
||||
wait_for_condition(
|
||||
lambda: coordinator.get_allocated_resources() == [{"CPU": 1}],
|
||||
retry_interval_ms=100,
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(ray, "wait", lambda refs, *a, **kw: (refs, []))
|
||||
monkeypatch.setattr(
|
||||
ray, "get", Mock(side_effect=ValueError("unexpected local error"))
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="unexpected local error"):
|
||||
coordinator.get_allocated_resources()
|
||||
|
||||
|
||||
def test_coordinator_accepts_zero_resource_for_missing_resource_type(
|
||||
autoscaling_coordinator_actor,
|
||||
):
|
||||
# This is a regression test for a bug where the coordinator crashes when you request
|
||||
# a resource type (e.g., GPU: 0) that doesn't exist on the cluster.
|
||||
coordinator = DefaultAutoscalingCoordinator(
|
||||
"spam", autoscaling_coordinator_actor=autoscaling_coordinator_actor
|
||||
)
|
||||
|
||||
coordinator.request_resources(resources=[{"CPU": 1, "GPU": 0}], expire_after_s=1)
|
||||
|
||||
wait_for_condition(
|
||||
lambda: coordinator.get_allocated_resources() == [{"CPU": 1, "GPU": 0}],
|
||||
retry_interval_ms=100,
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
|
||||
def test_fractional_bundles_are_forwarded_unchanged():
|
||||
"""Fractional bundle values needs be forwarded to the autoscaler SDK as-is.
|
||||
|
||||
Previously the coordinator rounded each value up to the next integer
|
||||
before forwarding (e.g. ``{"CPU": 0.1}`` became ``{"CPU": 1}``), which
|
||||
inflated the autoscaler's demand view by up to N× when training launched
|
||||
N workers with fractional ``resources_per_worker``."""
|
||||
mock_send = Mock()
|
||||
coord = _AutoscalingCoordinatorActor(
|
||||
get_current_time=lambda: 0,
|
||||
send_resources_request=mock_send,
|
||||
get_cluster_nodes=lambda: CLUSTER_NODES_WITHOUT_HEAD,
|
||||
)
|
||||
|
||||
coord.request_resources(
|
||||
requester_id="r", resources=[{"CPU": 0.1}], expire_after_s=1
|
||||
)
|
||||
mock_send.assert_called_once_with([{"CPU": 0.1}])
|
||||
|
||||
|
||||
def test_label_selectors_are_forwarded_to_sdk():
|
||||
"""Per-bundle label_selectors are forwarded as ``bundle_label_selectors``."""
|
||||
mock_send = Mock()
|
||||
coord = _AutoscalingCoordinatorActor(
|
||||
get_current_time=lambda: 0,
|
||||
send_resources_request=mock_send,
|
||||
get_cluster_nodes=lambda: CLUSTER_NODES_WITHOUT_HEAD,
|
||||
)
|
||||
|
||||
coord.request_resources(
|
||||
requester_id="r",
|
||||
resources=[{"CPU": 1}, {"CPU": 1}],
|
||||
label_selectors=[{"instance-type": "m6i.xlarge"}, {}],
|
||||
expire_after_s=10,
|
||||
)
|
||||
mock_send.assert_called_once_with(
|
||||
[{"CPU": 1}, {"CPU": 1}],
|
||||
label_selectors=[{"instance-type": "m6i.xlarge"}, {}],
|
||||
)
|
||||
|
||||
|
||||
def test_sdk_forwarding_merges_subcluster_into_each_bundle():
|
||||
"""Forwarded bundles union the per-bundle selector with the
|
||||
requester's subcluster."""
|
||||
mock_send = Mock()
|
||||
coord = _AutoscalingCoordinatorActor(
|
||||
get_current_time=lambda: 0,
|
||||
send_resources_request=mock_send,
|
||||
get_cluster_nodes=lambda: CLUSTER_NODES_WITHOUT_HEAD,
|
||||
)
|
||||
|
||||
coord.request_resources(
|
||||
requester_id="r",
|
||||
resources=[{"CPU": 1}, {"CPU": 1}],
|
||||
label_selectors=[
|
||||
# Non-subcluster key preserved alongside the subcluster.
|
||||
{"node_id": "n1"},
|
||||
# Empty per-bundle entry — should still receive the subcluster.
|
||||
{},
|
||||
],
|
||||
subcluster_selector={"ray-subcluster": "training"},
|
||||
expire_after_s=10,
|
||||
)
|
||||
mock_send.assert_called_once_with(
|
||||
[{"CPU": 1}, {"CPU": 1}],
|
||||
label_selectors=[
|
||||
{"node_id": "n1", "ray-subcluster": "training"},
|
||||
{"ray-subcluster": "training"},
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def test_label_selectors_length_mismatch_raises():
|
||||
coord = _AutoscalingCoordinatorActor(
|
||||
get_current_time=lambda: 0,
|
||||
send_resources_request=Mock(),
|
||||
get_cluster_nodes=lambda: CLUSTER_NODES_WITHOUT_HEAD,
|
||||
)
|
||||
with pytest.raises(ValueError, match="label_selectors length"):
|
||||
coord.request_resources(
|
||||
requester_id="r",
|
||||
resources=[{"CPU": 1}, {"CPU": 1}],
|
||||
label_selectors=[{"a": "b"}],
|
||||
expire_after_s=10,
|
||||
)
|
||||
|
||||
|
||||
def test_request_rejects_per_bundle_cross_subcluster():
|
||||
"""Per-bundle subcluster values that disagree with the requester's
|
||||
``subcluster_selector`` raise."""
|
||||
coord = _AutoscalingCoordinatorActor(
|
||||
get_current_time=lambda: 0,
|
||||
send_resources_request=Mock(),
|
||||
get_cluster_nodes=lambda: CLUSTER_NODES_WITHOUT_HEAD,
|
||||
)
|
||||
with pytest.raises(ValueError, match="cross-subcluster"):
|
||||
coord.request_resources(
|
||||
requester_id="r",
|
||||
resources=[{"CPU": 1}, {"CPU": 1}],
|
||||
label_selectors=[
|
||||
{"ray-subcluster": "training"},
|
||||
{"ray-subcluster": "validation"},
|
||||
],
|
||||
subcluster_selector={"ray-subcluster": "training"},
|
||||
expire_after_s=10,
|
||||
)
|
||||
|
||||
|
||||
def test_request_rejects_changing_subcluster_selector():
|
||||
"""A requester's ``subcluster_selector`` can't change between calls;
|
||||
the rejected call must also leave the registry untouched."""
|
||||
coord = _AutoscalingCoordinatorActor(
|
||||
get_current_time=lambda: 0,
|
||||
send_resources_request=Mock(),
|
||||
get_cluster_nodes=lambda: CLUSTER_NODES_WITHOUT_HEAD,
|
||||
)
|
||||
coord.request_resources(
|
||||
requester_id="r",
|
||||
resources=[{"CPU": 1}],
|
||||
subcluster_selector={"ray-subcluster": "training"},
|
||||
expire_after_s=10,
|
||||
)
|
||||
with pytest.raises(ValueError, match="Cannot change subcluster_selector"):
|
||||
coord.request_resources(
|
||||
requester_id="r",
|
||||
resources=[{"CPU": 1}],
|
||||
subcluster_selector={"ray-subcluster": "validation"},
|
||||
expire_after_s=10,
|
||||
)
|
||||
# Registry must be unchanged after the rejected call.
|
||||
assert coord._subcluster_selectors["r"] == {"ray-subcluster": "training"}
|
||||
|
||||
|
||||
def test_label_selector_change_triggers_resend():
|
||||
"""A request whose only change is the label selector should still be
|
||||
re-sent to the autoscaler."""
|
||||
mock_send = Mock()
|
||||
coord = _AutoscalingCoordinatorActor(
|
||||
get_current_time=lambda: 0,
|
||||
send_resources_request=mock_send,
|
||||
get_cluster_nodes=lambda: CLUSTER_NODES_WITHOUT_HEAD,
|
||||
)
|
||||
|
||||
coord.request_resources(
|
||||
requester_id="r",
|
||||
resources=[{"CPU": 1}],
|
||||
label_selectors=[{"zone": "a"}],
|
||||
expire_after_s=10,
|
||||
)
|
||||
coord.request_resources(
|
||||
requester_id="r",
|
||||
resources=[{"CPU": 1}],
|
||||
label_selectors=[{"zone": "b"}],
|
||||
expire_after_s=10,
|
||||
)
|
||||
assert mock_send.call_count == 2
|
||||
mock_send.assert_called_with([{"CPU": 1}], label_selectors=[{"zone": "b"}])
|
||||
|
||||
|
||||
LABELED_CLUSTER_NODES = [
|
||||
{
|
||||
"NodeID": "n-train-1",
|
||||
"Resources": {"CPU": 8, "object_store_memory": 1000},
|
||||
"Labels": {"ray-subcluster": "training"},
|
||||
"Alive": True,
|
||||
},
|
||||
{
|
||||
"NodeID": "n-train-2",
|
||||
"Resources": {"CPU": 8, "object_store_memory": 1000},
|
||||
"Labels": {"ray-subcluster": "training"},
|
||||
"Alive": True,
|
||||
},
|
||||
{
|
||||
"NodeID": "n-val-1",
|
||||
"Resources": {"CPU": 4, "object_store_memory": 500},
|
||||
"Labels": {"ray-subcluster": "validation"},
|
||||
"Alive": True,
|
||||
},
|
||||
{
|
||||
"NodeID": "n-default-1",
|
||||
"Resources": {"CPU": 2, "object_store_memory": 200},
|
||||
"Labels": {},
|
||||
"Alive": True,
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def _make_coordinator(nodes):
|
||||
return _AutoscalingCoordinatorActor(
|
||||
get_current_time=lambda: 0,
|
||||
send_resources_request=Mock(),
|
||||
get_cluster_nodes=lambda: nodes,
|
||||
)
|
||||
|
||||
|
||||
def test_label_selector_disjoint_requesters_dont_cross_talk():
|
||||
coord = _make_coordinator(LABELED_CLUSTER_NODES)
|
||||
coord.request_resources(
|
||||
requester_id="train",
|
||||
resources=[{"CPU": 4}],
|
||||
subcluster_selector={"ray-subcluster": "training"},
|
||||
expire_after_s=10,
|
||||
request_remaining=True,
|
||||
)
|
||||
coord.request_resources(
|
||||
requester_id="val",
|
||||
resources=[{"CPU": 4}],
|
||||
subcluster_selector={"ray-subcluster": "validation"},
|
||||
expire_after_s=10,
|
||||
request_remaining=True,
|
||||
)
|
||||
|
||||
train = coord.get_allocated_resources("train")
|
||||
val = coord.get_allocated_resources("val")
|
||||
assert {"CPU": 4} in train and {"CPU": 4} in val
|
||||
# Training bucket: 2 x 8 - 4 explicit = 12 leftover, all to train.
|
||||
assert sum(a["CPU"] for a in train if a != {"CPU": 4}) == 12
|
||||
# Validation bucket: 4 - 4 explicit = 0 leftover.
|
||||
assert sum(a["CPU"] for a in val if a != {"CPU": 4}) == 0
|
||||
|
||||
|
||||
def test_unlabeled_requester_only_sees_none_bucket():
|
||||
"""An unlabeled requester is only eligible for nodes in the ``None``
|
||||
bucket (no subcluster label). It must not get explicit allocations or
|
||||
leftover share from any labeled subcluster, even when it has
|
||||
``request_remaining=True``."""
|
||||
coord = _make_coordinator(LABELED_CLUSTER_NODES)
|
||||
coord.request_resources(
|
||||
requester_id="anon",
|
||||
resources=[{"CPU": 1}],
|
||||
# No label_selector -> effective subcluster = None -> only the
|
||||
# default-labeled node (2 CPU).
|
||||
expire_after_s=10,
|
||||
request_remaining=True,
|
||||
)
|
||||
|
||||
alloc = coord.get_allocated_resources("anon")
|
||||
total_cpu = sum(a["CPU"] for a in alloc)
|
||||
assert total_cpu == 2, (
|
||||
f"unlabeled requester should only see the 2-CPU None bucket; got "
|
||||
f"{total_cpu} (alloc={alloc})"
|
||||
)
|
||||
|
||||
|
||||
def test_labeled_and_unlabeled_requesters_are_isolated():
|
||||
coord = _make_coordinator(LABELED_CLUSTER_NODES)
|
||||
coord.request_resources(
|
||||
requester_id="train",
|
||||
resources=[{"CPU": 1}],
|
||||
subcluster_selector={"ray-subcluster": "training"},
|
||||
expire_after_s=10,
|
||||
request_remaining=True,
|
||||
)
|
||||
coord.request_resources(
|
||||
requester_id="anon",
|
||||
resources=[{"CPU": 1}],
|
||||
expire_after_s=10,
|
||||
request_remaining=True,
|
||||
)
|
||||
|
||||
train_total = sum(a["CPU"] for a in coord.get_allocated_resources("train"))
|
||||
anon_total = sum(a["CPU"] for a in coord.get_allocated_resources("anon"))
|
||||
# Training bucket: 2 x 8 = 16 CPU; anon gets none of it.
|
||||
assert train_total == 16
|
||||
# Default bucket: 1 x 2 = 2 CPU; train gets none of it.
|
||||
assert anon_total == 2
|
||||
|
||||
|
||||
def test_label_selector_unmatched_yields_no_allocation():
|
||||
"""A requester whose subcluster has no matching nodes gets no
|
||||
allocation this tick."""
|
||||
coord = _make_coordinator(LABELED_CLUSTER_NODES)
|
||||
coord.request_resources(
|
||||
requester_id="ghost",
|
||||
resources=[{"CPU": 1}],
|
||||
subcluster_selector={"ray-subcluster": "nonexistent"},
|
||||
expire_after_s=10,
|
||||
request_remaining=True,
|
||||
)
|
||||
assert coord.get_allocated_resources("ghost") == []
|
||||
|
||||
|
||||
def test_label_selector_partial_fit_when_demand_exceeds_capacity():
|
||||
"""When demand exceeds capacity in the matching bucket, only the
|
||||
bundles that fit get allocated this tick."""
|
||||
coord = _make_coordinator(LABELED_CLUSTER_NODES)
|
||||
coord.request_resources(
|
||||
requester_id="val",
|
||||
resources=[{"CPU": 3}, {"CPU": 3}, {"CPU": 3}],
|
||||
subcluster_selector={"ray-subcluster": "validation"},
|
||||
expire_after_s=10,
|
||||
)
|
||||
# Validation has one 4-CPU node; only the first 3-CPU bundle fits.
|
||||
assert coord.get_allocated_resources("val") == [{"CPU": 3}]
|
||||
|
||||
|
||||
def test_full_tick_exercises_update_merge_reallocate():
|
||||
"""A `_tick()` call runs update -> merge_and_forward -> reallocate, so
|
||||
a mid-stream node-list change is picked up after the next tick."""
|
||||
nodes = [
|
||||
{
|
||||
"NodeID": "n1",
|
||||
"Resources": {"CPU": 4},
|
||||
"Labels": {"ray-subcluster": "training"},
|
||||
"Alive": True,
|
||||
},
|
||||
]
|
||||
mock_send = Mock()
|
||||
coord = _AutoscalingCoordinatorActor(
|
||||
get_current_time=lambda: 0,
|
||||
send_resources_request=mock_send,
|
||||
get_cluster_nodes=lambda: nodes,
|
||||
)
|
||||
coord.request_resources(
|
||||
requester_id="train",
|
||||
resources=[{"CPU": 1}],
|
||||
subcluster_selector={"ray-subcluster": "training"},
|
||||
expire_after_s=10,
|
||||
request_remaining=True,
|
||||
)
|
||||
# Before the join: only 4 CPU in the training bucket; 1 used explicitly,
|
||||
# 3 leftover go to train.
|
||||
train_total = sum(a["CPU"] for a in coord.get_allocated_resources("train"))
|
||||
assert train_total == 4
|
||||
|
||||
# A new training node joins the cluster.
|
||||
nodes.append(
|
||||
{
|
||||
"NodeID": "n2",
|
||||
"Resources": {"CPU": 8},
|
||||
"Labels": {"ray-subcluster": "training"},
|
||||
"Alive": True,
|
||||
}
|
||||
)
|
||||
# Without a tick, the coordinator still sees the old snapshot.
|
||||
coord._tick()
|
||||
train_total = sum(a["CPU"] for a in coord.get_allocated_resources("train"))
|
||||
# Now: 4 + 8 = 12 total; 1 explicit + 11 leftover.
|
||||
assert train_total == 12
|
||||
|
||||
|
||||
def test_labeled_requester_with_empty_resources_stays_pinned():
|
||||
"""A labeled requester with empty resources + request_remaining=True is
|
||||
eligible only for leftovers from its own subcluster."""
|
||||
coord = _make_coordinator(LABELED_CLUSTER_NODES)
|
||||
|
||||
# Idle "train" requester: no bundles, still affiliated with training
|
||||
# via the requester-wide ``label_selector``.
|
||||
coord.request_resources(
|
||||
requester_id="train_idle",
|
||||
resources=[],
|
||||
subcluster_selector={"ray-subcluster": "training"},
|
||||
expire_after_s=10,
|
||||
request_remaining=True,
|
||||
)
|
||||
# Active "val" requester: asks for 2 CPU on validation. After explicit
|
||||
# allocation, validation has 2 CPU of leftover.
|
||||
coord.request_resources(
|
||||
requester_id="val_active",
|
||||
resources=[{"CPU": 2}],
|
||||
subcluster_selector={"ray-subcluster": "validation"},
|
||||
expire_after_s=10,
|
||||
request_remaining=True,
|
||||
)
|
||||
|
||||
train_idle_alloc = coord.get_allocated_resources("train_idle")
|
||||
val_active_alloc = coord.get_allocated_resources("val_active")
|
||||
|
||||
# Training bucket: 2 x 8 = 16 CPU, all leftover, all to train_idle.
|
||||
train_idle_cpu = sum(a.get("CPU", 0) for a in train_idle_alloc)
|
||||
assert train_idle_cpu == 16, (
|
||||
f"train_idle should get exactly 16 CPU from training only, got "
|
||||
f"{train_idle_cpu} (alloc={train_idle_alloc})"
|
||||
)
|
||||
# Validation bucket: 4 - 2 explicit = 2 leftover, all to val_active.
|
||||
# Total = 2 explicit + 2 leftover = 4.
|
||||
val_active_cpu = sum(a["CPU"] for a in val_active_alloc)
|
||||
assert val_active_cpu == 4, (
|
||||
f"val_active should get 2 explicit + 2 leftover = 4 CPU, got "
|
||||
f"{val_active_cpu} (alloc={val_active_alloc})"
|
||||
)
|
||||
|
||||
|
||||
def test_proxy_forwards_label_selector_from_init():
|
||||
"""``DefaultAutoscalingCoordinator`` forwards the ``label_selector``
|
||||
it was constructed with on every request, so the actor can store the
|
||||
requester's subcluster affiliation."""
|
||||
mock_actor = Mock()
|
||||
proxy = DefaultAutoscalingCoordinator(
|
||||
requester_id="r",
|
||||
autoscaling_coordinator_actor=mock_actor,
|
||||
subcluster_selector={"ray-subcluster": "training"},
|
||||
)
|
||||
proxy.request_resources(resources=[{"CPU": 1}, {"CPU": 2}], expire_after_s=10)
|
||||
kwargs = mock_actor.request_resources.remote.call_args.kwargs
|
||||
assert kwargs["subcluster_selector"] == {"ray-subcluster": "training"}
|
||||
|
||||
|
||||
def test_proxy_forwards_label_selector_on_empty_resources():
|
||||
"""The proxy carries its ``label_selector`` even on the empty /
|
||||
registration path, so the actor keeps the requester pinned to its
|
||||
subcluster for remaining-resources eligibility."""
|
||||
mock_actor = Mock()
|
||||
proxy = DefaultAutoscalingCoordinator(
|
||||
requester_id="r",
|
||||
autoscaling_coordinator_actor=mock_actor,
|
||||
subcluster_selector={"ray-subcluster": "training"},
|
||||
)
|
||||
proxy.request_resources(resources=[], expire_after_s=10, request_remaining=True)
|
||||
kwargs = mock_actor.request_resources.remote.call_args.kwargs
|
||||
assert kwargs["resources"] == []
|
||||
assert kwargs["subcluster_selector"] == {"ray-subcluster": "training"}
|
||||
|
||||
|
||||
def test_proxy_passes_caller_label_selectors_through():
|
||||
"""If the caller passes per-bundle ``label_selectors``, the proxy
|
||||
forwards them as-is (used by callers that want per-bundle
|
||||
constraints beyond subcluster, e.g. node pins)."""
|
||||
mock_actor = Mock()
|
||||
proxy = DefaultAutoscalingCoordinator(
|
||||
requester_id="r",
|
||||
autoscaling_coordinator_actor=mock_actor,
|
||||
subcluster_selector={"ray-subcluster": "training"},
|
||||
)
|
||||
proxy.request_resources(
|
||||
resources=[{"CPU": 1}],
|
||||
label_selectors=[{"node_id": "n1"}],
|
||||
expire_after_s=10,
|
||||
)
|
||||
kwargs = mock_actor.request_resources.remote.call_args.kwargs
|
||||
assert kwargs["label_selectors"] == [{"node_id": "n1"}]
|
||||
assert kwargs["subcluster_selector"] == {"ray-subcluster": "training"}
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,382 @@
|
||||
import time
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray._private.internal_api import memory_summary
|
||||
from ray.data._internal.execution.backpressure_policy.downstream_capacity_backpressure_policy import (
|
||||
DownstreamCapacityBackpressurePolicy,
|
||||
)
|
||||
from ray.data._internal.execution.util import memory_string
|
||||
from ray.data._internal.util import MiB
|
||||
from ray.data.block import BlockMetadata
|
||||
from ray.data.datasource import Datasource, ReadTask
|
||||
from ray.data.tests.conftest import (
|
||||
CoreExecutionMetrics,
|
||||
assert_core_execution_metrics_equals,
|
||||
get_initial_core_execution_metrics_snapshot,
|
||||
restore_data_context, # noqa: F401
|
||||
)
|
||||
from ray.tests.conftest import shutdown_only # noqa: F401
|
||||
|
||||
|
||||
def test_large_e2e_backpressure_no_spilling(
|
||||
shutdown_only, restore_data_context # noqa: F811
|
||||
):
|
||||
"""Test backpressure can prevent object spilling on a synthetic large-scale
|
||||
workload."""
|
||||
# The cluster has 10 CPUs and 200MB object store memory.
|
||||
#
|
||||
# Each produce task generates 10 blocks, each of which has 10MB data.
|
||||
# In total, there will be 10 * 10 * 10MB = 1000MB intermediate data.
|
||||
#
|
||||
# `ReservationOpResourceAllocator` should dynamically allocate resources to each
|
||||
# operator and prevent object spilling.
|
||||
NUM_CPUS = 10
|
||||
NUM_ROWS_PER_TASK = 10
|
||||
NUM_TASKS = 20
|
||||
NUM_ROWS_TOTAL = NUM_ROWS_PER_TASK * NUM_TASKS
|
||||
BLOCK_SIZE = 10 * MiB
|
||||
object_store_memory = 200 * MiB
|
||||
|
||||
print(f">>> Setting Object Store to {memory_string(object_store_memory)}")
|
||||
|
||||
ray.init(num_cpus=NUM_CPUS, object_store_memory=object_store_memory)
|
||||
|
||||
def produce(batch):
|
||||
print(">>> [Producer] Produce task started", batch["id"])
|
||||
time.sleep(0.1)
|
||||
for id in batch["id"]:
|
||||
print(f">>> [Producer] Producing row {id=}")
|
||||
yield {
|
||||
"id": [id],
|
||||
"image": [np.zeros(BLOCK_SIZE, dtype=np.uint8)],
|
||||
}
|
||||
|
||||
def consume(batch):
|
||||
print(">>> [Consumer] Consume task started", batch["id"])
|
||||
time.sleep(0.01)
|
||||
return {"id": batch["id"], "result": [0 for _ in batch["id"]]}
|
||||
|
||||
data_context = ray.data.DataContext.get_current()
|
||||
data_context.execution_options.verbose_progress = True
|
||||
data_context.target_max_block_size = BLOCK_SIZE
|
||||
|
||||
last_snapshot = get_initial_core_execution_metrics_snapshot()
|
||||
|
||||
ds = ray.data.range(NUM_ROWS_TOTAL, override_num_blocks=NUM_TASKS)
|
||||
ds = ds.map_batches(produce, batch_size=NUM_ROWS_PER_TASK)
|
||||
ds = ds.map_batches(consume, batch_size=None, num_cpus=0.9)
|
||||
# Check core execution metrics every 10 rows, because it's expensive.
|
||||
for _ in ds.iter_batches(batch_size=NUM_ROWS_PER_TASK):
|
||||
last_snapshot = assert_core_execution_metrics_equals(
|
||||
CoreExecutionMetrics(
|
||||
object_store_stats={
|
||||
"spilled_bytes_total": 0,
|
||||
"restored_bytes_total": 0,
|
||||
},
|
||||
),
|
||||
last_snapshot,
|
||||
)
|
||||
|
||||
|
||||
def _build_dataset(
|
||||
obj_store_limit,
|
||||
producer_num_cpus,
|
||||
consumer_num_cpus,
|
||||
num_blocks,
|
||||
block_size,
|
||||
insert_limit_op=False,
|
||||
):
|
||||
# Create a dataset with 2 operators:
|
||||
# - The producer op has only 1 task, which produces `num_blocks` blocks, each
|
||||
# of which has `block_size` data.
|
||||
# - The consumer op has `num_blocks` tasks, each of which consumes 1 block.
|
||||
ctx = ray.data.DataContext.get_current()
|
||||
ctx.target_max_block_size = block_size
|
||||
ctx.execution_options.resource_limits = ctx.execution_options.resource_limits.copy(
|
||||
object_store_memory=obj_store_limit
|
||||
)
|
||||
|
||||
def producer(batch):
|
||||
for i in range(num_blocks):
|
||||
print(f"[{time.time()}] Producing block #{i} ({block_size=})")
|
||||
yield {
|
||||
"id": [i],
|
||||
"data": [np.zeros(block_size, dtype=np.uint8)],
|
||||
}
|
||||
|
||||
def consumer(batch):
|
||||
assert len(batch["id"]) == 1
|
||||
print(f"[{time.time()}] Consuming block #{batch['id'][0]}")
|
||||
time.sleep(0.01)
|
||||
del batch["data"]
|
||||
return batch
|
||||
|
||||
ds = ray.data.range(1, override_num_blocks=1).materialize()
|
||||
ds = ds.map_batches(producer, batch_size=None, num_cpus=producer_num_cpus)
|
||||
# Add a limit op in the middle, to test that ReservationOpResourceAllocator
|
||||
# will account limit op's resource usage to the previous producer map op.
|
||||
if insert_limit_op:
|
||||
ds = ds.limit(num_blocks)
|
||||
ds = ds.map_batches(consumer, batch_size=None, num_cpus=consumer_num_cpus)
|
||||
if insert_limit_op:
|
||||
ds = ds.limit(num_blocks)
|
||||
return ds
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"cluster_cpus, cluster_obj_store_mem_mb",
|
||||
[
|
||||
(3, 500), # CPU not enough
|
||||
(4, 100), # Object store memory not enough
|
||||
(3, 100), # Both not enough
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("insert_limit_op", [False, True])
|
||||
def test_no_deadlock_on_small_cluster_resources(
|
||||
cluster_cpus,
|
||||
cluster_obj_store_mem_mb,
|
||||
insert_limit_op,
|
||||
shutdown_only, # noqa: F811
|
||||
restore_data_context, # noqa: F811
|
||||
):
|
||||
"""Test when cluster resources are not enough for launching one task per op,
|
||||
the execution can still proceed without deadlock.
|
||||
"""
|
||||
cluster_obj_store_mem_mb *= 1024**2
|
||||
ray.init(num_cpus=cluster_cpus, object_store_memory=cluster_obj_store_mem_mb)
|
||||
num_blocks = 10
|
||||
block_size = 100 * 1024 * 1024
|
||||
ds = _build_dataset(
|
||||
obj_store_limit=cluster_obj_store_mem_mb // 2,
|
||||
producer_num_cpus=3,
|
||||
consumer_num_cpus=1,
|
||||
num_blocks=num_blocks,
|
||||
block_size=block_size,
|
||||
insert_limit_op=insert_limit_op,
|
||||
)
|
||||
assert len(ds.take_all()) == num_blocks
|
||||
|
||||
|
||||
@pytest.mark.parametrize("insert_limit_op", [False, True])
|
||||
def test_no_deadlock_on_resource_contention(
|
||||
insert_limit_op, shutdown_only, restore_data_context # noqa: F811
|
||||
):
|
||||
"""Test when resources are preempted by non-Data code, the execution can
|
||||
still proceed without deadlock."""
|
||||
cluster_obj_store_mem = 1000 * 1024 * 1024
|
||||
ray.init(num_cpus=5, object_store_memory=cluster_obj_store_mem)
|
||||
# Create a non-Data actor that uses 4 CPUs, only 1 CPU
|
||||
# is left for Data. Currently Data StreamExecutor still
|
||||
# incorrectly assumes it has all the 5 CPUs.
|
||||
# Check that we don't deadlock in this case.
|
||||
|
||||
@ray.remote(num_cpus=4)
|
||||
class DummyActor:
|
||||
def foo(self):
|
||||
return None
|
||||
|
||||
dummy_actor = DummyActor.remote()
|
||||
ray.get(dummy_actor.foo.remote())
|
||||
|
||||
num_blocks = 10
|
||||
block_size = 50 * 1024 * 1024
|
||||
ds = _build_dataset(
|
||||
obj_store_limit=cluster_obj_store_mem // 2,
|
||||
producer_num_cpus=1,
|
||||
consumer_num_cpus=0.9,
|
||||
num_blocks=num_blocks,
|
||||
block_size=block_size,
|
||||
insert_limit_op=insert_limit_op,
|
||||
)
|
||||
|
||||
from ray.data._internal.execution.streaming_executor_state import IdleDetector
|
||||
|
||||
with patch.object(IdleDetector, "DETECTION_INTERVAL_S", 0.1):
|
||||
assert len(ds.take_all()) == num_blocks
|
||||
|
||||
|
||||
def test_no_deadlock_when_downstream_capacity_policy_zeros_limit(
|
||||
shutdown_only, restore_data_context # noqa: F811
|
||||
):
|
||||
"""Test when DownstreamCapacityBackpressurePolicy zeros the output limit,
|
||||
the execution can still proceed without deadlock."""
|
||||
cluster_obj_store_mem = 100 * MiB
|
||||
ray.init(num_cpus=2, object_store_memory=cluster_obj_store_mem)
|
||||
|
||||
num_blocks = 20
|
||||
block_size = 1 * MiB
|
||||
ds = _build_dataset(
|
||||
obj_store_limit=cluster_obj_store_mem // 2,
|
||||
producer_num_cpus=1,
|
||||
consumer_num_cpus=1,
|
||||
num_blocks=num_blocks,
|
||||
block_size=block_size,
|
||||
)
|
||||
|
||||
# Force DownstreamCapacityBackpressurePolicy to always return 0 to trigger unblock
|
||||
with patch.object(
|
||||
DownstreamCapacityBackpressurePolicy,
|
||||
"max_task_output_bytes_to_read",
|
||||
lambda self, op: 0,
|
||||
):
|
||||
# Without the escape hatch firing, this would hang.
|
||||
assert len(ds.take_all()) == num_blocks
|
||||
|
||||
|
||||
def test_no_deadlock_with_preserve_order(
|
||||
restore_data_context, shutdown_only # noqa: F811
|
||||
):
|
||||
"""Test backpressure won't cause deadlocks when `preserve_order=True`."""
|
||||
num_blocks = 20
|
||||
block_size = 10 * 1024 * 1024
|
||||
ray.init(num_cpus=num_blocks)
|
||||
data_context = ray.data.DataContext.get_current()
|
||||
data_context.target_max_block_size = block_size
|
||||
data_context._max_num_blocks_in_streaming_gen_buffer = 1
|
||||
data_context.execution_options.preserve_order = True
|
||||
data_context.execution_options.resource_limits = (
|
||||
data_context.execution_options.resource_limits.copy(
|
||||
object_store_memory=5 * block_size
|
||||
)
|
||||
)
|
||||
|
||||
# Some tasks are slower than others.
|
||||
# The faster tasks will finish first and occupy Map op's internal output buffer.
|
||||
# Test that we won't backpressure the operator in this case.
|
||||
def map_fn(batch):
|
||||
idx = batch["id"][0]
|
||||
print("map_fn", idx, time.time())
|
||||
if idx % 2 == 0:
|
||||
time.sleep(3)
|
||||
batch["data"] = [np.zeros(block_size, dtype=np.uint8)]
|
||||
return batch
|
||||
|
||||
ds = ray.data.range(num_blocks, override_num_blocks=num_blocks)
|
||||
ds = ds.map_batches(map_fn, batch_size=None, num_cpus=1)
|
||||
assert len(ds.take_all()) == num_blocks
|
||||
|
||||
|
||||
def test_input_backpressure_e2e(restore_data_context, shutdown_only): # noqa: F811
|
||||
# Tests that backpressure applies even when reading directly from the input
|
||||
# datasource. This relies on datasource metadata size estimation.
|
||||
|
||||
@ray.remote
|
||||
class Counter:
|
||||
def __init__(self):
|
||||
self.count = 0
|
||||
|
||||
def increment(self):
|
||||
self.count += 1
|
||||
|
||||
def get(self):
|
||||
return self.count
|
||||
|
||||
def reset(self):
|
||||
self.count = 0
|
||||
|
||||
class CountingRangeDatasource(Datasource):
|
||||
def __init__(self):
|
||||
self.counter = Counter.remote()
|
||||
|
||||
def prepare_read(self, parallelism):
|
||||
# Use 50 MiB blocks to exceed the 25 MiB output reservation
|
||||
# and trigger object store backpressure
|
||||
num_bytes = 50 * MiB
|
||||
|
||||
def range_(i):
|
||||
print(f">>> Read task: {i=}")
|
||||
|
||||
ray.get(self.counter.increment.remote())
|
||||
return [pd.DataFrame({"data": np.ones((num_bytes,), dtype=np.uint8)})]
|
||||
|
||||
print(f">>> Block size: {num_bytes}")
|
||||
|
||||
return [
|
||||
ReadTask(
|
||||
lambda i=i: range_(i),
|
||||
BlockMetadata(
|
||||
num_rows=1,
|
||||
size_bytes=num_bytes,
|
||||
input_files=None,
|
||||
exec_stats=None,
|
||||
),
|
||||
)
|
||||
for i in range(parallelism)
|
||||
]
|
||||
|
||||
source = CountingRangeDatasource()
|
||||
|
||||
ctx = ray.data.DataContext.get_current()
|
||||
ctx.execution_options.resource_limits = ctx.execution_options.resource_limits.copy(
|
||||
object_store_memory=100 * MiB,
|
||||
cpu=1,
|
||||
)
|
||||
ctx.target_max_block_size = 50 * MiB
|
||||
|
||||
# Create dataset with many blocks
|
||||
ds = ray.data.read_datasource(source, override_num_blocks=1000)
|
||||
it = iter(ds.iter_internal_ref_bundles())
|
||||
# Dequeue 1 block
|
||||
next(it)
|
||||
# Let it bake for some time
|
||||
time.sleep(3)
|
||||
|
||||
launched = ray.get(source.counter.get.remote())
|
||||
|
||||
# Clean up
|
||||
del it
|
||||
# With 50 MiB blocks and 100 MiB limit, backpressure should limit to ~2 tasks
|
||||
# because after 2 outputs (100 MiB), the budget is depleted
|
||||
assert launched == 2, launched
|
||||
|
||||
|
||||
def test_streaming_backpressure_e2e(
|
||||
shutdown_only, monkeypatch, restore_data_context # noqa: F811
|
||||
):
|
||||
# This test case is particularly challenging since there is a large input->output
|
||||
# increase in data size: https://github.com/ray-project/ray/issues/34041
|
||||
|
||||
# Increase the Ray Core spilling threshold to 100% to avoid flakiness.
|
||||
monkeypatch.setenv("RAY_object_spilling_threshold", "1")
|
||||
|
||||
class TestSlow:
|
||||
def __call__(self, df: np.ndarray):
|
||||
time.sleep(2)
|
||||
return {"id": np.random.randn(1, 20, 1024, 1024)}
|
||||
|
||||
class TestFast:
|
||||
def __call__(self, df: np.ndarray):
|
||||
time.sleep(0.5)
|
||||
return {"id": np.random.randn(1, 20, 1024, 1024)}
|
||||
|
||||
ctx = ray.init(object_store_memory=4e9)
|
||||
ds = ray.data.range_tensor(20, shape=(3, 1024, 1024), override_num_blocks=20)
|
||||
|
||||
pipe = ds.map_batches(
|
||||
TestFast,
|
||||
batch_size=1,
|
||||
num_cpus=0.5,
|
||||
compute=ray.data.ActorPoolStrategy(size=2),
|
||||
).map_batches(
|
||||
TestSlow,
|
||||
batch_size=1,
|
||||
compute=ray.data.ActorPoolStrategy(size=1),
|
||||
)
|
||||
|
||||
for batch in pipe.iter_batches(batch_size=1, prefetch_batches=2):
|
||||
...
|
||||
|
||||
# If backpressure is not working right, we will spill.
|
||||
meminfo = memory_summary(ctx.address_info["address"], stats_only=True)
|
||||
assert "Spilled" not in meminfo, meminfo
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,476 @@
|
||||
import functools
|
||||
import math
|
||||
import time
|
||||
import types
|
||||
import unittest
|
||||
from collections import defaultdict
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data._internal.execution.backpressure_policy import (
|
||||
ENABLED_BACKPRESSURE_POLICIES_CONFIG_KEY,
|
||||
ConcurrencyCapBackpressurePolicy,
|
||||
)
|
||||
from ray.data._internal.execution.operators.input_data_buffer import InputDataBuffer
|
||||
from ray.data._internal.execution.operators.task_pool_map_operator import (
|
||||
TaskPoolMapOperator,
|
||||
)
|
||||
from ray.data._internal.execution.resource_manager import ResourceManager
|
||||
from ray.data.context import DataContext
|
||||
from ray.data.tests.conftest import mock_all_to_all_op
|
||||
from ray.util.annotations import RayDeprecationWarning
|
||||
|
||||
|
||||
class TestConcurrencyCapBackpressurePolicy(unittest.TestCase):
|
||||
"""Tests for ConcurrencyCapBackpressurePolicy."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls._cluster_cpus = 10
|
||||
ray.init(num_cpus=cls._cluster_cpus)
|
||||
data_context = ray.data.DataContext.get_current()
|
||||
data_context.set_config(
|
||||
ENABLED_BACKPRESSURE_POLICIES_CONFIG_KEY,
|
||||
[ConcurrencyCapBackpressurePolicy],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
ray.shutdown()
|
||||
data_context = ray.data.DataContext.get_current()
|
||||
data_context.remove_config(ENABLED_BACKPRESSURE_POLICIES_CONFIG_KEY)
|
||||
|
||||
def _mock_resource_manager(self):
|
||||
"""Helper to create a resource manager mock with real method bindings."""
|
||||
rm = MagicMock()
|
||||
rm.is_op_eligible = types.MethodType(ResourceManager.is_op_eligible, rm)
|
||||
rm._get_downstream_ineligible_ops = types.MethodType(
|
||||
ResourceManager._get_downstream_ineligible_ops, rm
|
||||
)
|
||||
rm._is_blocking_materializing_op = types.MethodType(
|
||||
ResourceManager._is_blocking_materializing_op, rm
|
||||
)
|
||||
return rm
|
||||
|
||||
def test_basic(self):
|
||||
concurrency = 16
|
||||
input_op = InputDataBuffer(DataContext.get_current(), input_data=[MagicMock()])
|
||||
map_op_no_concurrency = TaskPoolMapOperator(
|
||||
map_transformer=MagicMock(),
|
||||
data_context=DataContext.get_current(),
|
||||
input_op=input_op,
|
||||
)
|
||||
map_op = TaskPoolMapOperator(
|
||||
map_transformer=MagicMock(),
|
||||
data_context=DataContext.get_current(),
|
||||
input_op=map_op_no_concurrency,
|
||||
max_concurrency=concurrency,
|
||||
)
|
||||
map_op.metrics.num_tasks_running = 0
|
||||
map_op.metrics.num_tasks_finished = 0
|
||||
topology = {
|
||||
map_op: MagicMock(),
|
||||
input_op: MagicMock(),
|
||||
map_op_no_concurrency: MagicMock(),
|
||||
}
|
||||
|
||||
mock_resource_manager = MagicMock()
|
||||
# Return None to skip dynamic output queue size backpressure check
|
||||
mock_resource_manager.get_op_usage.return_value = None
|
||||
mock_resource_manager.get_budget.return_value = None
|
||||
mock_resource_manager.is_op_eligible.return_value = False
|
||||
|
||||
policy = ConcurrencyCapBackpressurePolicy(
|
||||
DataContext.get_current(),
|
||||
topology,
|
||||
mock_resource_manager,
|
||||
)
|
||||
|
||||
self.assertEqual(policy._concurrency_caps[map_op], concurrency)
|
||||
self.assertTrue(math.isinf(policy._concurrency_caps[input_op]))
|
||||
self.assertTrue(math.isinf(policy._concurrency_caps[map_op_no_concurrency]))
|
||||
|
||||
# Gradually increase num_tasks_running to the cap.
|
||||
for i in range(1, concurrency + 1):
|
||||
self.assertTrue(policy.can_add_input(map_op))
|
||||
map_op.metrics.num_tasks_running = i
|
||||
# Now num_tasks_running reaches the cap, so can_add_input should return False.
|
||||
self.assertFalse(policy.can_add_input(map_op))
|
||||
|
||||
map_op.metrics.num_tasks_running = concurrency / 2
|
||||
self.assertEqual(policy.can_add_input(map_op), True)
|
||||
|
||||
def _create_record_time_actor(self):
|
||||
@ray.remote(num_cpus=0)
|
||||
class RecordTimeActor:
|
||||
def __init__(self):
|
||||
self._start_time = defaultdict(lambda: [])
|
||||
self._end_time = defaultdict(lambda: [])
|
||||
|
||||
def record_start_time(self, index):
|
||||
self._start_time[index].append(time.time())
|
||||
|
||||
def record_end_time(self, index):
|
||||
self._end_time[index].append(time.time())
|
||||
|
||||
def get_start_and_end_time_for_op(self, index):
|
||||
return min(self._start_time[index]), max(self._end_time[index])
|
||||
|
||||
def get_start_and_end_time_for_all_tasks_of_op(self, index):
|
||||
return self._start_time[index], self._end_time[index]
|
||||
|
||||
actor = RecordTimeActor.remote()
|
||||
return actor
|
||||
|
||||
def _get_map_func(self, actor, index):
|
||||
def map_func(data, actor, index):
|
||||
actor.record_start_time.remote(index)
|
||||
yield data
|
||||
actor.record_end_time.remote(index)
|
||||
|
||||
return functools.partial(map_func, actor=actor, index=index)
|
||||
|
||||
def test_e2e_normal(self):
|
||||
"""A simple E2E test with ConcurrencyCapBackpressurePolicy enabled."""
|
||||
actor = self._create_record_time_actor()
|
||||
map_func1 = self._get_map_func(actor, 1)
|
||||
map_func2 = self._get_map_func(actor, 2)
|
||||
|
||||
# Create a dataset with 2 map ops. Each map op has N tasks, where N is
|
||||
# the number of cluster CPUs.
|
||||
N = self.__class__._cluster_cpus
|
||||
ds = ray.data.range(N, override_num_blocks=N)
|
||||
# Use different `num_cpus` to make sure they don't fuse.
|
||||
ds = ds.map_batches(map_func1, batch_size=None, num_cpus=1, concurrency=1)
|
||||
ds = ds.map_batches(map_func2, batch_size=None, num_cpus=1.1, concurrency=1)
|
||||
res = ds.take_all()
|
||||
self.assertEqual(len(res), N)
|
||||
|
||||
# We recorded the start and end time of each op,
|
||||
# check that these 2 ops are executed interleavingly.
|
||||
# This means that the executor didn't allocate all resources to the first
|
||||
# op in the beginning.
|
||||
start1, end1 = ray.get(actor.get_start_and_end_time_for_op.remote(1))
|
||||
start2, end2 = ray.get(actor.get_start_and_end_time_for_op.remote(2))
|
||||
assert start1 < start2 < end1 < end2, (start1, start2, end1, end2)
|
||||
|
||||
def test_can_add_input_with_dynamic_output_queue_size_backpressure_disabled(self):
|
||||
"""Test can_add_input when dynamic output queue size backpressure is disabled."""
|
||||
input_op = InputDataBuffer(DataContext.get_current(), input_data=[MagicMock()])
|
||||
map_op = TaskPoolMapOperator(
|
||||
map_transformer=MagicMock(),
|
||||
data_context=DataContext.get_current(),
|
||||
input_op=input_op,
|
||||
max_concurrency=5,
|
||||
)
|
||||
map_op.metrics.num_tasks_running = 3
|
||||
|
||||
topology = {map_op: MagicMock(), input_op: MagicMock()}
|
||||
|
||||
# Create policy with dynamic output queue size backpressure disabled
|
||||
policy = ConcurrencyCapBackpressurePolicy(
|
||||
DataContext.get_current(),
|
||||
topology,
|
||||
MagicMock(), # resource_manager
|
||||
)
|
||||
policy.enable_dynamic_output_queue_size_backpressure = False
|
||||
|
||||
# Should only check against configured concurrency cap
|
||||
self.assertTrue(policy.can_add_input(map_op)) # 3 < 5
|
||||
|
||||
map_op.metrics.num_tasks_running = 5
|
||||
self.assertFalse(policy.can_add_input(map_op)) # 5 >= 5
|
||||
|
||||
def test_can_add_input_with_non_map_operator(self):
|
||||
"""Test can_add_input with non-MapOperator (should use basic cap check)."""
|
||||
input_op = InputDataBuffer(DataContext.get_current(), input_data=[MagicMock()])
|
||||
input_op.metrics.num_tasks_running = 1
|
||||
|
||||
topology = {input_op: MagicMock()}
|
||||
|
||||
policy = ConcurrencyCapBackpressurePolicy(
|
||||
DataContext.get_current(),
|
||||
topology,
|
||||
MagicMock(), # resource_manager
|
||||
)
|
||||
|
||||
# InputDataBuffer has infinite concurrency cap, so should always allow
|
||||
self.assertTrue(policy.can_add_input(input_op))
|
||||
|
||||
def test_can_add_input_with_ineligible_op(self):
|
||||
"""Test can_add_input when op is not eligible for backpressure."""
|
||||
input_op = InputDataBuffer(DataContext.get_current(), input_data=[MagicMock()])
|
||||
map_op = TaskPoolMapOperator(
|
||||
map_transformer=MagicMock(),
|
||||
data_context=DataContext.get_current(),
|
||||
input_op=input_op,
|
||||
max_concurrency=5,
|
||||
)
|
||||
map_op.metrics.num_tasks_running = 3
|
||||
|
||||
topology = {map_op: MagicMock(), input_op: MagicMock()}
|
||||
|
||||
mock_resource_manager = self._mock_resource_manager()
|
||||
# Override to test policy behavior when op is not eligible
|
||||
mock_resource_manager.is_op_eligible = MagicMock(return_value=False)
|
||||
|
||||
policy = ConcurrencyCapBackpressurePolicy(
|
||||
DataContext.get_current(),
|
||||
topology,
|
||||
mock_resource_manager,
|
||||
)
|
||||
policy.enable_dynamic_output_queue_size_backpressure = True
|
||||
|
||||
# Should skip dynamic backpressure and use basic cap check
|
||||
self.assertTrue(policy.can_add_input(map_op)) # 3 < 5
|
||||
|
||||
map_op.metrics.num_tasks_running = 5
|
||||
self.assertFalse(policy.can_add_input(map_op)) # 5 >= 5
|
||||
|
||||
def test_can_add_input_with_materializing_downstream_op(self):
|
||||
"""Test can_add_input when downstream op is a materializing operator."""
|
||||
input_op = InputDataBuffer(DataContext.get_current(), input_data=[MagicMock()])
|
||||
map_op = TaskPoolMapOperator(
|
||||
map_transformer=MagicMock(),
|
||||
data_context=DataContext.get_current(),
|
||||
input_op=input_op,
|
||||
max_concurrency=5,
|
||||
)
|
||||
map_op.metrics.num_tasks_running = 3
|
||||
|
||||
# Create materializing downstream op (automatically adds to map_op._output_dependencies)
|
||||
mock_all_to_all_op(map_op)
|
||||
|
||||
topology = {map_op: MagicMock(), input_op: MagicMock()}
|
||||
|
||||
mock_resource_manager = self._mock_resource_manager()
|
||||
|
||||
policy = ConcurrencyCapBackpressurePolicy(
|
||||
DataContext.get_current(),
|
||||
topology,
|
||||
mock_resource_manager,
|
||||
)
|
||||
policy.enable_dynamic_output_queue_size_backpressure = True
|
||||
|
||||
# Should skip dynamic backpressure and use basic cap check
|
||||
# to avoid starving materializing operators
|
||||
self.assertTrue(policy.can_add_input(map_op)) # 3 < 5
|
||||
|
||||
map_op.metrics.num_tasks_running = 5
|
||||
self.assertFalse(policy.can_add_input(map_op)) # 5 >= 5
|
||||
|
||||
@patch(
|
||||
"ray.data._internal.execution.backpressure_policy."
|
||||
"concurrency_cap_backpressure_policy.get_available_object_store_budget_fraction"
|
||||
)
|
||||
def test_can_add_input_with_object_store_memory_usage_ratio_above_threshold(
|
||||
self, mock_get_budget_fraction
|
||||
):
|
||||
"""Test can_add_input when object store memory usage ratio is above threshold."""
|
||||
input_op = InputDataBuffer(DataContext.get_current(), input_data=[MagicMock()])
|
||||
map_op = TaskPoolMapOperator(
|
||||
map_transformer=MagicMock(),
|
||||
data_context=DataContext.get_current(),
|
||||
input_op=input_op,
|
||||
max_concurrency=5,
|
||||
)
|
||||
map_op.metrics.num_tasks_running = 3
|
||||
|
||||
topology = {map_op: MagicMock(), input_op: MagicMock()}
|
||||
|
||||
mock_resource_manager = self._mock_resource_manager()
|
||||
|
||||
# Mock available object store memory budget fraction above threshold to skip dynamic backpressure
|
||||
threshold = (
|
||||
ConcurrencyCapBackpressurePolicy.AVAILABLE_OBJECT_STORE_BUDGET_THRESHOLD
|
||||
)
|
||||
# Set fraction above threshold to skip dynamic backpressure
|
||||
mock_get_budget_fraction.return_value = threshold + 0.05
|
||||
|
||||
policy = ConcurrencyCapBackpressurePolicy(
|
||||
DataContext.get_current(),
|
||||
topology,
|
||||
mock_resource_manager,
|
||||
)
|
||||
policy.enable_dynamic_output_queue_size_backpressure = True
|
||||
|
||||
# Initialize EWMA state to verify it's not updated when ratio > threshold
|
||||
initial_level = 100.0
|
||||
initial_dev = 20.0
|
||||
policy._q_level_nbytes[map_op] = initial_level
|
||||
policy._q_level_dev[map_op] = initial_dev
|
||||
|
||||
# Should skip dynamic backpressure and use basic cap check
|
||||
# EWMA state should not be updated (early return)
|
||||
self.assertTrue(policy.can_add_input(map_op)) # 3 < 5
|
||||
self.assertEqual(policy._q_level_nbytes[map_op], initial_level)
|
||||
self.assertEqual(policy._q_level_dev[map_op], initial_dev)
|
||||
|
||||
map_op.metrics.num_tasks_running = 5
|
||||
self.assertFalse(policy.can_add_input(map_op)) # 5 >= 5
|
||||
# EWMA state should still not be updated
|
||||
self.assertEqual(policy._q_level_nbytes[map_op], initial_level)
|
||||
self.assertEqual(policy._q_level_dev[map_op], initial_dev)
|
||||
|
||||
@patch(
|
||||
"ray.data._internal.execution.backpressure_policy."
|
||||
"concurrency_cap_backpressure_policy.get_available_object_store_budget_fraction"
|
||||
)
|
||||
def test_can_add_input_with_object_store_memory_usage_ratio_below_threshold(
|
||||
self, mock_get_budget_fraction
|
||||
):
|
||||
"""Test can_add_input when object store memory usage ratio is below threshold."""
|
||||
input_op = InputDataBuffer(DataContext.get_current(), input_data=[MagicMock()])
|
||||
map_op = TaskPoolMapOperator(
|
||||
map_transformer=MagicMock(),
|
||||
data_context=DataContext.get_current(),
|
||||
input_op=input_op,
|
||||
max_concurrency=5,
|
||||
)
|
||||
map_op.metrics.num_tasks_running = 3
|
||||
|
||||
topology = {map_op: MagicMock(), input_op: MagicMock()}
|
||||
|
||||
mock_resource_manager = self._mock_resource_manager()
|
||||
|
||||
# Mock available object store memory budget fraction below threshold to apply dynamic backpressure
|
||||
threshold = (
|
||||
ConcurrencyCapBackpressurePolicy.AVAILABLE_OBJECT_STORE_BUDGET_THRESHOLD
|
||||
)
|
||||
# Set fraction below threshold to apply dynamic backpressure
|
||||
mock_get_budget_fraction.return_value = threshold - 0.05
|
||||
|
||||
# Mock queue size methods
|
||||
mock_resource_manager.get_mem_op_internal.return_value = 100
|
||||
mock_resource_manager.get_mem_op_outputs.return_value = 200
|
||||
|
||||
policy = ConcurrencyCapBackpressurePolicy(
|
||||
DataContext.get_current(),
|
||||
topology,
|
||||
mock_resource_manager,
|
||||
)
|
||||
policy.enable_dynamic_output_queue_size_backpressure = True
|
||||
|
||||
# Should proceed with dynamic backpressure logic
|
||||
# Initialize EWMA state for the operator with a different level
|
||||
# so we can verify the update happens (queue size is 300)
|
||||
initial_level = 200.0
|
||||
initial_dev = 50.0
|
||||
policy._q_level_nbytes[map_op] = initial_level
|
||||
policy._q_level_dev[map_op] = initial_dev
|
||||
|
||||
result = policy.can_add_input(map_op)
|
||||
# With queue size 300, initial level=200, dev=50, bounds=[150, 250]
|
||||
# Queue size 300 is above the upper bound, so should backoff.
|
||||
# running=3, backoff by 1 -> effective_cap=2
|
||||
# running=3 < effective_cap=2 should be False
|
||||
self.assertFalse(result)
|
||||
# EWMA state should be updated when ratio < threshold
|
||||
# Level should move toward 300 (queue size)
|
||||
self.assertNotEqual(policy._q_level_nbytes[map_op], initial_level)
|
||||
# Dev should also be updated
|
||||
self.assertNotEqual(policy._q_level_dev[map_op], initial_dev)
|
||||
|
||||
@patch(
|
||||
"ray.data._internal.execution.backpressure_policy."
|
||||
"concurrency_cap_backpressure_policy.get_available_object_store_budget_fraction"
|
||||
)
|
||||
def test_can_add_input_effective_cap_calculation(self, mock_get_budget_fraction):
|
||||
"""Test that effective cap calculation works correctly with different queue sizes."""
|
||||
input_op = InputDataBuffer(DataContext.get_current(), input_data=[MagicMock()])
|
||||
map_op = TaskPoolMapOperator(
|
||||
map_transformer=MagicMock(),
|
||||
data_context=DataContext.get_current(),
|
||||
input_op=input_op,
|
||||
max_concurrency=8,
|
||||
)
|
||||
map_op.metrics.num_tasks_running = 4
|
||||
|
||||
topology = {map_op: MagicMock(), input_op: MagicMock()}
|
||||
|
||||
mock_resource_manager = self._mock_resource_manager()
|
||||
threshold = (
|
||||
ConcurrencyCapBackpressurePolicy.AVAILABLE_OBJECT_STORE_BUDGET_THRESHOLD
|
||||
)
|
||||
# Set fraction below threshold to apply dynamic backpressure
|
||||
mock_get_budget_fraction.return_value = threshold - 0.05
|
||||
|
||||
policy = ConcurrencyCapBackpressurePolicy(
|
||||
DataContext.get_current(),
|
||||
topology,
|
||||
mock_resource_manager,
|
||||
)
|
||||
policy.enable_dynamic_output_queue_size_backpressure = True
|
||||
|
||||
# Test different queue sizes using policy constants
|
||||
test_cases = [
|
||||
# (internal_usage, downstream_usage, level, dev, expected_result, description)
|
||||
(
|
||||
50,
|
||||
50,
|
||||
5000.0,
|
||||
200.0,
|
||||
True,
|
||||
"low_queue_below_lower_bound",
|
||||
), # 100 < 5000 - 2*200 = 4600, ramp up
|
||||
(
|
||||
200,
|
||||
200,
|
||||
400.0,
|
||||
50.0,
|
||||
False,
|
||||
"medium_queue_in_hold_region",
|
||||
), # 400 in [300, 500], hold
|
||||
(
|
||||
300,
|
||||
300,
|
||||
200.0,
|
||||
50.0,
|
||||
False,
|
||||
"high_queue_above_upper_bound",
|
||||
), # 600 > 200 + 2*50 = 300, backoff
|
||||
]
|
||||
|
||||
for (
|
||||
internal_usage,
|
||||
downstream_usage,
|
||||
level,
|
||||
dev,
|
||||
expected_result,
|
||||
description,
|
||||
) in test_cases:
|
||||
with self.subTest(description=description):
|
||||
mock_resource_manager.get_mem_op_internal.return_value = internal_usage
|
||||
mock_resource_manager.get_mem_op_outputs.return_value = downstream_usage
|
||||
mock_resource_manager.get_op_outputs_object_store_usage_with_downstream.return_value = (
|
||||
downstream_usage
|
||||
)
|
||||
|
||||
# Initialize EWMA state
|
||||
policy._q_level_nbytes[map_op] = level
|
||||
policy._q_level_dev[map_op] = dev
|
||||
|
||||
result = policy.can_add_input(map_op)
|
||||
assert (
|
||||
result == expected_result
|
||||
), f"Expected {expected_result} for {description}"
|
||||
|
||||
|
||||
def test_emits_deprecation_warning_when_dynamic_backpressure_enabled(
|
||||
restore_data_context,
|
||||
):
|
||||
ctx = DataContext.get_current()
|
||||
ctx.enable_dynamic_output_queue_size_backpressure = True
|
||||
input_op = InputDataBuffer(ctx, input_data=[MagicMock()])
|
||||
topology = {input_op: MagicMock()}
|
||||
|
||||
with pytest.warns(RayDeprecationWarning, match="deprecated"):
|
||||
ConcurrencyCapBackpressurePolicy(ctx, topology, MagicMock())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,391 @@
|
||||
import time
|
||||
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data._internal.arrow_block import ArrowBlockAccessor
|
||||
from ray.data._internal.arrow_ops.transform_pyarrow import try_combine_chunked_columns
|
||||
from ray.data._internal.batcher import (
|
||||
SHUFFLE_BUFFER_COMPACTION_THRESHOLD,
|
||||
Batcher,
|
||||
ShufflingBatcher,
|
||||
)
|
||||
from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
|
||||
from ray.data.block import BlockAccessor
|
||||
|
||||
|
||||
def gen_block(num_rows):
|
||||
return pa.table({"foo": [1] * num_rows})
|
||||
|
||||
|
||||
def test_shuffling_batcher():
|
||||
batch_size = 5
|
||||
buffer_size = 20
|
||||
|
||||
with pytest.raises(
|
||||
ValueError, match="Must specify a batch_size if using a local shuffle."
|
||||
):
|
||||
ShufflingBatcher(batch_size=None, shuffle_buffer_min_size=buffer_size)
|
||||
|
||||
# Should not raise error.
|
||||
ShufflingBatcher(batch_size=batch_size, shuffle_buffer_min_size=batch_size - 1)
|
||||
|
||||
batcher = ShufflingBatcher(
|
||||
batch_size=batch_size,
|
||||
shuffle_buffer_min_size=buffer_size,
|
||||
)
|
||||
|
||||
total_added = 0
|
||||
total_yielded = 0
|
||||
|
||||
def add_and_check(num_rows, expect_has_batch):
|
||||
"""Add a block and verify has_batch() matches expectation."""
|
||||
nonlocal total_added
|
||||
batcher.add(gen_block(num_rows))
|
||||
total_added += num_rows
|
||||
assert batcher.has_batch() == expect_has_batch, (
|
||||
f"after adding {num_rows}: has_batch()={batcher.has_batch()}, "
|
||||
f"expected {expect_has_batch} "
|
||||
f"(compacted={batcher._num_compacted_rows()}, "
|
||||
f"uncompacted={batcher._num_uncompacted_rows()}, "
|
||||
f"total={batcher._num_rows()})"
|
||||
)
|
||||
|
||||
def next_and_check(
|
||||
expect_full_batch=True,
|
||||
expect_has_batch_after=True,
|
||||
):
|
||||
"""Consume one batch and verify size and post-state."""
|
||||
nonlocal total_yielded
|
||||
batch = batcher.next_batch()
|
||||
total_yielded += len(batch)
|
||||
if expect_full_batch:
|
||||
assert (
|
||||
len(batch) == batch_size
|
||||
), f"expected full batch of {batch_size}, got {len(batch)}"
|
||||
else:
|
||||
assert len(batch) <= batch_size
|
||||
assert batcher.has_batch() == expect_has_batch_after, (
|
||||
f"after next_batch: has_batch()={batcher.has_batch()}, "
|
||||
f"expected {expect_has_batch_after} "
|
||||
f"(compacted={batcher._num_compacted_rows()}, "
|
||||
f"uncompacted={batcher._num_uncompacted_rows()}, "
|
||||
f"total={batcher._num_rows()})"
|
||||
)
|
||||
|
||||
# Before any data is added, there should be no batches.
|
||||
assert not batcher.has_batch()
|
||||
assert not batcher.has_any()
|
||||
|
||||
# Add blocks incrementally. All rows go into the pending buffer,
|
||||
# and has_batch will return False until enough rows accumulate.
|
||||
add_and_check(3, expect_has_batch=False) # total=3
|
||||
add_and_check(7, expect_has_batch=False) # total=10
|
||||
add_and_check(10, expect_has_batch=False) # total=20
|
||||
|
||||
# After adding 15 more (total=35), total - batch_size = 30 >= min_rows_to_trigger.
|
||||
add_and_check(15, expect_has_batch=True) # total=35
|
||||
|
||||
# All 35 rows are still uncompacted since no next_batch() has been called.
|
||||
assert batcher._shuffle_buffer is None
|
||||
assert batcher._builder.num_rows() == 35
|
||||
|
||||
# Consume one batch — this triggers the first compaction.
|
||||
next_and_check(expect_full_batch=True, expect_has_batch_after=True)
|
||||
assert batcher._shuffle_buffer is not None # compaction happened
|
||||
assert batcher._builder.num_rows() == 0 # all rows moved to compacted buffer
|
||||
|
||||
# Add more data while consuming.
|
||||
add_and_check(20, expect_has_batch=True) # total grows
|
||||
|
||||
# Consume batches. Each must be full since we're still streaming.
|
||||
while batcher.has_batch():
|
||||
batch = batcher.next_batch()
|
||||
assert len(batch) == batch_size
|
||||
total_yielded += batch_size
|
||||
|
||||
# Streaming exhausted: remaining rows <= batch_size (not enough to trigger
|
||||
# has_batch without more data or done_adding).
|
||||
assert batcher._num_rows() <= batch_size
|
||||
|
||||
# Add a partial amount and signal done.
|
||||
batcher.add(gen_block(8))
|
||||
total_added += 8
|
||||
batcher.done_adding()
|
||||
|
||||
# Drain remaining full batches via next_and_check.
|
||||
while batcher.has_batch():
|
||||
remaining_after = batcher._num_rows() - batch_size
|
||||
next_and_check(
|
||||
expect_full_batch=True,
|
||||
expect_has_batch_after=remaining_after >= batch_size,
|
||||
)
|
||||
|
||||
# Final partial batch.
|
||||
if batcher.has_any():
|
||||
next_and_check(expect_full_batch=False, expect_has_batch_after=False)
|
||||
|
||||
# All rows must be accounted for.
|
||||
assert total_yielded == total_added
|
||||
assert not batcher.has_any()
|
||||
|
||||
|
||||
def test_batching_pyarrow_table_with_many_chunks():
|
||||
"""Make sure batching a pyarrow table with many chunks is fast.
|
||||
|
||||
See https://github.com/ray-project/ray/issues/31108 for more details.
|
||||
"""
|
||||
num_chunks = 5000
|
||||
batch_size = 1024
|
||||
|
||||
batches = []
|
||||
for _ in range(num_chunks):
|
||||
batch = {}
|
||||
for i in range(10):
|
||||
batch[str(i)] = list(range(batch_size))
|
||||
batches.append(pa.Table.from_pydict(batch))
|
||||
|
||||
block = pa.concat_tables(batches, promote=True)
|
||||
|
||||
start = time.perf_counter()
|
||||
batcher = Batcher(batch_size, ensure_copy=False)
|
||||
batcher.add(block)
|
||||
batcher.done_adding()
|
||||
while batcher.has_any():
|
||||
batcher.next_batch()
|
||||
duration = time.perf_counter() - start
|
||||
assert duration < 10
|
||||
|
||||
start = time.perf_counter()
|
||||
shuffling_batcher = ShufflingBatcher(
|
||||
batch_size=batch_size, shuffle_buffer_min_size=batch_size
|
||||
)
|
||||
shuffling_batcher.add(block)
|
||||
shuffling_batcher.done_adding()
|
||||
while shuffling_batcher.has_any():
|
||||
shuffling_batcher.next_batch()
|
||||
duration = time.perf_counter() - start
|
||||
assert duration < 30
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_size,local_shuffle_buffer_size",
|
||||
[(1, 1), (10, 1), (1, 10), (10, 1000), (1000, 10)],
|
||||
)
|
||||
def test_shuffling_batcher_grid(batch_size, local_shuffle_buffer_size):
|
||||
ds = ray.data.range_tensor(10000, shape=(130,))
|
||||
start = time.time()
|
||||
count = 0
|
||||
for batch in ds.iter_batches(
|
||||
batch_size=batch_size, local_shuffle_buffer_size=local_shuffle_buffer_size
|
||||
):
|
||||
count += len(batch["data"])
|
||||
print((ds.size_bytes() / 1e9) / (time.time() - start), "GB/s")
|
||||
assert count == 10000
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_size,local_shuffle_buffer_size",
|
||||
[(1, 1), (10, 1), (1, 10), (10, 100), (100, 10)],
|
||||
)
|
||||
def test_local_shuffle_determinism(batch_size, local_shuffle_buffer_size):
|
||||
# Preserve order so that the blocks are in the same order prior to shuffling.
|
||||
ctx = ray.data.DataContext.get_current()
|
||||
ctx.execution_options.preserve_order = True
|
||||
TEST_ITERATIONS = 10
|
||||
|
||||
ds = ray.data.range(1000)
|
||||
batch_map = {}
|
||||
for i in range(TEST_ITERATIONS):
|
||||
for batch in ds.iter_batches(
|
||||
batch_size=batch_size,
|
||||
local_shuffle_buffer_size=local_shuffle_buffer_size,
|
||||
local_shuffle_seed=0,
|
||||
):
|
||||
if i == 0:
|
||||
batch_map[batch["id"][0]] = batch
|
||||
else:
|
||||
# Check that batch is the same as the first dataset's batch
|
||||
assert all(batch_map[batch["id"][0]]["id"] == batch["id"])
|
||||
|
||||
|
||||
def test_local_shuffle_buffer_warns_if_too_large(shutdown_only):
|
||||
ray.shutdown()
|
||||
ray.init(object_store_memory=128 * 1024 * 1024)
|
||||
|
||||
# Each row is 16 MiB * 8 = 128 MiB
|
||||
ds = ray.data.range_tensor(2, shape=(16, 1024, 1024))
|
||||
|
||||
# Test that Ray Data emits a warning if the local shuffle buffer size would cause
|
||||
# spilling.
|
||||
with pytest.warns(UserWarning, match="shuffle buffer"):
|
||||
# Each row is 128 MiB and the shuffle buffer size is 2 rows, so expect at least
|
||||
# 256 MiB of memory usage > 128 MiB total on node.
|
||||
batches = ds.iter_batches(batch_size=1, local_shuffle_buffer_size=2)
|
||||
next(iter(batches))
|
||||
|
||||
|
||||
def _collect_rows_full_method(blocks, batch_size, buffer_size, seed):
|
||||
"""Reference implementation using the old full-shuffle method.
|
||||
|
||||
Materializes a fully shuffled copy of the buffer on each compaction,
|
||||
then yields contiguous slices. Used to validate the incremental index method.
|
||||
"""
|
||||
shuffle_buffer_min_size = max(buffer_size, batch_size)
|
||||
|
||||
min_rows_to_yield_batch = max(
|
||||
1, int(shuffle_buffer_min_size * SHUFFLE_BUFFER_COMPACTION_THRESHOLD)
|
||||
)
|
||||
|
||||
builder = DelegatingBlockBuilder()
|
||||
shuffle_buffer = None
|
||||
batch_head = 0
|
||||
shuffle_seed = seed
|
||||
|
||||
for block in blocks:
|
||||
if BlockAccessor.for_block(block).num_rows() > 0:
|
||||
builder.add_block(block)
|
||||
|
||||
done_adding = True
|
||||
rows = []
|
||||
|
||||
while True:
|
||||
compacted = 0
|
||||
if shuffle_buffer is not None:
|
||||
compacted = max(
|
||||
0, BlockAccessor.for_block(shuffle_buffer).num_rows() - batch_head
|
||||
)
|
||||
uncompacted = builder.num_rows()
|
||||
num_rows = compacted + uncompacted
|
||||
|
||||
has_batch = num_rows >= batch_size
|
||||
has_any = num_rows > 0
|
||||
|
||||
if not (has_batch or (done_adding and has_any)):
|
||||
break
|
||||
|
||||
# Compaction: merge uncompacted rows into shuffle buffer.
|
||||
if uncompacted > 0 and (done_adding or compacted <= min_rows_to_yield_batch):
|
||||
if shuffle_buffer is not None:
|
||||
if batch_head > 0:
|
||||
block_acc = BlockAccessor.for_block(shuffle_buffer)
|
||||
shuffle_buffer = block_acc.slice(batch_head, block_acc.num_rows())
|
||||
builder.add_block(shuffle_buffer)
|
||||
shuffle_buffer = builder.build()
|
||||
shuffle_buffer = BlockAccessor.for_block(shuffle_buffer).random_shuffle(
|
||||
shuffle_seed
|
||||
)
|
||||
if shuffle_seed is not None:
|
||||
shuffle_seed += 1
|
||||
if isinstance(BlockAccessor.for_block(shuffle_buffer), ArrowBlockAccessor):
|
||||
shuffle_buffer = try_combine_chunked_columns(shuffle_buffer)
|
||||
builder = DelegatingBlockBuilder()
|
||||
batch_head = 0
|
||||
|
||||
buf_size = BlockAccessor.for_block(shuffle_buffer).num_rows()
|
||||
bs = min(batch_size, buf_size - batch_head)
|
||||
batch = BlockAccessor.for_block(shuffle_buffer).slice(
|
||||
batch_head, batch_head + bs
|
||||
)
|
||||
batch_head += bs
|
||||
rows.extend(batch.column("val").to_pylist())
|
||||
|
||||
return rows
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_size,buffer_size,num_blocks,block_size",
|
||||
[
|
||||
(5, 20, 10, 10),
|
||||
(1, 10, 5, 20),
|
||||
(10, 10, 3, 50),
|
||||
(7, 30, 8, 15),
|
||||
(100, 100, 2, 200),
|
||||
],
|
||||
)
|
||||
def test_incremental_index_matches_full_method(
|
||||
batch_size, buffer_size, num_blocks, block_size
|
||||
):
|
||||
"""Verify that the incremental index method yields the same multiset of
|
||||
rows as the old full-shuffle reference implementation."""
|
||||
|
||||
seed = 42
|
||||
blocks = [
|
||||
pa.table({"val": list(range(i * block_size, (i + 1) * block_size))})
|
||||
for i in range(num_blocks)
|
||||
]
|
||||
|
||||
# Incremental index method (current implementation).
|
||||
batcher = ShufflingBatcher(
|
||||
batch_size=batch_size,
|
||||
shuffle_buffer_min_size=buffer_size,
|
||||
shuffle_seed=seed,
|
||||
)
|
||||
for block in blocks:
|
||||
batcher.add(block)
|
||||
batcher.done_adding()
|
||||
rows_index = []
|
||||
while batcher.has_batch() or batcher.has_any():
|
||||
batch = batcher.next_batch()
|
||||
rows_index.extend(batch.column("val").to_pylist())
|
||||
|
||||
# Full-shuffle reference.
|
||||
rows_full = _collect_rows_full_method(blocks, batch_size, buffer_size, seed)
|
||||
|
||||
total_rows = num_blocks * block_size
|
||||
assert len(rows_index) == total_rows
|
||||
assert len(rows_full) == total_rows
|
||||
assert sorted(rows_index) == sorted(rows_full) == list(range(total_rows))
|
||||
|
||||
|
||||
def test_no_partial_batch_mid_stream():
|
||||
"""has_batch() must not return True when total rows < batch_size.
|
||||
|
||||
With SHUFFLE_BUFFER_COMPACTION_THRESHOLD < 1.0, _min_rows_to_yield_batch
|
||||
can be less than batch_size. If we drain the compacted buffer below
|
||||
batch_size while no uncompacted rows are available, has_batch() must
|
||||
return False — otherwise next_batch() would return a partial batch
|
||||
mid-stream.
|
||||
"""
|
||||
batch_size = 10
|
||||
buffer_size = 10 # common case: equal to batch_size
|
||||
|
||||
batcher = ShufflingBatcher(
|
||||
batch_size=batch_size,
|
||||
shuffle_buffer_min_size=buffer_size,
|
||||
shuffle_seed=0,
|
||||
)
|
||||
|
||||
# Add enough rows to trigger compaction and yield some batches.
|
||||
batcher.add(gen_block(35))
|
||||
|
||||
# Consume batches until the compacted buffer is partially drained.
|
||||
batches = []
|
||||
while batcher.has_batch():
|
||||
batch = batcher.next_batch()
|
||||
batches.append(batch)
|
||||
# Every batch returned mid-stream must be full.
|
||||
assert (
|
||||
len(batch) == batch_size
|
||||
), f"got partial batch of {len(batch)} rows mid-stream"
|
||||
|
||||
# At this point has_batch() is False. There may be leftover rows
|
||||
# (< batch_size) but they should not be yielded until done_adding.
|
||||
leftover = batcher._num_rows()
|
||||
assert leftover < batch_size
|
||||
|
||||
# After done_adding, the remaining rows should drain as a partial batch.
|
||||
batcher.done_adding()
|
||||
assert batcher.has_any()
|
||||
final_batch = batcher.next_batch()
|
||||
assert len(final_batch) == leftover
|
||||
|
||||
total = sum(len(b) for b in batches) + len(final_batch)
|
||||
assert total == 35
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,146 @@
|
||||
import gc
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray._common.test_utils import wait_for_condition
|
||||
from ray.data._internal.execution.block_ref_counter import BlockRefCounter
|
||||
from ray.data._internal.util import MiB
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
# Grace period for asserting a callback has NOT fired. Must be shorter than
|
||||
# the task sleep in test_task_ref_keeps_counter_alive (1.0s); 0.3s leaves
|
||||
# wide margin even on slow CI while still surfacing early-fire bugs.
|
||||
_EARLY_FIRE_GRACE_S = 0.3
|
||||
|
||||
|
||||
@ray.remote
|
||||
def _hold_ref_for(block_ref, sleep_s: float) -> bool:
|
||||
"""Hold *block_ref* as a task argument for *sleep_s* seconds, then return.
|
||||
|
||||
Ray keeps an object alive for the duration of any task that received it as
|
||||
an argument, so this lets tests assert the callback has not fired while the
|
||||
task is still running.
|
||||
"""
|
||||
time.sleep(sleep_s)
|
||||
return True
|
||||
|
||||
|
||||
@pytest.fixture(params=["inlined", "regular"])
|
||||
def make_block(request):
|
||||
"""Factory for a block (ObjectRef, size_bytes), parametrized over the two
|
||||
storage paths.
|
||||
|
||||
Ray Core inlines tiny objects in the in-process store and puts larger ones
|
||||
in the shared-memory object store; the out-of-scope callback must work for
|
||||
both. Returning a factory (rather than the ref itself) avoids pytest holding
|
||||
an extra reference that would keep the object alive past the test's own ``del``.
|
||||
"""
|
||||
|
||||
def _make() -> tuple["ray.ObjectRef", int]:
|
||||
if request.param == "inlined":
|
||||
data = np.zeros(1, dtype=np.uint8)
|
||||
else:
|
||||
data = np.zeros(1 * MiB, dtype=np.uint8)
|
||||
return ray.put(data), len(data)
|
||||
|
||||
return _make
|
||||
|
||||
|
||||
def _wait_for_counter(*, counter, producer_id, expected, timeout_s: float = 10.0):
|
||||
"""Wait until *counter* reports *expected* bytes for *producer_id*.
|
||||
|
||||
``gc.collect()`` runs on each poll so any pending Python-level ObjectRef
|
||||
destructors get a chance to run; the polling/timeout loop is delegated to
|
||||
``wait_for_condition`` (raises on timeout).
|
||||
"""
|
||||
|
||||
def _reached():
|
||||
gc.collect()
|
||||
return counter.get_object_store_memory_usage(producer_id) == expected
|
||||
|
||||
wait_for_condition(_reached, timeout=timeout_s)
|
||||
|
||||
|
||||
class TestBlockRefCounterLifecycle:
|
||||
def test_callback_fires_after_last_python_ref_deleted(
|
||||
self, ray_start_regular_shared, make_block
|
||||
):
|
||||
"""Counter reaches 0 once the only Python ObjectRef is GC'd."""
|
||||
counter = BlockRefCounter()
|
||||
ref, size_bytes = make_block()
|
||||
|
||||
counter.on_block_produced(ref, size_bytes, "op_basic")
|
||||
assert counter.get_object_store_memory_usage("op_basic") == size_bytes
|
||||
|
||||
del ref # last Python ref gone
|
||||
_wait_for_counter(counter=counter, producer_id="op_basic", expected=0)
|
||||
|
||||
def test_second_python_ref_keeps_counter_alive(
|
||||
self, ray_start_regular_shared, make_block
|
||||
):
|
||||
"""Counter stays non-zero while a second Python ObjectRef is alive.
|
||||
|
||||
Dropping one of two refs that point at the same ObjectID must NOT fire
|
||||
the callback. Only the final ref drop may do so.
|
||||
"""
|
||||
counter = BlockRefCounter()
|
||||
ref1, size_bytes = make_block()
|
||||
ref2 = ref1 # second Python ref to the same ObjectID
|
||||
|
||||
counter.on_block_produced(ref1, size_bytes, "op_two_refs")
|
||||
assert counter.get_object_store_memory_usage("op_two_refs") == size_bytes
|
||||
|
||||
del ref1
|
||||
gc.collect()
|
||||
time.sleep(_EARLY_FIRE_GRACE_S) # counter must still be non-zero
|
||||
|
||||
assert (
|
||||
counter.get_object_store_memory_usage("op_two_refs") == size_bytes
|
||||
), "Callback fired too early — counter decremented while ref2 was still alive"
|
||||
|
||||
del ref2 # last ref gone; callback must now fire
|
||||
_wait_for_counter(counter=counter, producer_id="op_two_refs", expected=0)
|
||||
|
||||
def test_task_ref_keeps_counter_alive_until_task_completes(
|
||||
self, ray_start_regular_shared
|
||||
):
|
||||
"""Counter stays non-zero while a running Ray task holds the block.
|
||||
|
||||
Ray keeps any object alive for the duration of a task that received it
|
||||
as an argument. The callback should not fire until both conditions hold:
|
||||
(a) the task has completed, and (b) all Python refs are dropped.
|
||||
|
||||
Uses a plasma (by-reference) object specifically: tiny objects are
|
||||
inlined into the task by value, so they would not get a task-argument
|
||||
reference and this lifetime-extension behavior would not apply.
|
||||
"""
|
||||
counter = BlockRefCounter()
|
||||
ref = ray.put(
|
||||
np.zeros(1 * MiB, dtype=np.uint8) # pyrefly: ignore[bad-argument-type]
|
||||
)
|
||||
|
||||
counter.on_block_produced(ref, 1 * MiB, "op_task")
|
||||
assert counter.get_object_store_memory_usage("op_task") == 1 * MiB
|
||||
|
||||
# Submit a task that sleeps while holding the block, then drop the Python
|
||||
# ref so only the task's argument reference remains.
|
||||
task_future = _hold_ref_for.remote(ref, 1.0)
|
||||
del ref
|
||||
gc.collect()
|
||||
time.sleep(_EARLY_FIRE_GRACE_S) # task still running; callback must NOT fire
|
||||
|
||||
assert (
|
||||
counter.get_object_store_memory_usage("op_task") == 1 * MiB
|
||||
), "Callback fired too early: counter decremented while task was still running"
|
||||
|
||||
ray.get(task_future) # task completes; now both refs are gone
|
||||
_wait_for_counter(counter=counter, producer_id="op_task", expected=0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,259 @@
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data.context import DataContext
|
||||
from ray.data.dataset import Dataset
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.data.tests.conftest import (
|
||||
assert_blocks_expected_in_plasma,
|
||||
get_initial_core_execution_metrics_snapshot,
|
||||
)
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
|
||||
def _assert_num_blocks(ds, expected, tolerance=0.5):
|
||||
actual = ds.num_blocks()
|
||||
assert (
|
||||
expected * (1 - tolerance) <= actual <= expected * (1 + tolerance)
|
||||
), f"Expected ~{expected} blocks (±{tolerance*100}%), got {actual}"
|
||||
|
||||
|
||||
def test_map(shutdown_only, restore_data_context):
|
||||
ray.init(
|
||||
_system_config={
|
||||
"max_direct_call_object_size": 10_000,
|
||||
},
|
||||
num_cpus=2,
|
||||
object_store_memory=int(100e6),
|
||||
)
|
||||
|
||||
ctx = DataContext.get_current()
|
||||
ctx.target_min_block_size = 10_000 * 8
|
||||
ctx.target_max_block_size = 10_000 * 8
|
||||
num_blocks_expected = 10
|
||||
|
||||
# Test read.
|
||||
ds = ray.data.range(100_000, override_num_blocks=1).materialize()
|
||||
_assert_num_blocks(ds, num_blocks_expected)
|
||||
|
||||
# Test read -> map.
|
||||
# NOTE(swang): For some reason BlockBuilder's estimated memory usage when a
|
||||
# map fn is used is 2x the actual memory usage.
|
||||
ds = (
|
||||
ray.data.range(100_000, override_num_blocks=1)
|
||||
.map(lambda row: row)
|
||||
.materialize()
|
||||
)
|
||||
_assert_num_blocks(ds, num_blocks_expected * 2)
|
||||
|
||||
# Test adjusted block size.
|
||||
ctx.target_max_block_size *= 2
|
||||
num_blocks_expected //= 2
|
||||
|
||||
# Test read.
|
||||
ds = ray.data.range(100_000, override_num_blocks=1).materialize()
|
||||
_assert_num_blocks(ds, num_blocks_expected)
|
||||
|
||||
# Test read -> map.
|
||||
ds = (
|
||||
ray.data.range(100_000, override_num_blocks=1)
|
||||
.map(lambda row: row)
|
||||
.materialize()
|
||||
)
|
||||
_assert_num_blocks(ds, num_blocks_expected * 2)
|
||||
|
||||
# Setting the shuffle block size prints a warning and actually resets
|
||||
# target_max_block_size
|
||||
ctx.target_shuffle_max_block_size = ctx.target_max_block_size / 2
|
||||
num_blocks_expected *= 2
|
||||
|
||||
# Test read.
|
||||
ds = ray.data.range(100_000, override_num_blocks=1).materialize()
|
||||
_assert_num_blocks(ds, num_blocks_expected)
|
||||
|
||||
# Test read -> map.
|
||||
ds = (
|
||||
ray.data.range(100_000, override_num_blocks=1)
|
||||
.map(lambda row: row)
|
||||
.materialize()
|
||||
)
|
||||
_assert_num_blocks(ds, num_blocks_expected * 2)
|
||||
|
||||
|
||||
# TODO: Test that map stage output blocks are the correct size for groupby and
|
||||
# repartition. Currently we only have access to the reduce stage output block
|
||||
# size.
|
||||
SHUFFLE_ALL_TO_ALL_OPS = [
|
||||
(Dataset.random_shuffle, {}, True),
|
||||
(Dataset.sort, {"key": "id"}, False),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"shuffle_op",
|
||||
SHUFFLE_ALL_TO_ALL_OPS,
|
||||
)
|
||||
def test_shuffle(shutdown_only, restore_data_context, shuffle_op):
|
||||
ray.init(
|
||||
_system_config={
|
||||
"max_direct_call_object_size": 250,
|
||||
},
|
||||
num_cpus=2,
|
||||
object_store_memory=int(100e6),
|
||||
)
|
||||
|
||||
# Test AllToAll and Map -> AllToAll Datasets. Check that Map inherits
|
||||
# AllToAll's target block size.
|
||||
ctx = DataContext.get_current()
|
||||
ctx.read_op_min_num_blocks = 1
|
||||
ctx.target_min_block_size = 1
|
||||
|
||||
N = 100_000
|
||||
mem_size = 800_000
|
||||
|
||||
shuffle_fn, kwargs, fusion_supported = shuffle_op
|
||||
|
||||
ctx.target_max_block_size = 10_000 * 8
|
||||
num_blocks_expected = mem_size // ctx.target_max_block_size
|
||||
last_snapshot = get_initial_core_execution_metrics_snapshot()
|
||||
|
||||
ds = shuffle_fn(ray.data.range(N), **kwargs).materialize()
|
||||
assert (
|
||||
num_blocks_expected
|
||||
<= ds._logical_plan.initial_num_blocks()
|
||||
<= num_blocks_expected * 1.5
|
||||
)
|
||||
|
||||
def _estimate_intermediate_blocks(fusion_supported: bool, num_blocks_expected: int):
|
||||
return num_blocks_expected**2 + num_blocks_expected * (
|
||||
2 if fusion_supported else 4
|
||||
)
|
||||
|
||||
# map * reduce intermediate blocks + 1 metadata ref per map/reduce task.
|
||||
# If fusion is not supported, the un-fused map stage produces 1 data and 1
|
||||
# metadata per task.
|
||||
num_intermediate_blocks = _estimate_intermediate_blocks(
|
||||
fusion_supported, num_blocks_expected
|
||||
)
|
||||
|
||||
print(f">>> Asserting {num_intermediate_blocks} blocks are in plasma")
|
||||
|
||||
last_snapshot = assert_blocks_expected_in_plasma(
|
||||
last_snapshot,
|
||||
# Dataset.sort produces some empty intermediate blocks because the
|
||||
# input range is already partially sorted.
|
||||
num_intermediate_blocks,
|
||||
)
|
||||
|
||||
ds = shuffle_fn(ray.data.range(N).map(lambda x: x), **kwargs).materialize()
|
||||
if not fusion_supported:
|
||||
# TODO(swang): For some reason BlockBuilder's estimated
|
||||
# memory usage for range(1000)->map is 2x the actual memory usage.
|
||||
# Remove once https://github.com/ray-project/ray/issues/40246 is fixed.
|
||||
num_blocks_expected = int(num_blocks_expected * 2.2)
|
||||
assert (
|
||||
num_blocks_expected
|
||||
<= ds._logical_plan.initial_num_blocks()
|
||||
<= num_blocks_expected * 1.5
|
||||
)
|
||||
num_intermediate_blocks = _estimate_intermediate_blocks(
|
||||
fusion_supported, num_blocks_expected
|
||||
)
|
||||
last_snapshot = assert_blocks_expected_in_plasma(
|
||||
last_snapshot,
|
||||
# Dataset.sort produces some empty intermediate blocks because the
|
||||
# input range is already partially sorted.
|
||||
num_intermediate_blocks,
|
||||
)
|
||||
|
||||
ctx.target_max_block_size //= 2
|
||||
num_blocks_expected = mem_size // ctx.target_max_block_size
|
||||
block_size_expected = ctx.target_max_block_size
|
||||
|
||||
ds = shuffle_fn(ray.data.range(N), **kwargs).materialize()
|
||||
assert (
|
||||
num_blocks_expected
|
||||
<= ds._logical_plan.initial_num_blocks()
|
||||
<= num_blocks_expected * 1.5
|
||||
)
|
||||
num_intermediate_blocks = _estimate_intermediate_blocks(
|
||||
fusion_supported, num_blocks_expected
|
||||
)
|
||||
last_snapshot = assert_blocks_expected_in_plasma(
|
||||
last_snapshot,
|
||||
num_intermediate_blocks,
|
||||
)
|
||||
|
||||
ds = shuffle_fn(ray.data.range(N).map(lambda x: x), **kwargs).materialize()
|
||||
if not fusion_supported:
|
||||
num_blocks_expected = int(num_blocks_expected * 2.2)
|
||||
block_size_expected //= 2.2
|
||||
assert (
|
||||
num_blocks_expected
|
||||
<= ds._logical_plan.initial_num_blocks()
|
||||
<= num_blocks_expected * 1.5
|
||||
)
|
||||
num_intermediate_blocks = _estimate_intermediate_blocks(
|
||||
fusion_supported, num_blocks_expected
|
||||
)
|
||||
last_snapshot = assert_blocks_expected_in_plasma(
|
||||
last_snapshot,
|
||||
num_intermediate_blocks,
|
||||
)
|
||||
|
||||
# Setting target max block size does not affect map ops when there is a
|
||||
# shuffle downstream.
|
||||
ctx.target_max_block_size = ctx.target_max_block_size * 2
|
||||
num_blocks_expected //= 2
|
||||
|
||||
ds = shuffle_fn(ray.data.range(N).map(lambda x: x), **kwargs).materialize()
|
||||
assert (
|
||||
num_blocks_expected
|
||||
<= ds._logical_plan.initial_num_blocks()
|
||||
<= num_blocks_expected * 1.5
|
||||
)
|
||||
|
||||
num_intermediate_blocks = _estimate_intermediate_blocks(
|
||||
fusion_supported, num_blocks_expected
|
||||
)
|
||||
|
||||
assert_blocks_expected_in_plasma(
|
||||
last_snapshot,
|
||||
num_intermediate_blocks,
|
||||
)
|
||||
|
||||
|
||||
def test_target_max_block_size_infinite_or_default_disables_splitting_globally(
|
||||
shutdown_only, restore_data_context
|
||||
):
|
||||
"""Test that setting target_max_block_size to None disables block splitting globally."""
|
||||
ray.init(num_cpus=2)
|
||||
|
||||
# Create a large dataset that would normally trigger block splitting
|
||||
N = 1_000_000 # ~8MB worth of data
|
||||
|
||||
# First, test with normal target_max_block_size (should split into multiple blocks)
|
||||
ctx = DataContext.get_current()
|
||||
ctx.target_max_block_size = 1_000_000 # ~1MB
|
||||
|
||||
ds_with_limit = ray.data.range(N, override_num_blocks=1).materialize()
|
||||
blocks_with_limit = ds_with_limit._logical_plan.initial_num_blocks()
|
||||
|
||||
# Now test with target_max_block_size = None (should not split)
|
||||
ctx.target_max_block_size = None # Disable block size limit
|
||||
|
||||
ds_unlimited = (
|
||||
ray.data.range(N, override_num_blocks=1).map(lambda x: x).materialize()
|
||||
)
|
||||
blocks_unlimited = ds_unlimited._logical_plan.initial_num_blocks()
|
||||
|
||||
# Verify that unlimited creates fewer blocks (no splitting)
|
||||
assert blocks_unlimited <= blocks_with_limit
|
||||
# With target_max_block_size=None, it should maintain the original block structure
|
||||
assert blocks_unlimited == 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-sv", __file__]))
|
||||
@@ -0,0 +1,149 @@
|
||||
from typing import Any
|
||||
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data._internal.execution.bundle_queue import create_bundle_queue
|
||||
from ray.data._internal.execution.interfaces import BlockEntry, RefBundle
|
||||
from ray.data.block import BlockAccessor
|
||||
|
||||
|
||||
def _create_bundle(data: Any) -> RefBundle:
|
||||
"""Create a RefBundle with a single row with the given data."""
|
||||
block = pa.Table.from_pydict({"data": [data]})
|
||||
block_ref = ray.put(block)
|
||||
metadata = BlockAccessor.for_block(block).get_metadata()
|
||||
schema = BlockAccessor.for_block(block).schema()
|
||||
return RefBundle(
|
||||
[BlockEntry(block_ref, metadata)], owns_blocks=False, schema=schema
|
||||
)
|
||||
|
||||
|
||||
# CVGA-start
|
||||
def test_add_and_length():
|
||||
queue = create_bundle_queue()
|
||||
queue.add(_create_bundle("test1"))
|
||||
queue.add(_create_bundle("test2"))
|
||||
assert len(queue) == 2
|
||||
|
||||
|
||||
def test_get_next():
|
||||
queue = create_bundle_queue()
|
||||
bundle1 = _create_bundle("test1")
|
||||
queue.add(bundle1)
|
||||
bundle2 = _create_bundle("test11")
|
||||
queue.add(bundle2)
|
||||
|
||||
popped_bundle = queue.get_next()
|
||||
assert popped_bundle is bundle1
|
||||
assert len(queue) == 1
|
||||
assert queue.num_blocks() == 1
|
||||
assert queue.num_rows() == 1
|
||||
assert queue.estimate_size_bytes() == bundle2.size_bytes()
|
||||
|
||||
|
||||
def test_peek_next():
|
||||
queue = create_bundle_queue()
|
||||
bundle1 = _create_bundle("test1")
|
||||
queue.add(bundle1)
|
||||
bundle2 = _create_bundle("test11")
|
||||
queue.add(bundle2)
|
||||
|
||||
peeked_bundle = queue.peek_next()
|
||||
assert peeked_bundle is bundle1
|
||||
assert len(queue) == 2 # Length should remain unchanged
|
||||
assert queue.num_blocks() == 2
|
||||
assert queue.num_rows() == 2
|
||||
assert queue.estimate_size_bytes() == bundle1.size_bytes() + bundle2.size_bytes()
|
||||
|
||||
|
||||
def test_get_next_empty_queue():
|
||||
queue = create_bundle_queue()
|
||||
with pytest.raises(IndexError):
|
||||
queue.get_next()
|
||||
|
||||
|
||||
def test_get_next_does_not_leak_objects():
|
||||
queue = create_bundle_queue()
|
||||
bundle1 = _create_bundle("test11")
|
||||
queue.add(bundle1)
|
||||
queue.get_next()
|
||||
assert len(queue) == 0
|
||||
assert queue.estimate_size_bytes() == 0
|
||||
assert queue.num_rows() == 0
|
||||
assert queue.num_blocks() == 0
|
||||
|
||||
|
||||
def test_peek_next_empty_queue():
|
||||
queue = create_bundle_queue()
|
||||
assert queue.peek_next() is None
|
||||
assert len(queue) == 0
|
||||
assert queue.num_blocks() == 0
|
||||
assert queue.estimate_size_bytes() == 0
|
||||
assert queue.num_rows() == 0
|
||||
|
||||
|
||||
def test_remove():
|
||||
queue = create_bundle_queue()
|
||||
bundle1 = _create_bundle("test1")
|
||||
bundle2 = _create_bundle("test11")
|
||||
queue.add(bundle1)
|
||||
queue.add(bundle2)
|
||||
|
||||
queue.remove(bundle1)
|
||||
assert len(queue) == 1
|
||||
assert queue.num_blocks() == 1
|
||||
assert queue.peek_next() is bundle2
|
||||
assert queue.estimate_size_bytes() == bundle2.size_bytes()
|
||||
assert queue.num_rows() == bundle2.num_rows()
|
||||
|
||||
|
||||
def test_remove_does_not_leak_objects():
|
||||
queue = create_bundle_queue()
|
||||
bundle1 = _create_bundle("test1")
|
||||
queue.add(bundle1)
|
||||
queue.remove(bundle1)
|
||||
assert len(queue) == 0
|
||||
assert queue.num_blocks() == 0
|
||||
assert queue.estimate_size_bytes() == 0
|
||||
assert queue.num_rows() == 0
|
||||
|
||||
|
||||
def test_add_and_remove_duplicates():
|
||||
queue = create_bundle_queue()
|
||||
bundle1 = _create_bundle("test1")
|
||||
bundle2 = _create_bundle("test11")
|
||||
queue.add(bundle1)
|
||||
queue.add(bundle2)
|
||||
queue.add(bundle1)
|
||||
|
||||
assert len(queue) == 3
|
||||
assert queue.num_rows() == 3
|
||||
assert queue.num_blocks() == 3
|
||||
queue.remove(bundle1)
|
||||
assert len(queue) == 2
|
||||
|
||||
assert queue.estimate_size_bytes() == bundle1.size_bytes() + bundle2.size_bytes()
|
||||
assert queue.num_rows() == 2
|
||||
assert queue.num_blocks() == 2
|
||||
assert queue.peek_next() is bundle2
|
||||
|
||||
|
||||
def test_clear():
|
||||
queue = create_bundle_queue()
|
||||
queue.add(_create_bundle("test1"))
|
||||
queue.add(_create_bundle("test11"))
|
||||
queue.clear()
|
||||
assert len(queue) == 0
|
||||
assert queue.estimate_size_bytes() == 0
|
||||
assert queue.num_blocks() == 0
|
||||
assert queue.num_rows() == 0
|
||||
|
||||
|
||||
# CVGA-end
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,413 @@
|
||||
"""Tests for the Catalog connector API (ray.data.catalog)."""
|
||||
|
||||
import contextlib
|
||||
import os
|
||||
import pickle
|
||||
from unittest import mock
|
||||
|
||||
import pyarrow as pa
|
||||
import pyarrow.fs as pafs
|
||||
import pyarrow.parquet as pq
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data.catalog import (
|
||||
Catalog,
|
||||
DatabricksUnityCatalog,
|
||||
ReaderFormat,
|
||||
ResolvedSource,
|
||||
)
|
||||
|
||||
# conftest provides ray_start_regular_shared
|
||||
from ray.data.tests.conftest import * # noqa: F401,F403
|
||||
|
||||
pytest.importorskip("databricks.sdk")
|
||||
from databricks.sdk.service.catalog import ( # noqa: E402
|
||||
AwsCredentials,
|
||||
AzureUserDelegationSas,
|
||||
DataSourceFormat,
|
||||
GcpOauthToken,
|
||||
GenerateTemporaryTableCredentialResponse,
|
||||
TableInfo,
|
||||
)
|
||||
|
||||
AWS_RESP = GenerateTemporaryTableCredentialResponse(
|
||||
url="s3://bucket/path",
|
||||
aws_temp_credentials=AwsCredentials(
|
||||
access_key_id="AKIA",
|
||||
secret_access_key="secret",
|
||||
session_token="token",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _mock_uc_sdk(*, data_source_format="DELTA", storage_location=None, creds=None):
|
||||
"""Patch DatabricksUnityCatalog._workspace_client to return canned SDK responses.
|
||||
|
||||
Replaces the catalog's two SDK calls (``tables.get`` /
|
||||
``generate_temporary_table_credentials``) so no Databricks workspace is hit.
|
||||
Returns an already-started patcher; the caller is responsible for ``stop()``.
|
||||
"""
|
||||
creds = creds if creds is not None else AWS_RESP
|
||||
table_info = TableInfo(
|
||||
table_id="tid-123",
|
||||
data_source_format=(
|
||||
DataSourceFormat(data_source_format) if data_source_format else None
|
||||
),
|
||||
storage_location=storage_location,
|
||||
)
|
||||
client = mock.MagicMock()
|
||||
client.tables.get.return_value = table_info
|
||||
gen = client.temporary_table_credentials.generate_temporary_table_credentials
|
||||
gen.return_value = creds
|
||||
|
||||
patcher = mock.patch.object(
|
||||
DatabricksUnityCatalog, "_workspace_client", return_value=client
|
||||
)
|
||||
patcher.start()
|
||||
return patcher
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def uc_catalog():
|
||||
return DatabricksUnityCatalog(
|
||||
url="https://dbc-test.cloud.databricks.com",
|
||||
token="dapi-test",
|
||||
region="us-west-2",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def isolated_env(monkeypatch):
|
||||
# resolve() exports vended creds to os.environ and may call ray.init.
|
||||
# Isolate the env mutation and skip the real ray.init for unit tests.
|
||||
monkeypatch.setattr("ray.is_initialized", lambda: True)
|
||||
monkeypatch.setattr(os, "environ", dict(os.environ))
|
||||
return os.environ
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# DatabricksUnityCatalog.resolve
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reader", [ReaderFormat.DELTA, ReaderFormat.PARQUET])
|
||||
def test_resolve_storage_aws(uc_catalog, isolated_env, reader):
|
||||
patcher = _mock_uc_sdk()
|
||||
try:
|
||||
resolved = uc_catalog.resolve("main.sales.txns", reader=reader)
|
||||
finally:
|
||||
patcher.stop()
|
||||
|
||||
assert resolved.path == "s3://bucket/path"
|
||||
# AWS creds are exported to the environment for both readers.
|
||||
assert isolated_env["AWS_ACCESS_KEY_ID"] == "AKIA"
|
||||
assert isolated_env["AWS_SESSION_TOKEN"] == "token"
|
||||
# Delta additionally gets an explicit S3FileSystem for the data scan;
|
||||
# Parquet reads via the environment-configured filesystem.
|
||||
if reader is ReaderFormat.DELTA:
|
||||
assert isinstance(resolved.filesystem, pafs.S3FileSystem)
|
||||
else:
|
||||
assert resolved.filesystem is None
|
||||
assert resolved.data_format is ReaderFormat.DELTA
|
||||
|
||||
|
||||
def test_resolve_aws_requires_region(isolated_env):
|
||||
catalog = DatabricksUnityCatalog(
|
||||
url="https://h.databricks.com", token="t"
|
||||
) # no region
|
||||
patcher = _mock_uc_sdk()
|
||||
try:
|
||||
with pytest.raises(ValueError, match="region"):
|
||||
catalog.resolve("main.sales.txns", reader=ReaderFormat.DELTA)
|
||||
finally:
|
||||
patcher.stop()
|
||||
|
||||
|
||||
def test_resolve_initializes_ray_with_runtime_env(uc_catalog, monkeypatch):
|
||||
# When Ray isn't running, vended creds are propagated via runtime_env.
|
||||
monkeypatch.setattr(os, "environ", dict(os.environ))
|
||||
monkeypatch.setattr("ray.is_initialized", lambda: False)
|
||||
init_kwargs = {}
|
||||
monkeypatch.setattr("ray.init", lambda **kw: init_kwargs.update(kw))
|
||||
|
||||
patcher = _mock_uc_sdk()
|
||||
try:
|
||||
uc_catalog.resolve("main.sales.txns", reader=ReaderFormat.PARQUET)
|
||||
finally:
|
||||
patcher.stop()
|
||||
|
||||
env_vars = init_kwargs["runtime_env"]["env_vars"]
|
||||
assert env_vars["AWS_ACCESS_KEY_ID"] == "AKIA"
|
||||
assert env_vars["AWS_SESSION_TOKEN"] == "token"
|
||||
|
||||
|
||||
def test_resolve_retries_once_on_401(uc_catalog, isolated_env):
|
||||
# On a 401, the provider is invalidated and the SDK call is retried once
|
||||
# (mirrors the old request_with_401_retry behavior).
|
||||
from databricks.sdk.errors import Unauthenticated
|
||||
|
||||
client = mock.MagicMock()
|
||||
client.tables.get.side_effect = [
|
||||
Unauthenticated("expired"),
|
||||
TableInfo(
|
||||
table_id="tid-123",
|
||||
data_source_format=DataSourceFormat("DELTA"),
|
||||
storage_location="s3://bucket/path",
|
||||
),
|
||||
]
|
||||
gen = client.temporary_table_credentials.generate_temporary_table_credentials
|
||||
gen.return_value = AWS_RESP
|
||||
|
||||
with mock.patch.object(
|
||||
DatabricksUnityCatalog, "_workspace_client", return_value=client
|
||||
), mock.patch.object(uc_catalog._provider, "invalidate") as invalidate:
|
||||
resolved = uc_catalog.resolve("main.sales.txns", reader=ReaderFormat.PARQUET)
|
||||
|
||||
invalidate.assert_called_once()
|
||||
assert client.tables.get.call_count == 2
|
||||
assert resolved.path == "s3://bucket/path"
|
||||
|
||||
|
||||
def test_resolve_iceberg(uc_catalog):
|
||||
# Iceberg resolution does not hit the credential-vending REST endpoints.
|
||||
resolved = uc_catalog.resolve("main.sales.txns", reader=ReaderFormat.ICEBERG)
|
||||
|
||||
assert resolved.path is None
|
||||
assert resolved.filesystem is None
|
||||
assert resolved.data_format is ReaderFormat.ICEBERG
|
||||
ckw = resolved.catalog_kwargs
|
||||
assert ckw["type"] == "rest"
|
||||
assert ckw["uri"] == (
|
||||
"https://dbc-test.cloud.databricks.com/api/2.1/unity-catalog/iceberg-rest"
|
||||
)
|
||||
assert ckw["token"] == "dapi-test"
|
||||
assert ckw["header.X-Iceberg-Access-Delegation"] == "vended-credentials"
|
||||
|
||||
|
||||
def test_resolve_azure_sets_env(isolated_env):
|
||||
catalog = DatabricksUnityCatalog(url="https://h.databricks.com", token="t")
|
||||
azure_resp = GenerateTemporaryTableCredentialResponse(
|
||||
url="abfss://c@acct.dfs.core.windows.net/path",
|
||||
azure_user_delegation_sas=AzureUserDelegationSas(sas_token="sv=2021&sig=abc"),
|
||||
)
|
||||
patcher = _mock_uc_sdk(creds=azure_resp)
|
||||
try:
|
||||
resolved = catalog.resolve("main.sales.txns", reader=ReaderFormat.DELTA)
|
||||
finally:
|
||||
patcher.stop()
|
||||
|
||||
# Azure creds flow via the environment (read by both pyarrow and the
|
||||
# deltalake object_store log reader); no filesystem/storage_options.
|
||||
assert resolved.filesystem is None
|
||||
assert resolved.storage_options is None
|
||||
assert isolated_env["AZURE_STORAGE_SAS_TOKEN"] == "sv=2021&sig=abc"
|
||||
|
||||
|
||||
def test_resolve_azure_strips_leading_question_mark(isolated_env):
|
||||
# UC may return the SAS as a full query string ("?sv=..."); the leading "?"
|
||||
# must be stripped for AZURE_STORAGE_SAS_TOKEN.
|
||||
catalog = DatabricksUnityCatalog(url="https://h.databricks.com", token="t")
|
||||
azure_resp = GenerateTemporaryTableCredentialResponse(
|
||||
url="abfss://c@acct.dfs.core.windows.net/path",
|
||||
azure_user_delegation_sas=AzureUserDelegationSas(sas_token="?sv=2021&sig=abc"),
|
||||
)
|
||||
patcher = _mock_uc_sdk(creds=azure_resp)
|
||||
try:
|
||||
catalog.resolve("main.sales.txns", reader=ReaderFormat.DELTA)
|
||||
finally:
|
||||
patcher.stop()
|
||||
|
||||
assert isolated_env["AZURE_STORAGE_SAS_TOKEN"] == "sv=2021&sig=abc"
|
||||
|
||||
|
||||
def _gcp_resp():
|
||||
return GenerateTemporaryTableCredentialResponse(
|
||||
url="gs://bucket/path",
|
||||
gcp_oauth_token=GcpOauthToken(oauth_token="ya29.tok"),
|
||||
expiration_time=4102444800000, # far-future epoch ms
|
||||
)
|
||||
|
||||
|
||||
def test_resolve_gcp_parquet_builds_filesystem(uc_catalog, isolated_env):
|
||||
# GCP vends an OAuth token (not a service-account JSON); for Parquet it rides
|
||||
# on an explicit GcsFileSystem (the data scan), never an env var.
|
||||
patcher = _mock_uc_sdk(creds=_gcp_resp(), storage_location="gs://bucket/path")
|
||||
try:
|
||||
resolved = uc_catalog.resolve("main.sales.txns", reader=ReaderFormat.PARQUET)
|
||||
finally:
|
||||
patcher.stop()
|
||||
|
||||
assert isinstance(resolved.filesystem, pafs.GcsFileSystem)
|
||||
assert "GOOGLE_APPLICATION_CREDENTIALS" not in isolated_env
|
||||
|
||||
|
||||
def test_resolve_gcp_delta_raises(uc_catalog, isolated_env):
|
||||
# deltalake's object_store can't use a GCS OAuth token, so GCP + Delta is
|
||||
# rejected up front with an actionable error instead of failing deep in the
|
||||
# log read against the GCE metadata server.
|
||||
patcher = _mock_uc_sdk(creds=_gcp_resp(), storage_location="gs://bucket/path")
|
||||
try:
|
||||
with pytest.raises(RuntimeError, match="GCP-backed Delta"):
|
||||
uc_catalog.resolve("main.sales.txns", reader=ReaderFormat.DELTA)
|
||||
finally:
|
||||
patcher.stop()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Reader integration via a fake catalog (no network)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class _FakeCatalog(Catalog):
|
||||
"""Returns a pre-baked ResolvedSource; records the reader it was asked for."""
|
||||
|
||||
def __init__(self, resolved):
|
||||
self._resolved = resolved
|
||||
self.calls = []
|
||||
|
||||
def resolve(self, table, *, reader):
|
||||
self.calls.append((table, reader))
|
||||
return self._resolved
|
||||
|
||||
|
||||
def test_read_parquet_with_catalog(ray_start_regular_shared, tmp_path):
|
||||
path = str(tmp_path / "data.parquet")
|
||||
pq.write_table(pa.table({"id": [1, 2, 3]}), path)
|
||||
|
||||
catalog = _FakeCatalog(ResolvedSource(path=path))
|
||||
ds = ray.data.read_parquet("main.db.tbl", catalog=catalog)
|
||||
|
||||
assert sorted(r["id"] for r in ds.take_all()) == [1, 2, 3]
|
||||
assert catalog.calls == [("main.db.tbl", ReaderFormat.PARQUET)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("reader", ["parquet", "delta"])
|
||||
def test_catalog_filesystem_overrides_with_warning(reader):
|
||||
# The catalog-resolved filesystem overrides a user-supplied one, but warns.
|
||||
# The warning fires at the top of the reader body; suppress any downstream
|
||||
# failure from the (intentionally unreachable) s3 path.
|
||||
if reader == "delta":
|
||||
pytest.importorskip("deltalake")
|
||||
fs = pafs.S3FileSystem(
|
||||
access_key="AKIA", secret_key="secret", session_token="t", region="us-west-2"
|
||||
)
|
||||
catalog = _FakeCatalog(ResolvedSource(path="s3://b/p", filesystem=fs))
|
||||
read_fn = ray.data.read_parquet if reader == "parquet" else ray.data.read_delta
|
||||
|
||||
with mock.patch.object(ray.data.read_api.logger, "warning") as warn:
|
||||
with contextlib.suppress(Exception):
|
||||
read_fn("main.db.tbl", catalog=catalog, filesystem=pafs.LocalFileSystem())
|
||||
|
||||
assert any(
|
||||
"Overriding the provided `filesystem`" in str(c) for c in warn.call_args_list
|
||||
)
|
||||
|
||||
|
||||
def test_read_delta_with_catalog(ray_start_regular_shared, tmp_path):
|
||||
deltalake = pytest.importorskip("deltalake")
|
||||
path = str(tmp_path / "delta-table")
|
||||
deltalake.write_deltalake(path, pa.table({"id": [1, 2, 3]}))
|
||||
|
||||
catalog = _FakeCatalog(ResolvedSource(path=path))
|
||||
ds = ray.data.read_delta("main.db.tbl", catalog=catalog)
|
||||
|
||||
assert sorted(r["id"] for r in ds.take_all()) == [1, 2, 3]
|
||||
assert catalog.calls == [("main.db.tbl", ReaderFormat.DELTA)]
|
||||
|
||||
|
||||
def test_read_iceberg_uses_catalog_resolved_kwargs():
|
||||
catalog = _FakeCatalog(
|
||||
ResolvedSource(catalog_kwargs={"type": "rest", "uri": "u", "token": "tk"})
|
||||
)
|
||||
with mock.patch(
|
||||
"ray.data._internal.datasource.iceberg_datasource.IcebergDatasource"
|
||||
) as ds_cls, mock.patch("ray.data.read_api.read_datasource"):
|
||||
ray.data.read_iceberg(table_identifier="main.db.tbl", catalog=catalog)
|
||||
|
||||
_, kwargs = ds_cls.call_args
|
||||
assert kwargs["catalog_kwargs"] == {"type": "rest", "uri": "u", "token": "tk"}
|
||||
assert catalog.calls == [("main.db.tbl", ReaderFormat.ICEBERG)]
|
||||
|
||||
|
||||
def test_read_iceberg_explicit_catalog_kwargs_take_precedence():
|
||||
# When both catalog and catalog_kwargs are given, catalog is ignored.
|
||||
catalog = _FakeCatalog(ResolvedSource(catalog_kwargs={"type": "rest", "uri": "u"}))
|
||||
with mock.patch(
|
||||
"ray.data._internal.datasource.iceberg_datasource.IcebergDatasource"
|
||||
) as ds_cls, mock.patch("ray.data.read_api.read_datasource"):
|
||||
ray.data.read_iceberg(
|
||||
table_identifier="main.db.tbl",
|
||||
catalog=catalog,
|
||||
catalog_kwargs={"type": "sql", "uri": "explicit"},
|
||||
)
|
||||
|
||||
_, kwargs = ds_cls.call_args
|
||||
assert kwargs["catalog_kwargs"] == {"type": "sql", "uri": "explicit"}
|
||||
assert catalog.calls == [] # catalog was not consulted
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Serialization
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_unity_catalog_is_picklable(uc_catalog):
|
||||
restored = pickle.loads(pickle.dumps(uc_catalog))
|
||||
assert isinstance(restored, DatabricksUnityCatalog)
|
||||
assert restored.region == "us-west-2"
|
||||
|
||||
|
||||
def test_resolved_source_with_filesystem_is_picklable():
|
||||
fs = pafs.S3FileSystem(
|
||||
access_key="AKIA", secret_key="secret", session_token="t", region="us-west-2"
|
||||
)
|
||||
src = ResolvedSource(path="s3://b/p", filesystem=fs, data_format=ReaderFormat.DELTA)
|
||||
restored = pickle.loads(pickle.dumps(src))
|
||||
assert restored.path == "s3://b/p"
|
||||
assert isinstance(restored.filesystem, pafs.S3FileSystem)
|
||||
assert restored.data_format is ReaderFormat.DELTA
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Deprecated read_unity_catalog shim
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_read_unity_catalog_deprecation_delegates():
|
||||
with mock.patch("ray.data.read_api.read_delta") as read_delta:
|
||||
with pytest.warns(DeprecationWarning, match="read_unity_catalog"):
|
||||
ray.data.read_unity_catalog(
|
||||
table="main.db.tbl",
|
||||
url="https://h.databricks.com",
|
||||
token="t",
|
||||
data_format="delta",
|
||||
)
|
||||
|
||||
read_delta.assert_called_once()
|
||||
_, kwargs = read_delta.call_args
|
||||
assert isinstance(kwargs["catalog"], DatabricksUnityCatalog)
|
||||
|
||||
|
||||
def test_read_unity_catalog_infers_format_from_cred_url():
|
||||
# Metadata omits both data_source_format and storage_location; the vended
|
||||
# credential URL extension must still identify the format.
|
||||
creds = GenerateTemporaryTableCredentialResponse(url="s3://bucket/data.parquet")
|
||||
patcher = _mock_uc_sdk(data_source_format=None, storage_location=None, creds=creds)
|
||||
try:
|
||||
with mock.patch("ray.data.read_api.read_parquet") as read_parquet, pytest.warns(
|
||||
DeprecationWarning
|
||||
):
|
||||
ray.data.read_unity_catalog(
|
||||
table="main.db.tbl", url="https://h.databricks.com", token="t"
|
||||
)
|
||||
finally:
|
||||
patcher.stop()
|
||||
|
||||
read_parquet.assert_called_once()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,283 @@
|
||||
import pyarrow as pa
|
||||
import pyarrow.compute as pc
|
||||
|
||||
import ray
|
||||
from ray.data._internal.compute import ActorPoolStrategy
|
||||
from ray.data._internal.logical.interfaces import LogicalPlan
|
||||
from ray.data._internal.logical.operators import (
|
||||
CSE_TEMP_COLUMN_PREFIX,
|
||||
Project,
|
||||
)
|
||||
from ray.data._internal.logical.operators.input_data_operator import InputData
|
||||
from ray.data._internal.logical.optimizers import LogicalOptimizer
|
||||
from ray.data._internal.logical.rules import CommonSubExprElimination
|
||||
from ray.data.datatype import DataType
|
||||
from ray.data.expressions import (
|
||||
AliasExpr,
|
||||
BinaryExpr,
|
||||
PyArrowComputeUDFExpr,
|
||||
UDFExpr,
|
||||
col,
|
||||
udf,
|
||||
)
|
||||
|
||||
|
||||
def _input_op():
|
||||
return InputData(input_data=[])
|
||||
|
||||
|
||||
def _apply_cse(project: Project) -> Project:
|
||||
plan = LogicalPlan(project, ray.data.DataContext.get_current())
|
||||
optimized = CommonSubExprElimination().apply(plan)
|
||||
assert isinstance(optimized.dag, Project)
|
||||
return optimized.dag
|
||||
|
||||
|
||||
def _temp_names(project: Project) -> list[str]:
|
||||
return [expr.name for expr in project.get_common_sub_exprs()]
|
||||
|
||||
|
||||
@udf(return_dtype=DataType.int64())
|
||||
def add_one(x: pa.Array) -> pa.Array:
|
||||
return pc.add(x, 1)
|
||||
|
||||
|
||||
def _unwrap_alias(expr):
|
||||
return expr.expr if isinstance(expr, AliasExpr) else expr
|
||||
|
||||
|
||||
def test_repeated_udf_in_one_project():
|
||||
subexpr = add_one(col("a"))
|
||||
project = Project(
|
||||
exprs=[(subexpr + subexpr + subexpr).alias("y")],
|
||||
input_dependencies=[_input_op()],
|
||||
)
|
||||
|
||||
optimized = _apply_cse(project)
|
||||
|
||||
assert len(optimized.get_common_sub_exprs()) == 1
|
||||
temp_name = _temp_names(optimized)[0]
|
||||
assert temp_name.startswith(CSE_TEMP_COLUMN_PREFIX)
|
||||
assert isinstance(_unwrap_alias(optimized.get_common_sub_exprs()[0]), UDFExpr)
|
||||
assert isinstance(optimized.exprs[0], AliasExpr)
|
||||
assert optimized.exprs[0].name == "y"
|
||||
assert temp_name in repr(optimized.exprs[0])
|
||||
|
||||
|
||||
def test_structurally_equal_separately_constructed_udf_calls():
|
||||
project = Project(
|
||||
exprs=[(add_one(col("a")) + add_one(col("a"))).alias("y")],
|
||||
input_dependencies=[_input_op()],
|
||||
)
|
||||
|
||||
optimized = _apply_cse(project)
|
||||
|
||||
assert len(optimized.get_common_sub_exprs()) == 1
|
||||
|
||||
|
||||
def test_nested_common_expressions_materialize_bottom_up():
|
||||
leaf_1 = add_one(col("a"))
|
||||
leaf_2 = add_one(col("a"))
|
||||
parent_1 = leaf_1 + leaf_2
|
||||
|
||||
leaf_3 = add_one(col("a"))
|
||||
leaf_4 = add_one(col("a"))
|
||||
parent_2 = leaf_3 + leaf_4
|
||||
|
||||
project = Project(
|
||||
exprs=[(parent_1 + parent_2).alias("y")],
|
||||
input_dependencies=[_input_op()],
|
||||
)
|
||||
|
||||
optimized = _apply_cse(project)
|
||||
|
||||
common_exprs = optimized.get_common_sub_exprs()
|
||||
assert len(common_exprs) == 2
|
||||
first_temp, second_temp = _temp_names(optimized)
|
||||
assert isinstance(_unwrap_alias(common_exprs[0]), UDFExpr)
|
||||
assert first_temp in repr(common_exprs[1])
|
||||
assert second_temp in repr(optimized.exprs[0])
|
||||
|
||||
|
||||
def test_alias_root_is_ignored_but_child_is_extracted():
|
||||
project = Project(
|
||||
exprs=[
|
||||
add_one(col("a")).alias("x"),
|
||||
add_one(col("a")).alias("y"),
|
||||
],
|
||||
input_dependencies=[_input_op()],
|
||||
)
|
||||
|
||||
optimized = _apply_cse(project)
|
||||
|
||||
assert len(optimized.get_common_sub_exprs()) == 1
|
||||
assert isinstance(_unwrap_alias(optimized.get_common_sub_exprs()[0]), UDFExpr)
|
||||
assert [expr.name for expr in optimized.exprs] == ["x", "y"]
|
||||
|
||||
|
||||
def test_columns_and_literals_alone_are_not_materialized():
|
||||
project = Project(
|
||||
exprs=[
|
||||
col("a").alias("a1"),
|
||||
col("a").alias("a2"),
|
||||
(col("b") + 1).alias("b1"),
|
||||
(col("b") + 1).alias("b2"),
|
||||
],
|
||||
input_dependencies=[_input_op()],
|
||||
)
|
||||
|
||||
optimized = _apply_cse(project)
|
||||
|
||||
assert len(optimized.get_common_sub_exprs()) == 1
|
||||
common_inner = _unwrap_alias(optimized.get_common_sub_exprs()[0])
|
||||
assert isinstance(common_inner, BinaryExpr)
|
||||
|
||||
|
||||
def test_pyarrow_compute_udf_reuse():
|
||||
project = Project(
|
||||
exprs=[(col("a").abs() + col("a").abs()).alias("y")],
|
||||
input_dependencies=[_input_op()],
|
||||
)
|
||||
|
||||
optimized = _apply_cse(project)
|
||||
|
||||
assert len(optimized.get_common_sub_exprs()) == 1
|
||||
common_inner = _unwrap_alias(optimized.get_common_sub_exprs()[0])
|
||||
assert isinstance(common_inner, PyArrowComputeUDFExpr)
|
||||
|
||||
|
||||
def test_output_schema_uses_visible_expressions_only(ray_start_regular_shared_2_cpus):
|
||||
ds = ray.data.from_arrow(pa.table({"a": [1, 2]}))
|
||||
ds = ds.with_column("x", add_one(col("a")))
|
||||
ds = ds.with_column("y", col("x") + col("x"))
|
||||
ds = ds.select_columns(["y"])
|
||||
|
||||
optimized = LogicalOptimizer().optimize(ds._logical_plan)
|
||||
project = optimized.dag
|
||||
assert isinstance(project, Project)
|
||||
assert project.get_common_sub_exprs()
|
||||
assert project.infer_schema().names == ["y"]
|
||||
assert all(
|
||||
not name.startswith(CSE_TEMP_COLUMN_PREFIX)
|
||||
for name in project.infer_schema().names
|
||||
)
|
||||
|
||||
|
||||
def test_compute_strategy_is_preserved():
|
||||
@udf(return_dtype=DataType.int64())
|
||||
class AddOffset:
|
||||
def __init__(self, offset):
|
||||
self.offset = offset
|
||||
|
||||
def __call__(self, x: pa.Array) -> pa.Array:
|
||||
return pc.add(x, self.offset)
|
||||
|
||||
add_ten = AddOffset(10)
|
||||
subexpr = add_ten(col("a"))
|
||||
project = Project(
|
||||
exprs=[(subexpr + subexpr).alias("y")],
|
||||
input_dependencies=[_input_op()],
|
||||
)
|
||||
|
||||
optimized = _apply_cse(project)
|
||||
|
||||
assert isinstance(optimized.compute, ActorPoolStrategy)
|
||||
|
||||
|
||||
def test_cse_rule_is_idempotent():
|
||||
project = Project(
|
||||
exprs=[(add_one(col("a")) + add_one(col("a"))).alias("y")],
|
||||
input_dependencies=[_input_op()],
|
||||
)
|
||||
|
||||
once = _apply_cse(project)
|
||||
twice = CommonSubExprElimination._try_optimize_project(once)
|
||||
|
||||
assert twice is once
|
||||
|
||||
|
||||
def test_logical_post_optimize_cse_executes_without_pushing_temps_into_read(
|
||||
ray_start_regular_shared_2_cpus,
|
||||
):
|
||||
ds = ray.data.from_arrow(pa.table({"a": [1, 2]}))
|
||||
ds = ds.with_column("x", add_one(col("a")))
|
||||
ds = ds.with_column("y", col("x") + col("x"))
|
||||
ds = ds.select_columns(["y"])
|
||||
|
||||
optimized = LogicalOptimizer().optimize(ds._logical_plan)
|
||||
project = optimized.dag
|
||||
assert isinstance(project, Project)
|
||||
assert project.get_common_sub_exprs()
|
||||
assert ds.take_all() == [{"y": 4}, {"y": 6}]
|
||||
|
||||
|
||||
def test_udf_is_called_once_per_block(ray_start_regular_shared_2_cpus):
|
||||
@ray.remote
|
||||
class Counter:
|
||||
def __init__(self):
|
||||
self.count = 0
|
||||
|
||||
def inc(self, n):
|
||||
self.count += n
|
||||
|
||||
def get(self):
|
||||
return self.count
|
||||
|
||||
counter = Counter.remote()
|
||||
|
||||
@udf(return_dtype=DataType.int64())
|
||||
def counted_add_one(x: pa.Array) -> pa.Array:
|
||||
ray.get(counter.inc.remote(1))
|
||||
return pc.add(x, 1)
|
||||
|
||||
expr = counted_add_one(col("id"))
|
||||
ds = ray.data.range(6, override_num_blocks=3)
|
||||
ds = ds.with_column("x", expr)
|
||||
ds = ds.with_column("y", col("x") + col("x") + col("x"))
|
||||
ds = ds.select_columns(["y"])
|
||||
|
||||
assert ds.take_all() == [
|
||||
{"y": 3},
|
||||
{"y": 6},
|
||||
{"y": 9},
|
||||
{"y": 12},
|
||||
{"y": 15},
|
||||
{"y": 18},
|
||||
]
|
||||
assert ray.get(counter.get.remote()) == 3
|
||||
|
||||
|
||||
def test_callable_class_udf_still_initializes(ray_start_regular_shared_2_cpus):
|
||||
@udf(return_dtype=DataType.int64())
|
||||
class AddOffset:
|
||||
def __init__(self, offset):
|
||||
self.offset = offset
|
||||
|
||||
def __call__(self, x: pa.Array) -> pa.Array:
|
||||
return pc.add(x, self.offset)
|
||||
|
||||
add_ten = AddOffset(10)
|
||||
expr = add_ten(col("id"))
|
||||
ds = ray.data.range(4, override_num_blocks=2)
|
||||
ds = ds.with_column("x", expr)
|
||||
ds = ds.with_column("y", col("x") + col("x"))
|
||||
ds = ds.select_columns(["y"])
|
||||
|
||||
rows = ds.take_all()
|
||||
assert rows == [
|
||||
{"y": 20},
|
||||
{"y": 22},
|
||||
{"y": 24},
|
||||
{"y": 26},
|
||||
]
|
||||
assert all(
|
||||
not any(key.startswith(CSE_TEMP_COLUMN_PREFIX) for key in row) for row in rows
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,884 @@
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data._internal.block_builder import BlockBuilder
|
||||
from ray.data._internal.datasource.csv_datasink import CSVDatasink
|
||||
from ray.data._internal.datasource.csv_datasource import CSVDatasource
|
||||
from ray.data._internal.datasource.range_datasource import RangeDatasource
|
||||
from ray.data._internal.execution.interfaces.ref_bundle import (
|
||||
_ref_bundles_iterator_to_block_refs_list,
|
||||
)
|
||||
from ray.data.block import BlockAccessor
|
||||
from ray.data.dataset import Dataset, MaterializedDataset
|
||||
from ray.data.datasource.util import (
|
||||
_validate_head_node_resources_for_local_scheduling,
|
||||
)
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.data.tests.conftest import (
|
||||
CoreExecutionMetrics,
|
||||
assert_core_execution_metrics_equals,
|
||||
get_initial_core_execution_metrics_snapshot,
|
||||
)
|
||||
from ray.data.tests.util import column_udf, extract_values
|
||||
from ray.tests.conftest import * # noqa
|
||||
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
||||
|
||||
|
||||
def test_schema(ray_start_regular):
|
||||
last_snapshot = get_initial_core_execution_metrics_snapshot()
|
||||
|
||||
ds2 = ray.data.range(10, override_num_blocks=10)
|
||||
ds3 = ds2.repartition(5)
|
||||
ds3 = ds3.materialize()
|
||||
last_snapshot = assert_core_execution_metrics_equals(
|
||||
CoreExecutionMetrics(
|
||||
task_count={
|
||||
"ReadRange": 10,
|
||||
"reduce": 5,
|
||||
}
|
||||
),
|
||||
last_snapshot,
|
||||
)
|
||||
|
||||
ds4 = ds3.map(lambda x: {"a": "hi", "b": 1.0}).limit(5).repartition(1)
|
||||
ds4 = ds4.materialize()
|
||||
last_snapshot = assert_core_execution_metrics_equals(
|
||||
CoreExecutionMetrics(
|
||||
task_count={
|
||||
"Map(<lambda>)": lambda count: count <= 5,
|
||||
"slice_fn": 1,
|
||||
"reduce": 1,
|
||||
}
|
||||
),
|
||||
last_snapshot,
|
||||
)
|
||||
|
||||
ds2_schema = ds2.schema(fetch_if_missing=False)
|
||||
assert ds2_schema is not None
|
||||
assert ds2_schema.names == ["id"]
|
||||
assert not isinstance(ds2, MaterializedDataset)
|
||||
last_snapshot = assert_core_execution_metrics_equals(
|
||||
CoreExecutionMetrics(task_count={}), last_snapshot
|
||||
)
|
||||
|
||||
ds3_schema = ds3.schema(fetch_if_missing=False)
|
||||
assert ds3_schema is not None
|
||||
assert ds3_schema.names == ["id"]
|
||||
assert isinstance(ds3, MaterializedDataset)
|
||||
last_snapshot = assert_core_execution_metrics_equals(
|
||||
CoreExecutionMetrics(task_count={}), last_snapshot
|
||||
)
|
||||
ds4_schema = ds4.schema(fetch_if_missing=False)
|
||||
assert ds4_schema is not None
|
||||
assert ds4_schema.names == ["a", "b"]
|
||||
assert isinstance(ds4, MaterializedDataset)
|
||||
last_snapshot = assert_core_execution_metrics_equals(
|
||||
CoreExecutionMetrics(task_count={}), last_snapshot
|
||||
)
|
||||
|
||||
|
||||
def test_schema_no_execution(ray_start_regular):
|
||||
last_snapshot = get_initial_core_execution_metrics_snapshot()
|
||||
ds = ray.data.range(100, override_num_blocks=10)
|
||||
last_snapshot = assert_core_execution_metrics_equals(
|
||||
CoreExecutionMetrics(task_count={}),
|
||||
last_snapshot,
|
||||
)
|
||||
# We do not kick off the read task by default.
|
||||
schema = ds.schema()
|
||||
assert schema.names == ["id"]
|
||||
|
||||
# Fetching the schema does not trigger execution, since
|
||||
# the schema is known beforehand for RangeDatasource.
|
||||
last_snapshot = assert_core_execution_metrics_equals(
|
||||
CoreExecutionMetrics(task_count={}), last_snapshot
|
||||
)
|
||||
# Fetching the schema should not trigger execution of extra read tasks.
|
||||
|
||||
|
||||
def test_schema_cached(ray_start_regular):
|
||||
def check_schema_cached(ds, expected_task_count, last_snapshot):
|
||||
schema = ds.schema()
|
||||
last_snapshot = assert_core_execution_metrics_equals(
|
||||
CoreExecutionMetrics(expected_task_count), last_snapshot
|
||||
)
|
||||
assert schema.names == ["a"]
|
||||
cached_schema = ds.schema(fetch_if_missing=False)
|
||||
assert cached_schema is not None
|
||||
assert schema == cached_schema
|
||||
last_snapshot = assert_core_execution_metrics_equals(
|
||||
CoreExecutionMetrics({}), last_snapshot
|
||||
)
|
||||
return last_snapshot
|
||||
|
||||
last_snapshot = get_initial_core_execution_metrics_snapshot()
|
||||
ds = ray.data.from_items([{"a": i} for i in range(100)], override_num_blocks=10)
|
||||
last_snapshot = check_schema_cached(ds, {}, last_snapshot)
|
||||
|
||||
# Add a map_batches operator so that we are forced to compute the schema.
|
||||
ds = ds.map_batches(lambda x: x)
|
||||
last_snapshot = check_schema_cached(
|
||||
ds,
|
||||
{
|
||||
"MapBatches(<lambda>)": lambda count: count <= 5,
|
||||
"slice_fn": 1,
|
||||
},
|
||||
last_snapshot,
|
||||
)
|
||||
|
||||
|
||||
def test_avoid_placement_group_capture(shutdown_only):
|
||||
ray.init(num_cpus=2)
|
||||
|
||||
@ray.remote
|
||||
def run():
|
||||
ds = ray.data.range(5)
|
||||
assert sorted(
|
||||
extract_values("id", ds.map(column_udf("id", lambda x: x + 1)).take())
|
||||
) == [1, 2, 3, 4, 5]
|
||||
assert ds.count() == 5
|
||||
assert sorted(extract_values("id", ds.iter_rows())) == [0, 1, 2, 3, 4]
|
||||
|
||||
pg = ray.util.placement_group([{"CPU": 1}])
|
||||
ray.get(
|
||||
run.options(
|
||||
scheduling_strategy=PlacementGroupSchedulingStrategy(
|
||||
placement_group=pg, placement_group_capture_child_tasks=True
|
||||
)
|
||||
).remote()
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def remove_named_placement_groups():
|
||||
yield
|
||||
for info in ray.util.placement_group_table().values():
|
||||
if info["name"]:
|
||||
pg = ray.util.get_placement_group(info["name"])
|
||||
ray.util.remove_placement_group(pg)
|
||||
|
||||
|
||||
def test_ray_remote_args_fn(shutdown_only, remove_named_placement_groups):
|
||||
ray.init()
|
||||
|
||||
pg = ray.util.placement_group([{"CPU": 1}], name="test_pg")
|
||||
|
||||
def ray_remote_args_fn():
|
||||
scheduling_strategy = PlacementGroupSchedulingStrategy(placement_group=pg)
|
||||
return {"scheduling_strategy": scheduling_strategy}
|
||||
|
||||
class ActorClass:
|
||||
def __call__(self, batch):
|
||||
assert ray.util.get_current_placement_group() == pg
|
||||
return batch
|
||||
|
||||
ray.data.range(1).map_batches(
|
||||
ActorClass, concurrency=1, ray_remote_args_fn=ray_remote_args_fn
|
||||
).take_all()
|
||||
|
||||
|
||||
def test_dataset_lineage_serialization(shutdown_only):
|
||||
ray.init()
|
||||
ds = ray.data.range(10)
|
||||
ds = ds.map(column_udf("id", lambda x: x + 1))
|
||||
ds = ds.map(column_udf("id", lambda x: x + 1))
|
||||
ds = ds.random_shuffle()
|
||||
uuid = ds._get_uuid()
|
||||
plan_uuid = ds._uuid
|
||||
|
||||
serialized_ds = ds.serialize_lineage()
|
||||
|
||||
ray.shutdown()
|
||||
ray.init()
|
||||
|
||||
ds = Dataset.deserialize_lineage(serialized_ds)
|
||||
# Check Dataset state.
|
||||
assert ds._get_uuid() == uuid
|
||||
assert ds._uuid == plan_uuid
|
||||
# Check Dataset content.
|
||||
assert ds.count() == 10
|
||||
assert sorted(extract_values("id", ds.take())) == list(range(2, 12))
|
||||
|
||||
|
||||
def test_dataset_lineage_serialization_unsupported(shutdown_only):
|
||||
ray.init()
|
||||
# In-memory data sources not supported.
|
||||
ds = ray.data.from_items(list(range(10)))
|
||||
ds = ds.map(column_udf("item", lambda x: x + 1))
|
||||
ds = ds.map(column_udf("item", lambda x: x + 1))
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
ds.serialize_lineage()
|
||||
|
||||
# In-memory data source unions not supported.
|
||||
ds = ray.data.from_items(list(range(10)))
|
||||
ds1 = ray.data.from_items(list(range(10, 20)))
|
||||
ds2 = ds.union(ds1)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
ds2.serialize_lineage()
|
||||
|
||||
# Lazy read unions supported.
|
||||
ds = ray.data.range(10)
|
||||
ds1 = ray.data.range(20)
|
||||
ds2 = ds.union(ds1)
|
||||
|
||||
serialized_ds = ds2.serialize_lineage()
|
||||
ds3 = Dataset.deserialize_lineage(serialized_ds)
|
||||
assert set(extract_values("id", ds3.take(30))) == set(
|
||||
list(range(10)) + list(range(20))
|
||||
)
|
||||
|
||||
# Zips not supported.
|
||||
ds = ray.data.from_items(list(range(10)))
|
||||
ds1 = ray.data.from_items(list(range(10, 20)))
|
||||
ds2 = ds.zip(ds1)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
ds2.serialize_lineage()
|
||||
|
||||
|
||||
def test_basic(ray_start_regular_shared):
|
||||
ds = ray.data.range(5)
|
||||
assert sorted(
|
||||
extract_values("id", ds.map(column_udf("id", lambda x: x + 1)).take())
|
||||
) == [1, 2, 3, 4, 5]
|
||||
assert ds.count() == 5
|
||||
assert sorted(extract_values("id", ds.iter_rows())) == [0, 1, 2, 3, 4]
|
||||
|
||||
|
||||
def test_range(ray_start_regular_shared):
|
||||
ds = ray.data.range(10, override_num_blocks=10)
|
||||
assert ds._logical_plan.initial_num_blocks() == 10
|
||||
assert ds.count() == 10
|
||||
assert ds.take() == [{"id": i} for i in range(10)]
|
||||
|
||||
ds = ray.data.range(10, override_num_blocks=2)
|
||||
assert ds._logical_plan.initial_num_blocks() == 2
|
||||
assert ds.count() == 10
|
||||
assert ds.take() == [{"id": i} for i in range(10)]
|
||||
|
||||
|
||||
def test_empty_dataset(ray_start_regular_shared):
|
||||
ds = ray.data.range(0)
|
||||
assert ds.count() == 0
|
||||
assert ds.size_bytes() == 0
|
||||
assert ds.schema() is None
|
||||
|
||||
ds = ray.data.range(1)
|
||||
ds = ds.filter(lambda x: x["id"] > 1)
|
||||
ds = ds.materialize()
|
||||
assert (
|
||||
str(ds)
|
||||
== "MaterializedDataset(num_blocks=1, num_rows=0, schema=Unknown schema)"
|
||||
)
|
||||
|
||||
# Test map on empty dataset.
|
||||
ds = ray.data.from_items([])
|
||||
ds = ds.map(lambda x: x)
|
||||
ds = ds.materialize()
|
||||
assert ds.count() == 0
|
||||
|
||||
# Test filter on empty dataset.
|
||||
ds = ray.data.from_items([])
|
||||
ds = ds.filter(lambda: True)
|
||||
ds = ds.materialize()
|
||||
assert ds.count() == 0
|
||||
|
||||
|
||||
@ray.remote
|
||||
class Counter:
|
||||
def __init__(self):
|
||||
self.value = 0
|
||||
|
||||
def increment(self):
|
||||
self.value += 1
|
||||
return self.value
|
||||
|
||||
|
||||
def test_cache_dataset(ray_start_regular_shared):
|
||||
|
||||
c = Counter.remote()
|
||||
|
||||
def inc(x):
|
||||
ray.get(c.increment.remote())
|
||||
return x
|
||||
|
||||
ds = ray.data.range(1)
|
||||
ds = ds.map(inc)
|
||||
assert not isinstance(ds, MaterializedDataset)
|
||||
ds2 = ds.materialize()
|
||||
assert isinstance(ds2, MaterializedDataset)
|
||||
assert not isinstance(ds, MaterializedDataset)
|
||||
|
||||
# Tests standard iteration uses the materialized blocks.
|
||||
for _ in range(10):
|
||||
ds2.take_all()
|
||||
|
||||
assert ray.get(c.increment.remote()) == 2
|
||||
|
||||
# Tests streaming iteration uses the materialized blocks.
|
||||
for _ in range(10):
|
||||
list(ds2.streaming_split(1)[0].iter_batches())
|
||||
|
||||
assert ray.get(c.increment.remote()) == 3
|
||||
|
||||
|
||||
def test_columns(ray_start_regular_shared):
|
||||
ds = ray.data.range(1)
|
||||
assert ds.columns() == ds.schema().names
|
||||
assert ds.columns() == ["id"]
|
||||
|
||||
ds = ds.map(lambda x: x)
|
||||
assert ds.columns(fetch_if_missing=False) is None
|
||||
|
||||
|
||||
def test_schema_repr(ray_start_regular_shared):
|
||||
ds = ray.data.from_items([{"text": "spam", "number": 0}])
|
||||
# fmt: off
|
||||
expected_repr = (
|
||||
"Column Type\n"
|
||||
"------ ----\n"
|
||||
"text string\n"
|
||||
"number int64"
|
||||
)
|
||||
# fmt:on
|
||||
assert repr(ds.schema()) == expected_repr
|
||||
|
||||
ds = ray.data.from_items([{"long_column_name": "spam"}])
|
||||
# fmt: off
|
||||
expected_repr = (
|
||||
"Column Type\n"
|
||||
"------ ----\n"
|
||||
"long_column_name string"
|
||||
)
|
||||
# fmt: on
|
||||
assert repr(ds.schema()) == expected_repr
|
||||
|
||||
|
||||
def _check_none_computed(ds):
|
||||
# In streaming executor, ds.take() will not invoke partial execution
|
||||
# in LazyBlocklist.
|
||||
assert not ds._has_computed_output()
|
||||
|
||||
|
||||
def test_lazy_loading_exponential_rampup(ray_start_regular_shared):
|
||||
ds = ray.data.range(100, override_num_blocks=20)
|
||||
_check_none_computed(ds)
|
||||
assert extract_values("id", ds.take(10)) == list(range(10))
|
||||
_check_none_computed(ds)
|
||||
assert extract_values("id", ds.take(20)) == list(range(20))
|
||||
_check_none_computed(ds)
|
||||
assert extract_values("id", ds.take(30)) == list(range(30))
|
||||
_check_none_computed(ds)
|
||||
assert extract_values("id", ds.take(50)) == list(range(50))
|
||||
_check_none_computed(ds)
|
||||
assert extract_values("id", ds.take(100)) == list(range(100))
|
||||
_check_none_computed(ds)
|
||||
|
||||
|
||||
def test_dataset_repr_not_materialized(ray_start_regular_shared, restore_data_context):
|
||||
ds = ray.data.range(5)
|
||||
assert repr(ds) == (
|
||||
"shape: (5, 1)\n"
|
||||
"╭───────╮\n"
|
||||
"│ id │\n"
|
||||
"│ --- │\n"
|
||||
"│ int64 │\n"
|
||||
"╰───────╯\n"
|
||||
"(Dataset isn't materialized)"
|
||||
)
|
||||
|
||||
|
||||
def test_dataset_repr_materialized(ray_start_regular_shared, restore_data_context):
|
||||
materialized = ray.data.range(5).materialize()
|
||||
assert repr(materialized) == (
|
||||
"shape: (5, 1)\n"
|
||||
"╭───────╮\n"
|
||||
"│ id │\n"
|
||||
"│ --- │\n"
|
||||
"│ int64 │\n"
|
||||
"╞═══════╡\n"
|
||||
"│ 0 │\n"
|
||||
"│ 1 │\n"
|
||||
"│ 2 │\n"
|
||||
"│ 3 │\n"
|
||||
"│ 4 │\n"
|
||||
"╰───────╯\n"
|
||||
"(Showing 5 of 5 rows)"
|
||||
)
|
||||
|
||||
|
||||
def test_dataset_repr_gap(ray_start_regular_shared, restore_data_context):
|
||||
ds_with_gap = ray.data.range(20).materialize()
|
||||
assert repr(ds_with_gap) == (
|
||||
"shape: (20, 1)\n"
|
||||
"╭───────╮\n"
|
||||
"│ id │\n"
|
||||
"│ --- │\n"
|
||||
"│ int64 │\n"
|
||||
"╞═══════╡\n"
|
||||
"│ 0 │\n"
|
||||
"│ 1 │\n"
|
||||
"│ 2 │\n"
|
||||
"│ 3 │\n"
|
||||
"│ 4 │\n"
|
||||
"│ … │\n"
|
||||
"│ 15 │\n"
|
||||
"│ 16 │\n"
|
||||
"│ 17 │\n"
|
||||
"│ 18 │\n"
|
||||
"│ 19 │\n"
|
||||
"╰───────╯\n"
|
||||
"(Showing 10 of 20 rows)"
|
||||
)
|
||||
|
||||
|
||||
def test_dataset_explain(ray_start_regular_shared, capsys):
|
||||
ds = ray.data.range(10, override_num_blocks=10)
|
||||
ds = ds.map(lambda x: x)
|
||||
|
||||
ds.explain()
|
||||
captured = capsys.readouterr()
|
||||
assert captured.out.strip() == (
|
||||
"-------- Logical Plan --------\n"
|
||||
"MapRows[Map(<lambda>)]\n"
|
||||
"+- Read[ReadRange]\n"
|
||||
"\n-------- Logical Plan (Optimized) --------\n"
|
||||
"MapRows[Map(<lambda>)]\n"
|
||||
"+- Read[ReadRange]\n"
|
||||
"\n-------- Physical Plan --------\n"
|
||||
"TaskPoolMapOperator[Map(<lambda>)]\n"
|
||||
"+- TaskPoolMapOperator[ReadRange]\n"
|
||||
" +- InputDataBuffer[Input]\n"
|
||||
"\n-------- Physical Plan (Optimized) --------\n"
|
||||
"TaskPoolMapOperator[ReadRange->Map(<lambda>)]\n"
|
||||
"+- InputDataBuffer[Input]"
|
||||
)
|
||||
|
||||
ds = ds.filter(lambda x: x["id"] > 0)
|
||||
ds.explain()
|
||||
captured = capsys.readouterr()
|
||||
assert captured.out.strip() == (
|
||||
"-------- Logical Plan --------\n"
|
||||
"Filter[Filter(<lambda>)]\n"
|
||||
"+- MapRows[Map(<lambda>)]\n"
|
||||
" +- Read[ReadRange]\n"
|
||||
"\n-------- Logical Plan (Optimized) --------\n"
|
||||
"Filter[Filter(<lambda>)]\n"
|
||||
"+- MapRows[Map(<lambda>)]\n"
|
||||
" +- Read[ReadRange]\n"
|
||||
"\n-------- Physical Plan --------\n"
|
||||
"TaskPoolMapOperator[Filter(<lambda>)]\n"
|
||||
"+- TaskPoolMapOperator[Map(<lambda>)]\n"
|
||||
" +- TaskPoolMapOperator[ReadRange]\n"
|
||||
" +- InputDataBuffer[Input]\n"
|
||||
"\n-------- Physical Plan (Optimized) --------\n"
|
||||
"TaskPoolMapOperator[ReadRange->Map(<lambda>)->Filter(<lambda>)]\n"
|
||||
"+- InputDataBuffer[Input]"
|
||||
)
|
||||
ds = ds.random_shuffle().map(lambda x: x)
|
||||
ds.explain()
|
||||
captured = capsys.readouterr()
|
||||
assert captured.out.strip() == (
|
||||
"-------- Logical Plan --------\n"
|
||||
"MapRows[Map(<lambda>)]\n"
|
||||
"+- RandomShuffle[RandomShuffle]\n"
|
||||
" +- Filter[Filter(<lambda>)]\n"
|
||||
" +- MapRows[Map(<lambda>)]\n"
|
||||
" +- Read[ReadRange]\n"
|
||||
"\n-------- Logical Plan (Optimized) --------\n"
|
||||
"MapRows[Map(<lambda>)]\n"
|
||||
"+- RandomShuffle[RandomShuffle]\n"
|
||||
" +- Filter[Filter(<lambda>)]\n"
|
||||
" +- MapRows[Map(<lambda>)]\n"
|
||||
" +- Read[ReadRange]\n"
|
||||
"\n-------- Physical Plan --------\n"
|
||||
"TaskPoolMapOperator[Map(<lambda>)]\n"
|
||||
"+- AllToAllOperator[RandomShuffle]\n"
|
||||
" +- TaskPoolMapOperator[Filter(<lambda>)]\n"
|
||||
" +- TaskPoolMapOperator[Map(<lambda>)]\n"
|
||||
" +- TaskPoolMapOperator[ReadRange]\n"
|
||||
" +- InputDataBuffer[Input]\n"
|
||||
"\n-------- Physical Plan (Optimized) --------\n"
|
||||
"TaskPoolMapOperator[Map(<lambda>)]\n"
|
||||
"+- AllToAllOperator[ReadRange->Map(<lambda>)->Filter(<lambda>)->RandomShuffle]\n"
|
||||
" +- InputDataBuffer[Input]"
|
||||
)
|
||||
|
||||
|
||||
def test_convert_types(ray_start_regular_shared):
|
||||
plain_ds = ray.data.range(1)
|
||||
arrow_ds = plain_ds.map(lambda x: {"a": x["id"]})
|
||||
assert arrow_ds.take() == [{"a": 0}]
|
||||
assert "dict" in str(arrow_ds.map(lambda x: {"out": str(type(x))}).take()[0])
|
||||
|
||||
arrow_ds = ray.data.range(1)
|
||||
assert arrow_ds.map(lambda x: {"out": "plain_{}".format(x["id"])}).take() == [
|
||||
{"out": "plain_0"}
|
||||
]
|
||||
assert arrow_ds.map(lambda x: {"a": (x["id"],)}).take() == [{"a": [0]}]
|
||||
|
||||
|
||||
def test_take_batch(ray_start_regular_shared):
|
||||
ds = ray.data.range(10, override_num_blocks=2)
|
||||
assert ds.take_batch(3)["id"].tolist() == [0, 1, 2]
|
||||
assert ds.take_batch(6)["id"].tolist() == [0, 1, 2, 3, 4, 5]
|
||||
assert ds.take_batch(100)["id"].tolist() == [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
|
||||
assert isinstance(ds.take_batch(3, batch_format="pandas"), pd.DataFrame)
|
||||
assert isinstance(ds.take_batch(3, batch_format="numpy"), dict)
|
||||
|
||||
ds = ray.data.range_tensor(10, override_num_blocks=2)
|
||||
assert np.all(ds.take_batch(3)["data"] == np.array([[0], [1], [2]]))
|
||||
assert isinstance(ds.take_batch(3, batch_format="pandas"), pd.DataFrame)
|
||||
assert isinstance(ds.take_batch(3, batch_format="numpy"), dict)
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
ray.data.range(0).take_batch()
|
||||
|
||||
|
||||
def test_take_all(ray_start_regular_shared):
|
||||
assert extract_values("id", ray.data.range(5).take_all()) == [0, 1, 2, 3, 4]
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
assert ray.data.range(5).take_all(4)
|
||||
|
||||
|
||||
def test_union(ray_start_regular_shared, restore_data_context):
|
||||
# Set aggregator CPU to 0 to avoid deadlock in resource-constrained test env.
|
||||
# Without this, the shuffle task (1 CPU) + aggregator actor (~0.25 CPU) would
|
||||
# exceed the 1 CPU available in the test cluster, causing scheduler deadlock.
|
||||
restore_data_context.hash_aggregate_operator_actor_num_cpus_override = 0
|
||||
|
||||
ds = ray.data.range(20, override_num_blocks=10).materialize()
|
||||
|
||||
# Test lazy union.
|
||||
ds = ds.union(ds, ds, ds, ds)
|
||||
assert ds._logical_plan.initial_num_blocks() == 50
|
||||
assert ds.count() == 100
|
||||
assert ds.sum() == 950
|
||||
|
||||
ds = ds.union(ds)
|
||||
assert ds.count() == 200
|
||||
assert ds.sum() == (950 * 2)
|
||||
|
||||
# Test materialized union.
|
||||
ds2 = ray.data.from_items([1, 2, 3, 4, 5])
|
||||
assert ds2.count() == 5
|
||||
ds2 = ds2.union(ds2)
|
||||
assert ds2.count() == 10
|
||||
ds2 = ds2.union(ds)
|
||||
assert ds2.count() == 210
|
||||
|
||||
|
||||
def test_block_builder_for_block(ray_start_regular_shared):
|
||||
# pandas dataframe
|
||||
builder = BlockBuilder.for_block(pd.DataFrame())
|
||||
b1 = pd.DataFrame({"A": [1], "B": ["a"]})
|
||||
builder.add_block(b1)
|
||||
assert builder.build().equals(b1)
|
||||
b2 = pd.DataFrame({"A": [2, 3], "B": ["c", "d"]})
|
||||
builder.add_block(b2)
|
||||
expected = pd.DataFrame({"A": [1, 2, 3], "B": ["a", "c", "d"]})
|
||||
assert builder.build().equals(expected)
|
||||
|
||||
# pyarrow table
|
||||
builder = BlockBuilder.for_block(pa.Table.from_arrays(list()))
|
||||
b1 = pa.Table.from_pydict({"A": [1], "B": ["a"]})
|
||||
builder.add_block(b1)
|
||||
builder.build().equals(b1)
|
||||
b2 = pa.Table.from_pydict({"A": [2, 3], "B": ["c", "d"]})
|
||||
builder.add_block(b2)
|
||||
expected = pa.Table.from_pydict({"A": [1, 2, 3], "B": ["a", "c", "d"]})
|
||||
builder.build().equals(expected)
|
||||
|
||||
# wrong type
|
||||
with pytest.raises(TypeError):
|
||||
BlockBuilder.for_block(str())
|
||||
|
||||
|
||||
def test_len(ray_start_regular_shared):
|
||||
ds = ray.data.range(1)
|
||||
with pytest.raises(AttributeError):
|
||||
len(ds)
|
||||
|
||||
|
||||
def test_pandas_block_select():
|
||||
df = pd.DataFrame({"one": [10, 11, 12], "two": [11, 12, 13], "three": [14, 15, 16]})
|
||||
block_accessor = BlockAccessor.for_block(df)
|
||||
|
||||
block = block_accessor.select(["two"])
|
||||
assert block.equals(df[["two"]])
|
||||
|
||||
block = block_accessor.select(["two", "one"])
|
||||
assert block.equals(df[["two", "one"]])
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
block = block_accessor.select([lambda x: x % 3, "two"])
|
||||
|
||||
|
||||
# NOTE: All tests above share a Ray cluster, while the tests below do not. These
|
||||
# tests should only be carefully reordered to retain this invariant!
|
||||
|
||||
|
||||
def test_read_write_local_node_ray_client(ray_start_cluster_enabled):
|
||||
cluster = ray_start_cluster_enabled
|
||||
cluster.add_node(num_cpus=4)
|
||||
cluster.head_node._ray_params.ray_client_server_port = "10004"
|
||||
cluster.head_node.start_ray_client_server()
|
||||
address = "ray://localhost:10004"
|
||||
|
||||
import tempfile
|
||||
|
||||
data_path = tempfile.mkdtemp()
|
||||
df = pd.DataFrame({"one": list(range(0, 10)), "two": list(range(10, 20))})
|
||||
path = os.path.join(data_path, "test.parquet")
|
||||
df.to_parquet(path)
|
||||
|
||||
# Read/write from Ray Client will result in error.
|
||||
ray.init(address)
|
||||
with pytest.raises(ValueError):
|
||||
ds = ray.data.read_parquet("local://" + path).materialize()
|
||||
ds = ray.data.from_pandas(df)
|
||||
with pytest.raises(ValueError):
|
||||
ds.write_parquet("local://" + data_path).materialize()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
sys.version_info >= (3, 12), reason="No tensorflow for Python 3.12+"
|
||||
)
|
||||
def test_read_warning_large_parallelism(
|
||||
ray_start_regular_shared, propagate_logs, caplog
|
||||
):
|
||||
with caplog.at_level(logging.WARNING, logger="ray.data.read_api"):
|
||||
ray.data.range(5000, override_num_blocks=5000).materialize()
|
||||
assert (
|
||||
"The requested number of read blocks of 5000 is "
|
||||
"more than 4x the number of available CPU slots in the cluster" in caplog.text
|
||||
), caplog.text
|
||||
|
||||
|
||||
def test_read_write_local_node(ray_start_cluster):
|
||||
cluster = ray_start_cluster
|
||||
cluster.add_node(
|
||||
resources={"bar:1": 100},
|
||||
num_cpus=10,
|
||||
_system_config={"max_direct_call_object_size": 0},
|
||||
)
|
||||
cluster.add_node(resources={"bar:2": 100}, num_cpus=10)
|
||||
cluster.add_node(resources={"bar:3": 100}, num_cpus=10)
|
||||
|
||||
ray.shutdown()
|
||||
ray.init(cluster.address)
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
|
||||
data_path = tempfile.mkdtemp()
|
||||
num_files = 5
|
||||
for idx in range(num_files):
|
||||
df = pd.DataFrame(
|
||||
{"one": list(range(idx, idx + 10)), "two": list(range(idx + 10, idx + 20))}
|
||||
)
|
||||
path = os.path.join(data_path, f"test{idx}.parquet")
|
||||
df.to_parquet(path)
|
||||
|
||||
ctx = ray.data.context.DataContext.get_current()
|
||||
ctx.read_write_local_node = True
|
||||
|
||||
def check_dataset_is_local(ds):
|
||||
bundles = ds.iter_internal_ref_bundles()
|
||||
block_refs = _ref_bundles_iterator_to_block_refs_list(bundles)
|
||||
ray.wait(block_refs, num_returns=len(block_refs), fetch_local=False)
|
||||
location_data = ray.experimental.get_object_locations(block_refs)
|
||||
locations = []
|
||||
for block in block_refs:
|
||||
locations.extend(location_data[block]["node_ids"])
|
||||
assert set(locations) == {ray.get_runtime_context().get_node_id()}
|
||||
|
||||
local_path = "local://" + data_path
|
||||
# Plain read.
|
||||
ds = ray.data.read_parquet(local_path).materialize()
|
||||
check_dataset_is_local(ds)
|
||||
|
||||
# SPREAD scheduling got overridden when read local scheme.
|
||||
ds = ray.data.read_parquet(
|
||||
local_path, ray_remote_args={"scheduling_strategy": "SPREAD"}
|
||||
).materialize()
|
||||
check_dataset_is_local(ds)
|
||||
|
||||
# With fusion.
|
||||
ds = (
|
||||
ray.data.read_parquet(local_path, override_num_blocks=1)
|
||||
.map(lambda x: x)
|
||||
.materialize()
|
||||
)
|
||||
check_dataset_is_local(ds)
|
||||
|
||||
# Write back to local scheme.
|
||||
output = os.path.join(local_path, "test_read_write_local_node")
|
||||
ds.write_parquet(output)
|
||||
assert "1 nodes used" in ds.stats(), ds.stats()
|
||||
ray.data.read_parquet(output).take_all() == ds.take_all()
|
||||
|
||||
# Mixing paths of local and non-local scheme is invalid.
|
||||
with pytest.raises(ValueError):
|
||||
ds = ray.data.read_parquet(
|
||||
[local_path + "/test1.parquet", data_path + "/test2.parquet"]
|
||||
).materialize()
|
||||
with pytest.raises(ValueError):
|
||||
ds = ray.data.read_parquet(
|
||||
[local_path + "/test1.parquet", "example://iris.parquet"]
|
||||
).materialize()
|
||||
with pytest.raises(ValueError):
|
||||
ds = ray.data.read_parquet(
|
||||
["example://iris.parquet", local_path + "/test1.parquet"]
|
||||
).materialize()
|
||||
|
||||
|
||||
def test_validate_head_node_resources_zero_head_cpu(ray_start_cluster):
|
||||
|
||||
cluster = ray_start_cluster
|
||||
cluster.add_node(num_cpus=0)
|
||||
cluster.wait_for_nodes()
|
||||
|
||||
ray.shutdown()
|
||||
ray.init(address=cluster.address)
|
||||
|
||||
with pytest.raises(ValueError, match=r"head node doesn't have enough resources"):
|
||||
_validate_head_node_resources_for_local_scheduling(
|
||||
{}, op_description="read local"
|
||||
)
|
||||
|
||||
|
||||
class FlakyCSVDatasource(CSVDatasource):
|
||||
def __init__(self, paths, **csv_datasource_kwargs):
|
||||
super().__init__(paths, **csv_datasource_kwargs)
|
||||
|
||||
self.counter = Counter.remote()
|
||||
|
||||
def _read_stream(self, f: "pa.NativeFile", path: str):
|
||||
count = self.counter.increment.remote()
|
||||
if ray.get(count) == 1:
|
||||
raise ValueError("oops")
|
||||
else:
|
||||
for block in CSVDatasource._read_stream(self, f, path):
|
||||
yield block
|
||||
|
||||
|
||||
class FlakyCSVDatasink(CSVDatasink):
|
||||
def __init__(self, path, **csv_datasink_kwargs):
|
||||
super().__init__(path, **csv_datasink_kwargs)
|
||||
|
||||
self.counter = Counter.remote()
|
||||
|
||||
def write_block_to_file(self, block: BlockAccessor, file):
|
||||
count = self.counter.increment.remote()
|
||||
if ray.get(count) == 1:
|
||||
raise ValueError("oops")
|
||||
else:
|
||||
super().write_block_to_file(block, file)
|
||||
|
||||
|
||||
def test_datasource(ray_start_regular):
|
||||
source = ray.data.datasource.RandomIntRowDatasource(n=10, num_columns=2)
|
||||
assert len(ray.data.read_datasource(source).take()) == 10
|
||||
source = RangeDatasource(n=10)
|
||||
assert extract_values(
|
||||
"value",
|
||||
ray.data.read_datasource(source).take(),
|
||||
) == list(range(10))
|
||||
|
||||
|
||||
@pytest.mark.skip(reason="")
|
||||
def test_polars_lazy_import(shutdown_only):
|
||||
import sys
|
||||
|
||||
ctx = ray.data.context.DataContext.get_current()
|
||||
|
||||
try:
|
||||
original_use_polars = ctx.use_polars
|
||||
ctx.use_polars = True
|
||||
|
||||
num_items = 100
|
||||
parallelism = 4
|
||||
ray.init(num_cpus=4)
|
||||
|
||||
@ray.remote
|
||||
def f(should_import_polars):
|
||||
# Sleep to spread the tasks.
|
||||
time.sleep(1)
|
||||
polars_imported = "polars" in sys.modules.keys()
|
||||
return polars_imported == should_import_polars
|
||||
|
||||
# We should not use polars for non-Arrow sort.
|
||||
_ = ray.data.range(num_items, override_num_blocks=parallelism).sort()
|
||||
assert all(ray.get([f.remote(False) for _ in range(parallelism)]))
|
||||
|
||||
a = range(100)
|
||||
dfs = []
|
||||
partition_size = num_items // parallelism
|
||||
for i in range(parallelism):
|
||||
dfs.append(
|
||||
pd.DataFrame({"a": a[i * partition_size : (i + 1) * partition_size]})
|
||||
)
|
||||
# At least one worker should have imported polars.
|
||||
_ = (
|
||||
ray.data.from_pandas(dfs)
|
||||
.map_batches(lambda t: t, batch_format="pyarrow", batch_size=None)
|
||||
.sort(key="a")
|
||||
.materialize()
|
||||
)
|
||||
assert any(ray.get([f.remote(True) for _ in range(parallelism)]))
|
||||
|
||||
finally:
|
||||
ctx.use_polars = original_use_polars
|
||||
|
||||
|
||||
def test_batch_formats(shutdown_only):
|
||||
ds = ray.data.range(100)
|
||||
assert isinstance(next(iter(ds.iter_batches(batch_format=None))), pa.Table)
|
||||
assert isinstance(next(iter(ds.iter_batches(batch_format="default"))), dict)
|
||||
assert isinstance(next(iter(ds.iter_batches(batch_format="pandas"))), pd.DataFrame)
|
||||
assert isinstance(next(iter(ds.iter_batches(batch_format="pyarrow"))), pa.Table)
|
||||
assert isinstance(next(iter(ds.iter_batches(batch_format="numpy"))), dict)
|
||||
|
||||
ds = ray.data.range_tensor(100)
|
||||
assert isinstance(next(iter(ds.iter_batches(batch_format=None))), pa.Table)
|
||||
assert isinstance(next(iter(ds.iter_batches(batch_format="default"))), dict)
|
||||
assert isinstance(next(iter(ds.iter_batches(batch_format="pandas"))), pd.DataFrame)
|
||||
assert isinstance(next(iter(ds.iter_batches(batch_format="pyarrow"))), pa.Table)
|
||||
assert isinstance(next(iter(ds.iter_batches(batch_format="numpy"))), dict)
|
||||
|
||||
df = pd.DataFrame({"foo": ["a", "b"], "bar": [0, 1]})
|
||||
ds = ray.data.from_pandas(df)
|
||||
assert isinstance(next(iter(ds.iter_batches(batch_format=None))), pd.DataFrame)
|
||||
assert isinstance(next(iter(ds.iter_batches(batch_format="default"))), dict)
|
||||
assert isinstance(next(iter(ds.iter_batches(batch_format="pandas"))), pd.DataFrame)
|
||||
assert isinstance(next(iter(ds.iter_batches(batch_format="pyarrow"))), pa.Table)
|
||||
assert isinstance(next(iter(ds.iter_batches(batch_format="numpy"))), dict)
|
||||
|
||||
|
||||
def test_dataset_schema_after_read_stats(ray_start_cluster):
|
||||
cluster = ray_start_cluster
|
||||
cluster.add_node(num_cpus=1)
|
||||
ray.init(cluster.address)
|
||||
cluster.add_node(num_cpus=1, resources={"foo": 1})
|
||||
ds = ray.data.read_csv(
|
||||
"example://iris.csv", ray_remote_args={"resources": {"foo": 1}}
|
||||
)
|
||||
schema = ds.schema()
|
||||
ds.stats()
|
||||
assert schema == ds.schema()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,42 @@
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
|
||||
|
||||
def test_write_file_retry_on_errors_emits_deprecation_warning(caplog):
|
||||
ctx = ray.data.DataContext.get_current()
|
||||
with pytest.warns(DeprecationWarning):
|
||||
ctx.write_file_retry_on_errors = []
|
||||
|
||||
|
||||
def test_data_context_current_context_manager():
|
||||
import copy
|
||||
|
||||
from ray.data.context import DataContext
|
||||
|
||||
original = DataContext.get_current()
|
||||
ctx1 = copy.deepcopy(original)
|
||||
ctx1.set_config("level", "1")
|
||||
|
||||
ctx2 = copy.deepcopy(original)
|
||||
ctx2.set_config("level", "2")
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
with DataContext.current(ctx1):
|
||||
assert DataContext.get_current() is ctx1
|
||||
# Nested context manager
|
||||
with DataContext.current(ctx2):
|
||||
assert DataContext.get_current().get_config("level") == "2"
|
||||
|
||||
assert DataContext.get_current().get_config("level") == "1"
|
||||
|
||||
# Test that raising will reset context too
|
||||
raise ValueError("boom")
|
||||
|
||||
assert DataContext.get_current() is original
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,198 @@
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray._common.test_utils import run_string_as_driver
|
||||
from ray.data.block import BlockMetadata
|
||||
from ray.data.context import DataContext
|
||||
from ray.data.datasource import Datasource, ReadTask
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
# Auto-use `restore_data_context` for each test.
|
||||
pytestmark = pytest.mark.usefixtures("restore_data_context")
|
||||
|
||||
|
||||
def test_context_saved_when_dataset_created(
|
||||
ray_start_regular_shared,
|
||||
):
|
||||
ctx = DataContext.get_current()
|
||||
ctx.set_config("foo", 1)
|
||||
d1 = ray.data.range(10)
|
||||
d2 = ray.data.range(10)
|
||||
assert ctx.get_config("foo") == 1
|
||||
assert d1.context.get_config("foo") == 1
|
||||
assert d2.context.get_config("foo") == 1
|
||||
|
||||
# Changing `d1.context` should not affect `d2.context` or the global context.
|
||||
d1.context.set_config("foo", 2)
|
||||
assert d1.context.get_config("foo") == 2
|
||||
assert d2.context.get_config("foo") == 1
|
||||
assert ctx.get_config("foo") == 1
|
||||
|
||||
# Changed value can be propagated to remote tasks.
|
||||
@ray.remote(num_cpus=0)
|
||||
def check(d1, d2):
|
||||
assert d1.context.get_config("foo") == 2
|
||||
assert d2.context.get_config("foo") == 1
|
||||
|
||||
ray.get(check.remote(d1, d2))
|
||||
|
||||
# Changing the global context should not affect `d1.context` or `d2.context`.
|
||||
ctx.set_config("foo", 3)
|
||||
assert ctx.get_config("foo") == 3
|
||||
assert d1.context.get_config("foo") == 2
|
||||
assert d2.context.get_config("foo") == 1
|
||||
|
||||
|
||||
def test_context_inheritance(ray_start_regular_shared):
|
||||
ds = ray.data.range(10)
|
||||
ds.context.set_config("foo", 1)
|
||||
assert DataContext.get_current().get_config("foo", None) is None
|
||||
|
||||
# Test that applying a new operator to an existing dataset
|
||||
# inherits the context.
|
||||
ds2 = ds.map_batches(lambda batch: batch)
|
||||
assert ds2.context.get_config("foo") == 1
|
||||
|
||||
# Test that materializing a dataset also inherits the context.
|
||||
mds = ds.materialize()
|
||||
assert mds.context.get_config("foo") == 1
|
||||
|
||||
# Test that the iterator also inherits the context.
|
||||
iter = ds.iterator()
|
||||
assert iter.get_context().get_config("foo") == 1
|
||||
assert iter.materialize().context.get_config("foo") == 1
|
||||
|
||||
|
||||
def _test_updating_context_after_dataset_creation(gen_ds):
|
||||
context = DataContext.get_current()
|
||||
context.set_config("foo", 1)
|
||||
ds = gen_ds()
|
||||
assert ds.take_all()[0]["id"] == 1
|
||||
# DataContext is supposed to be sealed when a Dataset is created.
|
||||
# Test that updating the current DataContext doesn't affect existing Datasets.
|
||||
context.set_config("foo", 2)
|
||||
assert ds.take_all()[0]["id"] == 1
|
||||
|
||||
|
||||
def test_read(
|
||||
ray_start_regular_shared,
|
||||
):
|
||||
class CustomDatasource(Datasource):
|
||||
def prepare_read(self, parallelism: int):
|
||||
def read_fn():
|
||||
value = DataContext.get_current().get_config("foo")
|
||||
return [pd.DataFrame({"id": [value]})]
|
||||
|
||||
meta = BlockMetadata(
|
||||
num_rows=1, size_bytes=8, input_files=None, exec_stats=None
|
||||
)
|
||||
return [ReadTask(read_fn, meta)]
|
||||
|
||||
_test_updating_context_after_dataset_creation(
|
||||
lambda: ray.data.read_datasource(CustomDatasource()),
|
||||
)
|
||||
|
||||
|
||||
def test_map(
|
||||
ray_start_regular_shared,
|
||||
):
|
||||
_test_updating_context_after_dataset_creation(
|
||||
lambda: ray.data.range(1).map(
|
||||
lambda _: {"id": DataContext.get_current().get_config("foo")}
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_flat_map(
|
||||
ray_start_regular_shared,
|
||||
):
|
||||
_test_updating_context_after_dataset_creation(
|
||||
lambda: ray.data.range(1).flat_map(
|
||||
lambda _: [{"id": DataContext.get_current().get_config("foo")}]
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_map_batches(
|
||||
ray_start_regular_shared,
|
||||
):
|
||||
_test_updating_context_after_dataset_creation(
|
||||
lambda: ray.data.range(1).map_batches(
|
||||
lambda x: {"id": [DataContext.get_current().get_config("foo")]}
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_filter(
|
||||
ray_start_regular_shared,
|
||||
):
|
||||
_test_updating_context_after_dataset_creation(
|
||||
lambda: ray.data.from_items([1])
|
||||
.filter(lambda x: x["item"] == DataContext.get_current().get_config("foo"))
|
||||
.rename_columns({"item": "id"})
|
||||
)
|
||||
|
||||
|
||||
def test_streaming_split(
|
||||
ray_start_regular_shared,
|
||||
):
|
||||
# Tests that custom DataContext can be properly propagated
|
||||
# when using `streaming_split()`.
|
||||
block_size = 123 * 1024 * 1024
|
||||
data_context = DataContext.get_current()
|
||||
data_context.target_max_block_size = block_size
|
||||
data_context.set_config("foo", "bar")
|
||||
|
||||
def f(x):
|
||||
assert DataContext.get_current().target_max_block_size == block_size
|
||||
assert DataContext.get_current().get_config("foo") == "bar"
|
||||
return x
|
||||
|
||||
num_splits = 2
|
||||
splits = (
|
||||
ray.data.range(10, override_num_blocks=10).map(f).streaming_split(num_splits)
|
||||
)
|
||||
|
||||
@ray.remote(num_cpus=0)
|
||||
def consume(split):
|
||||
for _ in split.iter_rows():
|
||||
pass
|
||||
|
||||
assert ray.get([consume.remote(split) for split in splits]) == [None] * num_splits
|
||||
|
||||
|
||||
def test_context_placement_group():
|
||||
driver_code = """
|
||||
import ray
|
||||
from ray.data.context import DataContext
|
||||
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
|
||||
from ray._private.test_utils import placement_group_assert_no_leak
|
||||
|
||||
ray.init(num_cpus=1)
|
||||
|
||||
context = DataContext.get_current()
|
||||
# This placement group will take up all cores of the local cluster.
|
||||
placement_group = ray.util.placement_group(
|
||||
name="core_hog",
|
||||
strategy="SPREAD",
|
||||
bundles=[
|
||||
{"CPU": 1},
|
||||
],
|
||||
)
|
||||
ray.get(placement_group.ready())
|
||||
context.scheduling_strategy = PlacementGroupSchedulingStrategy(placement_group)
|
||||
ds = ray.data.range(100, override_num_blocks=2).map(lambda x: {"id": x["id"] + 1})
|
||||
assert ds.take_all() == [{"id": x} for x in range(1, 101)]
|
||||
placement_group_assert_no_leak([placement_group])
|
||||
ray.shutdown()
|
||||
"""
|
||||
|
||||
# Successful exit is sufficient to verify this test.
|
||||
run_string_as_driver(driver_code)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,293 @@
|
||||
"""Tests for cuDF batch format support in Ray Data.
|
||||
|
||||
These tests require cuDF to be installed and run on GPU CI. Use pytest.importorskip
|
||||
so the file is skipped when cudf is missing (e.g. local CPU runs).
|
||||
|
||||
Uses cudf.testing.assert_eq for comparisons (see cuDF developer guide:
|
||||
https://docs.rapids.ai/api/cudf/latest/developer_guide/testing/).
|
||||
"""
|
||||
|
||||
import pandas as pd
|
||||
import pyarrow as pa
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data._internal.block_batching.block_batching import batch_blocks
|
||||
from ray.data.expressions import col
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
|
||||
cudf = pytest.importorskip("cudf")
|
||||
|
||||
|
||||
def block_generator(num_rows: int, num_blocks: int):
|
||||
"""Yield Arrow blocks for testing batch_blocks."""
|
||||
for i in range(num_blocks):
|
||||
yield pa.table({"foo": list(range(i * num_rows, (i + 1) * num_rows))})
|
||||
|
||||
|
||||
def _make_dataset(data_source: str, shutdown_only, **kwargs):
|
||||
"""Create dataset from data_source type."""
|
||||
if data_source == "range":
|
||||
return ray.data.range(
|
||||
kwargs.get("n", 100), override_num_blocks=kwargs.get("blocks", 2)
|
||||
)
|
||||
elif data_source == "range_tensor":
|
||||
return ray.data.range_tensor(
|
||||
kwargs.get("n", 100), override_num_blocks=kwargs.get("blocks", 2)
|
||||
)
|
||||
elif data_source == "from_pandas":
|
||||
df = kwargs.get(
|
||||
"df",
|
||||
pd.DataFrame({"foo": ["a", "b"], "bar": [0, 1]}),
|
||||
)
|
||||
return ray.data.from_pandas(df)
|
||||
else:
|
||||
raise ValueError(f"Unknown data_source: {data_source}")
|
||||
|
||||
|
||||
# TODO(elliot-barn): "range_tensor" is disabled because cudf does not support Ray's
|
||||
# TensorDtype (TypeError: Unrecognized dtype: TensorDtype). This was not caught before
|
||||
# because cudf was not installed in the docgpu CI image — pytest.importorskip skipped
|
||||
# the entire file. Now that the depset installs cudf-cu12, these tests run for real.
|
||||
@pytest.mark.parametrize(
|
||||
"data_source",
|
||||
["range", "from_pandas"],
|
||||
ids=["range", "from_pandas"],
|
||||
)
|
||||
class TestCudfIterBatches:
|
||||
"""Tests for iter_batches with batch_format='cudf'."""
|
||||
|
||||
def test_iter_batches_returns_cudf(self, shutdown_only, data_source):
|
||||
ds = _make_dataset(data_source, shutdown_only)
|
||||
batch = next(iter(ds.iter_batches(batch_format="cudf")))
|
||||
assert isinstance(batch, cudf.DataFrame)
|
||||
assert len(batch) > 0
|
||||
|
||||
def test_iter_batches_columns(self, shutdown_only, data_source):
|
||||
ds = _make_dataset(data_source, shutdown_only)
|
||||
batch = next(iter(ds.iter_batches(batch_format="cudf")))
|
||||
if data_source == "range":
|
||||
assert list(batch.columns) == ["id"]
|
||||
elif data_source == "range_tensor":
|
||||
assert "data" in batch.columns
|
||||
else:
|
||||
assert list(batch.columns) == ["foo", "bar"]
|
||||
cudf.testing.assert_eq(
|
||||
batch, cudf.DataFrame({"foo": ["a", "b"], "bar": [0, 1]})
|
||||
)
|
||||
|
||||
|
||||
# TODO(elliot-barn): "range_tensor" is disabled — see comment on TestCudfIterBatches.
|
||||
@pytest.mark.parametrize(
|
||||
"data_source",
|
||||
["range"],
|
||||
ids=["range"],
|
||||
)
|
||||
class TestCudfTakeBatch:
|
||||
"""Tests for take_batch with batch_format='cudf'."""
|
||||
|
||||
def test_take_batch_returns_cudf(self, ray_start_regular_shared, data_source):
|
||||
ds = _make_dataset(data_source, None, n=10, blocks=2)
|
||||
batch = ds.take_batch(3, batch_format="cudf")
|
||||
assert isinstance(batch, cudf.DataFrame)
|
||||
|
||||
def test_take_batch_data(self, ray_start_regular_shared, data_source):
|
||||
ds = _make_dataset(data_source, None, n=10, blocks=2)
|
||||
batch = ds.take_batch(3, batch_format="cudf")
|
||||
if data_source == "range":
|
||||
cudf.testing.assert_eq(batch["id"], cudf.Series([0, 1, 2], name="id"))
|
||||
else:
|
||||
# Tensor columns are stored as list-type in cudf; compare via Arrow
|
||||
assert batch["data"].to_arrow().to_pylist() == [[0], [1], [2]]
|
||||
|
||||
|
||||
class TestCudfBatchBlocks:
|
||||
"""Tests for batch_blocks with batch_format='cudf'."""
|
||||
|
||||
def test_batch_blocks_cudf(self):
|
||||
blocks = block_generator(num_rows=3, num_blocks=2)
|
||||
batches = list(batch_blocks(blocks, batch_format="cudf"))
|
||||
|
||||
assert len(batches) == 2
|
||||
assert isinstance(batches[0], cudf.DataFrame)
|
||||
assert isinstance(batches[1], cudf.DataFrame)
|
||||
cudf.testing.assert_eq(batches[0], cudf.DataFrame({"foo": [0, 1, 2]}))
|
||||
cudf.testing.assert_eq(batches[1], cudf.DataFrame({"foo": [3, 4, 5]}))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_format",
|
||||
["cudf", "pandas", "pyarrow"],
|
||||
ids=["cudf", "pandas", "pyarrow"],
|
||||
)
|
||||
class TestCudfMapBatches:
|
||||
"""Tests for map_batches with various batch formats (cuDF in/out)."""
|
||||
|
||||
def test_map_batches_cudf_receive_and_return(
|
||||
self, ray_start_regular_shared, batch_format
|
||||
):
|
||||
"""UDF receives batches in requested format; test cudf round-trip."""
|
||||
ds = ray.data.range(10, override_num_blocks=2)
|
||||
|
||||
def add_one(batch):
|
||||
if batch_format == "cudf":
|
||||
assert isinstance(batch, cudf.DataFrame)
|
||||
batch = batch.copy()
|
||||
batch["id"] = batch["id"] + 1
|
||||
return batch
|
||||
# For pandas/pyarrow input, convert to cudf, transform, return cudf
|
||||
cudf_batch = (
|
||||
cudf.from_pandas(batch)
|
||||
if batch_format == "pandas"
|
||||
else cudf.DataFrame.from_arrow(batch)
|
||||
)
|
||||
cudf_batch["id"] = cudf_batch["id"] + 1
|
||||
return cudf_batch
|
||||
|
||||
result = ds.map_batches(
|
||||
add_one,
|
||||
batch_format=batch_format,
|
||||
batch_size=10,
|
||||
num_gpus=0.001,
|
||||
).take()
|
||||
assert result == [{"id": i} for i in range(1, 11)]
|
||||
|
||||
def test_map_batches_udf_returns_cudf(self, ray_start_regular_shared, batch_format):
|
||||
"""UDF returns cudf.DataFrame regardless of input format (batch_to_block)."""
|
||||
if batch_format == "cudf":
|
||||
pytest.skip("Already testing cudf in/out above")
|
||||
ds = ray.data.range(5, override_num_blocks=1)
|
||||
|
||||
def to_cudf_and_double(batch):
|
||||
cudf_batch = (
|
||||
cudf.from_pandas(batch)
|
||||
if batch_format == "pandas"
|
||||
else cudf.DataFrame.from_arrow(batch)
|
||||
)
|
||||
cudf_batch["id"] = cudf_batch["id"] * 2
|
||||
return cudf_batch
|
||||
|
||||
result = ds.map_batches(
|
||||
to_cudf_and_double,
|
||||
batch_format=batch_format,
|
||||
batch_size=5,
|
||||
num_gpus=0.001,
|
||||
).take()
|
||||
assert result == [{"id": 0}, {"id": 2}, {"id": 4}, {"id": 6}, {"id": 8}]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"predicate_expr, test_data, expected_ids",
|
||||
[
|
||||
(col("id") > 5, None, list(range(6, 10))),
|
||||
(col("id") >= 3, None, list(range(3, 10))),
|
||||
(col("id") < 3, None, [0, 1, 2]),
|
||||
(col("id") <= 2, None, [0, 1, 2]),
|
||||
(col("id") == 4, None, [4]),
|
||||
(col("id") != 4, None, [0, 1, 2, 3, 5, 6, 7, 8, 9]),
|
||||
((col("id") >= 2) & (col("id") < 6), None, [2, 3, 4, 5]),
|
||||
((col("id") < 2) | (col("id") > 7), None, [0, 1, 8, 9]),
|
||||
(
|
||||
col("value").is_not_null(),
|
||||
[{"value": None}, {"value": 1}, {"value": 2}],
|
||||
[1, 2],
|
||||
),
|
||||
],
|
||||
ids=[
|
||||
"gt",
|
||||
"gte",
|
||||
"lt",
|
||||
"lte",
|
||||
"eq",
|
||||
"neq",
|
||||
"and",
|
||||
"or",
|
||||
"is_not_null",
|
||||
],
|
||||
)
|
||||
class TestCudfFilterExpressions:
|
||||
"""Tests for filter with expressions on cuDF blocks."""
|
||||
|
||||
def test_filter_expr_after_map_batches_cudf(
|
||||
self, ray_start_regular_shared, predicate_expr, test_data, expected_ids
|
||||
):
|
||||
"""filter(expr=...) works on cuDF blocks from map_batches(batch_format='cudf')."""
|
||||
if test_data is not None:
|
||||
ds = ray.data.from_items(test_data)
|
||||
ds = ds.map_batches(
|
||||
lambda x: x,
|
||||
batch_format="cudf",
|
||||
batch_size=3,
|
||||
num_gpus=0.001,
|
||||
)
|
||||
result = ds.filter(expr=predicate_expr).take()
|
||||
result_ids = [r.get("value", r.get("id", r)) for r in result]
|
||||
else:
|
||||
ds = ray.data.range(10, override_num_blocks=2)
|
||||
ds = ds.map_batches(
|
||||
lambda x: x,
|
||||
batch_format="cudf",
|
||||
batch_size=10,
|
||||
num_gpus=0.001,
|
||||
)
|
||||
result = ds.filter(expr=predicate_expr).take()
|
||||
result_ids = [r["id"] for r in result]
|
||||
|
||||
assert result_ids == expected_ids
|
||||
|
||||
def test_map_batches_after_filter_expr(
|
||||
self, ray_start_regular_shared, predicate_expr, test_data, expected_ids
|
||||
):
|
||||
"""map_batches(batch_format='cudf') after filter(expr=...) works."""
|
||||
if test_data is not None:
|
||||
ds = ray.data.from_items(test_data)
|
||||
ds = ds.filter(expr=predicate_expr)
|
||||
ds = ds.map_batches(
|
||||
lambda x: x,
|
||||
batch_format="cudf",
|
||||
batch_size=3,
|
||||
num_gpus=0.001,
|
||||
)
|
||||
result = ds.take()
|
||||
result_ids = [r.get("value", r.get("id", r)) for r in result]
|
||||
else:
|
||||
ds = ray.data.range(10, override_num_blocks=2)
|
||||
ds = ds.filter(expr=predicate_expr)
|
||||
ds = ds.map_batches(
|
||||
lambda x: x,
|
||||
batch_format="cudf",
|
||||
batch_size=10,
|
||||
num_gpus=0.001,
|
||||
)
|
||||
result = ds.take()
|
||||
result_ids = [r["id"] for r in result]
|
||||
|
||||
assert result_ids == expected_ids
|
||||
|
||||
|
||||
class TestCudfAddColumn:
|
||||
"""Tests for add_column with batch_format='cudf'."""
|
||||
|
||||
def test_add_column_cudf(self, ray_start_regular_shared):
|
||||
"""add_column with batch_format='cudf' adds column to cudf batches."""
|
||||
ds = ray.data.range(5).add_column(
|
||||
"doubled",
|
||||
lambda x: x["id"] * 2,
|
||||
batch_format="cudf",
|
||||
batch_size=5,
|
||||
num_gpus=0.001,
|
||||
)
|
||||
result = ds.take()
|
||||
assert result == [
|
||||
{"id": 0, "doubled": 0},
|
||||
{"id": 1, "doubled": 2},
|
||||
{"id": 2, "doubled": 4},
|
||||
{"id": 3, "doubled": 6},
|
||||
{"id": 4, "doubled": 8},
|
||||
]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
@@ -0,0 +1,234 @@
|
||||
import math
|
||||
import random
|
||||
import sys
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
import ray
|
||||
from ray.data.tests.conftest import * # noqa
|
||||
from ray.data.tests.conftest import (
|
||||
CoreExecutionMetrics,
|
||||
assert_core_execution_metrics_equals,
|
||||
)
|
||||
from ray.tests.conftest import * # noqa
|
||||
|
||||
|
||||
def test_count(ray_start_regular):
|
||||
ds = ray.data.range(100, override_num_blocks=10)
|
||||
# We do not kick off the read task by default.
|
||||
assert ds.count() == 100
|
||||
# Getting number of rows should not trigger execution of any read tasks
|
||||
# for ray.data.range(), as the number of rows is known beforehand.
|
||||
|
||||
assert_core_execution_metrics_equals(CoreExecutionMetrics(task_count={}))
|
||||
|
||||
|
||||
def test_count_edge_case(ray_start_regular):
|
||||
# Test this edge case: https://github.com/ray-project/ray/issues/44509.
|
||||
ds = ray.data.range(10)
|
||||
ds.count()
|
||||
|
||||
actual_count = ds.filter(fn=lambda row: row["id"] % 2 == 0).count()
|
||||
|
||||
assert actual_count == 5
|
||||
|
||||
|
||||
def test_count_after_caching_after_execution(ray_start_regular):
|
||||
SCALE_FACTOR = 5
|
||||
FILE_ROW_COUNT = 150
|
||||
DS_ROW_COUNT = FILE_ROW_COUNT * SCALE_FACTOR
|
||||
paths = ["example://iris.csv"] * SCALE_FACTOR
|
||||
ds = ray.data.read_csv(paths)
|
||||
# Row count should be unknown before execution.
|
||||
assert "num_rows=?" in str(ds)
|
||||
# After iterating over bundles and completing execution, row count should be known.
|
||||
list(ds.iter_internal_ref_bundles())
|
||||
assert ds.count() == DS_ROW_COUNT
|
||||
assert ds._cache._num_rows == DS_ROW_COUNT
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_parts", [1, 30])
|
||||
@pytest.mark.parametrize("ds_format", ["arrow", "pandas"])
|
||||
def test_global_tabular_min(ray_start_regular_shared_2_cpus, ds_format, num_parts):
|
||||
seed = int(time.time())
|
||||
print(f"Seeding RNG for test_global_arrow_min with: {seed}")
|
||||
random.seed(seed)
|
||||
xs = list(range(100))
|
||||
random.shuffle(xs)
|
||||
|
||||
def _to_pandas(ds):
|
||||
return ds.map_batches(lambda x: x, batch_size=None, batch_format="pandas")
|
||||
|
||||
# Test built-in global min aggregation
|
||||
ds = ray.data.from_items([{"A": x} for x in xs]).repartition(num_parts)
|
||||
if ds_format == "pandas":
|
||||
ds = _to_pandas(ds)
|
||||
assert ds.min("A") == 0
|
||||
|
||||
# Test empty dataset
|
||||
# Note: we explicitly set parallelism here to ensure there are no empty
|
||||
# input blocks.
|
||||
ds = ray.data.range(10, override_num_blocks=10)
|
||||
if ds_format == "pandas":
|
||||
ds = _to_pandas(ds)
|
||||
assert ds.filter(lambda r: r["id"] > 10).min("id") is None
|
||||
|
||||
# Test built-in global min aggregation with nans
|
||||
nan_ds = ray.data.from_items([{"A": x} for x in xs] + [{"A": None}]).repartition(
|
||||
num_parts
|
||||
)
|
||||
if ds_format == "pandas":
|
||||
nan_ds = _to_pandas(nan_ds)
|
||||
assert nan_ds.min("A") == 0
|
||||
# Test ignore_nulls=False
|
||||
assert pd.isnull(nan_ds.min("A", ignore_nulls=False))
|
||||
# Test all nans
|
||||
nan_ds = ray.data.from_items([{"A": None}] * len(xs)).repartition(num_parts)
|
||||
if ds_format == "pandas":
|
||||
nan_ds = _to_pandas(nan_ds)
|
||||
assert pd.isnull(nan_ds.min("A"))
|
||||
assert pd.isnull(nan_ds.min("A", ignore_nulls=False))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_parts", [1, 30])
|
||||
@pytest.mark.parametrize("ds_format", ["arrow", "pandas"])
|
||||
def test_global_tabular_max(ray_start_regular_shared_2_cpus, ds_format, num_parts):
|
||||
seed = int(time.time())
|
||||
print(f"Seeding RNG for test_global_arrow_max with: {seed}")
|
||||
random.seed(seed)
|
||||
xs = list(range(100))
|
||||
random.shuffle(xs)
|
||||
|
||||
def _to_pandas(ds):
|
||||
return ds.map_batches(lambda x: x, batch_size=None, batch_format="pandas")
|
||||
|
||||
# Test built-in global max aggregation
|
||||
ds = ray.data.from_items([{"A": x} for x in xs]).repartition(num_parts)
|
||||
if ds_format == "pandas":
|
||||
ds = _to_pandas(ds)
|
||||
assert ds.max("A") == 99
|
||||
|
||||
# Test empty dataset
|
||||
# Note: we explicitly set parallelism here to ensure there are no empty
|
||||
# input blocks.
|
||||
ds = ray.data.range(10, override_num_blocks=10)
|
||||
if ds_format == "pandas":
|
||||
ds = _to_pandas(ds)
|
||||
assert ds.filter(lambda r: r["id"] > 10).max("id") is None
|
||||
|
||||
# Test built-in global max aggregation with nans
|
||||
nan_ds = ray.data.from_items([{"A": x} for x in xs] + [{"A": None}]).repartition(
|
||||
num_parts
|
||||
)
|
||||
if ds_format == "pandas":
|
||||
nan_ds = _to_pandas(nan_ds)
|
||||
assert nan_ds.max("A") == 99
|
||||
# Test ignore_nulls=False
|
||||
assert pd.isnull(nan_ds.max("A", ignore_nulls=False))
|
||||
# Test all nans
|
||||
nan_ds = ray.data.from_items([{"A": None}] * len(xs)).repartition(num_parts)
|
||||
if ds_format == "pandas":
|
||||
nan_ds = _to_pandas(nan_ds)
|
||||
assert pd.isnull(nan_ds.max("A"))
|
||||
assert pd.isnull(nan_ds.max("A", ignore_nulls=False))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_parts", [1, 30])
|
||||
@pytest.mark.parametrize("ds_format", ["arrow", "pandas"])
|
||||
def test_global_tabular_mean(ray_start_regular_shared_2_cpus, ds_format, num_parts):
|
||||
seed = int(time.time())
|
||||
print(f"Seeding RNG for test_global_arrow_mean with: {seed}")
|
||||
random.seed(seed)
|
||||
xs = list(range(100))
|
||||
random.shuffle(xs)
|
||||
|
||||
def _to_pandas(ds):
|
||||
return ds.map_batches(lambda x: x, batch_size=None, batch_format="pandas")
|
||||
|
||||
# Test built-in global mean aggregation
|
||||
ds = ray.data.from_items([{"A": x} for x in xs]).repartition(num_parts)
|
||||
if ds_format == "pandas":
|
||||
ds = _to_pandas(ds)
|
||||
assert ds.mean("A") == 49.5
|
||||
|
||||
# Test empty dataset
|
||||
# Note: we explicitly set parallelism here to ensure there are no empty
|
||||
# input blocks.
|
||||
ds = ray.data.range(10, override_num_blocks=10)
|
||||
if ds_format == "pandas":
|
||||
ds = _to_pandas(ds)
|
||||
assert ds.filter(lambda r: r["id"] > 10).mean("id") is None
|
||||
|
||||
# Test built-in global mean aggregation with nans
|
||||
nan_ds = ray.data.from_items([{"A": x} for x in xs] + [{"A": None}]).repartition(
|
||||
num_parts
|
||||
)
|
||||
if ds_format == "pandas":
|
||||
nan_ds = _to_pandas(nan_ds)
|
||||
assert nan_ds.mean("A") == 49.5
|
||||
# Test ignore_nulls=False
|
||||
assert pd.isnull(nan_ds.mean("A", ignore_nulls=False))
|
||||
# Test all nans
|
||||
nan_ds = ray.data.from_items([{"A": None}] * len(xs)).repartition(num_parts)
|
||||
if ds_format == "pandas":
|
||||
nan_ds = _to_pandas(nan_ds)
|
||||
assert pd.isnull(nan_ds.mean("A"))
|
||||
assert pd.isnull(nan_ds.mean("A", ignore_nulls=False))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_parts", [1, 30])
|
||||
@pytest.mark.parametrize("ds_format", ["arrow", "pandas"])
|
||||
def test_global_tabular_std(ray_start_regular_shared_2_cpus, ds_format, num_parts):
|
||||
# NOTE: Do not change the seed
|
||||
seed = 1740035705
|
||||
|
||||
random.seed(seed)
|
||||
xs = list(range(100))
|
||||
random.shuffle(xs)
|
||||
|
||||
def _to_arrow(ds):
|
||||
return ds.map_batches(lambda x: x, batch_size=None, batch_format="pyarrow")
|
||||
|
||||
def _to_pandas(ds):
|
||||
return ds.map_batches(lambda x: x, batch_size=None, batch_format="pandas")
|
||||
|
||||
# Test built-in global max aggregation
|
||||
df = pd.DataFrame({"A": xs})
|
||||
ds = ray.data.from_pandas(df).repartition(num_parts)
|
||||
if ds_format == "arrow":
|
||||
ds = _to_arrow(ds)
|
||||
assert math.isclose(ds.std("A"), df["A"].std())
|
||||
assert math.isclose(ds.std("A", ddof=0), df["A"].std(ddof=0))
|
||||
|
||||
# Test empty dataset
|
||||
ds = ray.data.from_pandas(pd.DataFrame({"A": []}))
|
||||
if ds_format == "arrow":
|
||||
ds = _to_arrow(ds)
|
||||
assert pd.isnull(ds.std("A"))
|
||||
# Test edge cases
|
||||
ds = ray.data.from_pandas(pd.DataFrame({"A": [3]}))
|
||||
if ds_format == "arrow":
|
||||
ds = _to_arrow(ds)
|
||||
assert np.isnan(ds.std("A"))
|
||||
|
||||
# Test built-in global std aggregation with nans
|
||||
nan_df = pd.DataFrame({"A": xs + [None]})
|
||||
nan_ds = ray.data.from_pandas(nan_df).repartition(num_parts)
|
||||
if ds_format == "arrow":
|
||||
nan_ds = _to_arrow(nan_ds)
|
||||
assert math.isclose(nan_ds.std("A"), nan_df["A"].std())
|
||||
# Test ignore_nulls=False
|
||||
assert pd.isnull(nan_ds.std("A", ignore_nulls=False))
|
||||
# Test all nans
|
||||
nan_ds = ray.data.from_items([{"A": None}] * len(xs)).repartition(num_parts)
|
||||
if ds_format == "pandas":
|
||||
nan_ds = _to_pandas(nan_ds)
|
||||
assert pd.isnull(nan_ds.std("A"))
|
||||
assert pd.isnull(nan_ds.std("A", ignore_nulls=False))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main(["-v", __file__]))
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user