wehub-resource-sync
431 lines
16 KiB
Python
431 lines
16 KiB
Python
import io
|
|
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
import ray
|
|
from ray.data.expressions import DownloadExpr, col, download
|
|
|
|
|
|
class TestDownloadExpressionStructure:
|
|
"""Test DownloadExpr structural equality and basic properties."""
|
|
|
|
def test_download_expression_creation(self):
|
|
"""Test that download() creates a DownloadExpr with correct properties."""
|
|
expr = download("uri_column")
|
|
|
|
assert isinstance(expr, DownloadExpr)
|
|
assert expr.uri_column_name == "uri_column"
|
|
|
|
def test_download_expression_structural_equality(self):
|
|
"""Test structural equality comparison for download expressions."""
|
|
# Same expressions should be equal
|
|
expr1 = download("uri")
|
|
expr2 = download("uri")
|
|
assert expr1.structurally_equals(expr2)
|
|
assert expr2.structurally_equals(expr1)
|
|
|
|
# Different URI column names should not be equal
|
|
expr3 = download("different_uri")
|
|
assert not expr1.structurally_equals(expr3)
|
|
assert not expr3.structurally_equals(expr1)
|
|
|
|
# Compare with non-DownloadExpr
|
|
non_download_expr = col("uri")
|
|
assert not expr1.structurally_equals(non_download_expr)
|
|
assert not non_download_expr.structurally_equals(expr1)
|
|
|
|
|
|
class TestDownloadExpressionFunctionality:
|
|
"""Test actual download functionality with real and mocked data."""
|
|
|
|
def test_download_expression_with_local_files(self, tmp_path):
|
|
"""Test basic download expression functionality with local files."""
|
|
# Create sample files with different content types
|
|
sample_data = [
|
|
b"This is test file 1 content",
|
|
b"Different content for file 2",
|
|
b"File 3 has some binary data: \x00\x01\x02\x03",
|
|
]
|
|
|
|
file_paths = []
|
|
for i, data in enumerate(sample_data):
|
|
file_path = tmp_path / f"test_file_{i}.txt"
|
|
file_path.write_bytes(data)
|
|
file_paths.append(str(file_path))
|
|
|
|
# Create dataset with file URIs and metadata
|
|
table = pa.Table.from_arrays(
|
|
[
|
|
pa.array([f"local://{path}" for path in file_paths]),
|
|
pa.array([f"id_{i}" for i in range(len(file_paths))]),
|
|
pa.array([f"metadata_{i}" for i in range(len(file_paths))]),
|
|
pa.array(range(len(file_paths))),
|
|
],
|
|
names=["file_uri", "file_id", "metadata", "index"],
|
|
)
|
|
|
|
ds = ray.data.from_arrow(table)
|
|
|
|
# Add download column using expression
|
|
ds_with_downloads = ds.with_column("file_bytes", download("file_uri"))
|
|
|
|
# Verify results
|
|
results = ds_with_downloads.take_all()
|
|
assert len(results) == len(sample_data)
|
|
|
|
for i, result in enumerate(results):
|
|
# Download column should be added correctly
|
|
assert "file_bytes" in result
|
|
assert result["file_bytes"] == sample_data[i]
|
|
|
|
# All original columns should be preserved
|
|
assert result["file_id"] == f"id_{i}"
|
|
assert result["metadata"] == f"metadata_{i}"
|
|
assert result["index"] == i
|
|
assert result["file_uri"] == f"local://{file_paths[i]}"
|
|
|
|
def test_download_expression_empty_dataset(self):
|
|
"""Test download expression with empty dataset."""
|
|
# Create empty dataset with correct schema
|
|
table = pa.Table.from_arrays(
|
|
[
|
|
pa.array([], type=pa.string()),
|
|
],
|
|
names=["uri"],
|
|
)
|
|
|
|
ds = ray.data.from_arrow(table)
|
|
ds_with_downloads = ds.with_column("bytes", download("uri"))
|
|
|
|
results = ds_with_downloads.take_all()
|
|
assert len(results) == 0
|
|
|
|
def test_download_expression_with_different_file_types(self, tmp_path):
|
|
"""Test download expression with various file types including actual images."""
|
|
# Create a small 8x8 RGB image
|
|
small_image = Image.new("RGB", (8, 8), color=(255, 0, 0)) # Red 8x8 image
|
|
image_buffer = io.BytesIO()
|
|
small_image.save(image_buffer, format="PNG")
|
|
image_bytes = image_buffer.getvalue()
|
|
|
|
# Create files with different types of content
|
|
test_files = [
|
|
("text_file.txt", b"Simple text content"),
|
|
("binary_file.dat", b"\x00\x01\x02\x03\x04\x05"),
|
|
("json_file.json", b'{"key": "value", "number": 123}'),
|
|
("small_image.png", image_bytes), # Actual PNG image (primary use case)
|
|
("empty_file.txt", b""), # Empty file edge case
|
|
]
|
|
|
|
file_paths = []
|
|
expected_data = []
|
|
for filename, content in test_files:
|
|
file_path = tmp_path / filename
|
|
file_path.write_bytes(content)
|
|
file_paths.append(str(file_path))
|
|
expected_data.append(content)
|
|
|
|
# Create dataset
|
|
table = pa.Table.from_arrays(
|
|
[
|
|
pa.array([f"local://{path}" for path in file_paths]),
|
|
pa.array(
|
|
[f.split(".")[0] for f, _ in test_files]
|
|
), # filename without extension
|
|
],
|
|
names=["file_uri", "file_type"],
|
|
)
|
|
|
|
ds = ray.data.from_arrow(table)
|
|
ds_with_downloads = ds.with_column("content", download("file_uri"))
|
|
|
|
results = ds_with_downloads.take_all()
|
|
assert len(results) == len(test_files)
|
|
|
|
for i, result in enumerate(results):
|
|
assert result["content"] == expected_data[i]
|
|
assert result["file_type"] == test_files[i][0].split(".")[0]
|
|
|
|
# Special verification for image file - ensure it can be loaded as an image
|
|
if test_files[i][0].endswith(".png"):
|
|
downloaded_image = Image.open(io.BytesIO(result["content"]))
|
|
assert downloaded_image.size == (8, 8)
|
|
assert downloaded_image.mode == "RGB"
|
|
|
|
def test_chained_download_expressions(self, tmp_path):
|
|
"""Test chained download expressions functionality."""
|
|
# Create sample files with different content
|
|
sample_data = [
|
|
b"Content for file 1",
|
|
b"Content for file 2",
|
|
b"Content for file 3",
|
|
]
|
|
|
|
file_paths = []
|
|
for i, data in enumerate(sample_data):
|
|
file_path = tmp_path / f"test_file_{i}.txt"
|
|
file_path.write_bytes(data)
|
|
file_paths.append(str(file_path))
|
|
|
|
# Create dataset with file URIs
|
|
table = pa.Table.from_arrays(
|
|
[
|
|
pa.array([f"local://{path}" for path in file_paths]),
|
|
pa.array([f"id_{i}" for i in range(len(file_paths))]),
|
|
],
|
|
names=["file_uri", "file_id"],
|
|
)
|
|
|
|
ds = ray.data.from_arrow(table)
|
|
|
|
# Chain multiple download expressions from the same URI column
|
|
ds_with_chained_downloads = (
|
|
ds.with_column("file_bytes_1", download("file_uri"))
|
|
.with_column("file_bytes_2", download("file_uri"))
|
|
.with_column("file_bytes_3", download("file_uri"))
|
|
)
|
|
|
|
# Verify results
|
|
results = ds_with_chained_downloads.take_all()
|
|
assert len(results) == len(sample_data)
|
|
|
|
for i, result in enumerate(results):
|
|
# All download columns should have the same content
|
|
assert "file_bytes_1" in result
|
|
assert "file_bytes_2" in result
|
|
assert "file_bytes_3" in result
|
|
assert result["file_bytes_1"] == sample_data[i]
|
|
assert result["file_bytes_2"] == sample_data[i]
|
|
assert result["file_bytes_3"] == sample_data[i]
|
|
|
|
# Original columns should be preserved
|
|
assert result["file_id"] == f"id_{i}"
|
|
assert result["file_uri"] == f"local://{file_paths[i]}"
|
|
|
|
def test_download_expression_with_pandas_blocks(self, tmp_path):
|
|
"""Test download with pandas blocks to ensure arrow conversion works.
|
|
|
|
This tests the code path in PartitionActor / AsyncPartitionActor.__call__ where non-arrow
|
|
blocks are converted to arrow format before processing.
|
|
"""
|
|
ctx = ray.data.context.DataContext.get_current()
|
|
old_enable_pandas_block = ctx.enable_pandas_block
|
|
ctx.enable_pandas_block = True
|
|
try:
|
|
# Create test files
|
|
sample_data = [
|
|
b"Pandas block test content 1",
|
|
b"Pandas block test content 2",
|
|
]
|
|
|
|
file_paths = []
|
|
for i, data in enumerate(sample_data):
|
|
file_path = tmp_path / f"pandas_test_{i}.txt"
|
|
file_path.write_bytes(data)
|
|
file_paths.append(str(file_path))
|
|
|
|
# Create dataset with pandas blocks (not arrow)
|
|
df = pd.DataFrame(
|
|
{
|
|
"file_uri": [f"local://{path}" for path in file_paths],
|
|
"file_id": [f"id_{i}" for i in range(len(file_paths))],
|
|
}
|
|
)
|
|
ds = ray.data.from_pandas(df)
|
|
|
|
# Apply download - this should trigger arrow conversion in the partition actor
|
|
ds_with_downloads = ds.with_column("content", download("file_uri"))
|
|
|
|
# Verify results
|
|
results = ds_with_downloads.take_all()
|
|
assert len(results) == len(sample_data)
|
|
|
|
for i, result in enumerate(results):
|
|
assert result["content"] == sample_data[i]
|
|
assert result["file_id"] == f"id_{i}"
|
|
assert result["file_uri"] == f"local://{file_paths[i]}"
|
|
finally:
|
|
ctx.enable_pandas_block = old_enable_pandas_block
|
|
|
|
def test_download_expression_with_custom_filesystem(self, tmp_path):
|
|
import pyarrow.fs as pafs
|
|
|
|
# 1. Setup paths
|
|
subdir = tmp_path / "data"
|
|
subdir.mkdir()
|
|
|
|
file_name = "test_file.txt"
|
|
file_path = subdir / file_name
|
|
sample_content = b"File content with custom fs"
|
|
file_path.write_bytes(sample_content)
|
|
|
|
# 2. Setup SubTreeFileSystem
|
|
# This treats 'subdir' as the root '/'
|
|
base_fs = pafs.LocalFileSystem()
|
|
custom_fs = pafs.SubTreeFileSystem(str(subdir), base_fs)
|
|
|
|
# 3. Create Dataset
|
|
# Note: We use the relative 'file_name' because the FS is rooted at 'subdir'
|
|
ds = ray.data.from_items([{"file_uri": file_name, "file_id": 0}])
|
|
|
|
# 4. Execute Download
|
|
ds_with_downloads = ds.with_column(
|
|
"content", download("file_uri", filesystem=custom_fs)
|
|
)
|
|
|
|
# 5. Assertions
|
|
results = ds_with_downloads.take_all()
|
|
|
|
assert len(results) == 1
|
|
assert results[0]["content"] == sample_content
|
|
assert results[0]["file_id"] == 0
|
|
|
|
|
|
class TestDownloadExpressionErrors:
|
|
"""Test error conditions and edge cases for download expressions."""
|
|
|
|
def test_download_expression_invalid_uri_column(self):
|
|
"""Test download expression with non-existent URI column."""
|
|
table = pa.Table.from_arrays(
|
|
[
|
|
pa.array(["local://test.txt"]),
|
|
],
|
|
names=["existing_column"],
|
|
)
|
|
|
|
ds = ray.data.from_arrow(table)
|
|
ds_with_downloads = ds.with_column("bytes", download("non_existent_column"))
|
|
|
|
# Should raise error when trying to execute
|
|
with pytest.raises(ValueError):
|
|
ds_with_downloads.take_all()
|
|
|
|
def test_download_expression_with_null_uris(self):
|
|
"""Test download expression handling of null/empty URIs."""
|
|
table = pa.Table.from_arrays(
|
|
[
|
|
pa.array(["local://test.txt", None, ""]),
|
|
],
|
|
names=["uri"],
|
|
)
|
|
|
|
ds = ray.data.from_arrow(table)
|
|
ds_with_downloads = ds.with_column("bytes", download("uri"))
|
|
|
|
# Should handle nulls gracefully (exact behavior may vary)
|
|
# This test mainly ensures no crash occurs
|
|
try:
|
|
results = ds_with_downloads.take_all()
|
|
# If it succeeds, verify structure is reasonable
|
|
assert len(results) == 3
|
|
for result in results:
|
|
assert "bytes" in result
|
|
except Exception as e:
|
|
# If it fails, should be a reasonable error (not a crash)
|
|
assert isinstance(e, (ValueError, KeyError, RuntimeError))
|
|
|
|
def test_download_expression_with_malformed_uris(self, tmp_path):
|
|
"""Test download expression with malformed URIs.
|
|
|
|
This tests that various malformed URIs are caught and return None
|
|
instead of crashing.
|
|
|
|
All of the URIs should be malformed in order to test the ZeroDivisionError
|
|
described in https://github.com/ray-project/ray/issues/58462.
|
|
"""
|
|
malformed_uris = [
|
|
f"local://{tmp_path}/nonexistent.txt", # File doesn't exist
|
|
"local:///this/path/does/not/exist/file.txt", # Invalid path
|
|
"", # Empty URI
|
|
"foobar", # Random string
|
|
# TODO(xyuzh): Currently, using the below URIs raises an exception
|
|
# in _resolve_paths_and_filesystem. We need to fix that issue and
|
|
# add the tests in.
|
|
# "file:///\x00/null/byte", # Null byte
|
|
# "http://host/path\n\r", # Line breaks
|
|
# "foo://bar", # Invalid scheme
|
|
# "://no-scheme", # Missing scheme
|
|
# "http://host/path?query=<script>", # Injection attempts
|
|
]
|
|
|
|
ds = ray.data.from_items([{"uri": uri} for uri in malformed_uris])
|
|
ds_with_downloads = ds.with_column("bytes", download("uri"))
|
|
results = ds_with_downloads.take_all()
|
|
|
|
# All malformed URIs should return None
|
|
assert len(results) == len(malformed_uris)
|
|
for result in results:
|
|
assert result["bytes"] is None
|
|
|
|
def test_download_expression_mixed_valid_and_invalid_uris(self, tmp_path):
|
|
"""Test download expression when some but not all of the URIs are invalid."""
|
|
# Create one valid file
|
|
valid_file = tmp_path / "valid.txt"
|
|
valid_file.write_bytes(b"valid content")
|
|
|
|
# Create URIs: one valid and one non-existent file.
|
|
ds = ray.data.from_items(
|
|
[
|
|
{"uri": str(valid_file), "id": 0},
|
|
{"uri": str(tmp_path / "nonexistent.txt"), "id": 1},
|
|
]
|
|
)
|
|
ds_with_downloads = ds.with_column("bytes", download("uri"))
|
|
|
|
# Should not crash - failed downloads return None
|
|
results = sorted(ds_with_downloads.take_all(), key=lambda row: row["id"])
|
|
assert len(results) == 2
|
|
|
|
# First URI should succeed
|
|
assert results[0]["bytes"] == b"valid content"
|
|
|
|
# Second URI should fail gracefully (return None)
|
|
assert results[1]["bytes"] is None
|
|
|
|
|
|
class TestDownloadExpressionIntegration:
|
|
"""Integration tests combining download expressions with other Ray Data operations."""
|
|
|
|
def test_download_expression_with_map_batches(self, tmpdir):
|
|
"""Test download expression followed by map_batches processing."""
|
|
# Create a test file
|
|
test_file = tmpdir.join("test.txt")
|
|
test_content = b"Hello, World!"
|
|
test_file.write_binary(test_content)
|
|
|
|
# Create dataset
|
|
table = pa.Table.from_arrays(
|
|
[
|
|
pa.array([f"local://{test_file}"]),
|
|
],
|
|
names=["uri"],
|
|
)
|
|
|
|
ds = ray.data.from_arrow(table)
|
|
|
|
# Download then process
|
|
ds_with_content = ds.with_column("raw_bytes", download("uri"))
|
|
|
|
def decode_bytes(batch):
|
|
# Access the specific column containing the bytes data
|
|
batch["decoded_text"] = [
|
|
data.decode("utf-8") for data in batch["raw_bytes"]
|
|
]
|
|
return batch
|
|
|
|
ds_decoded = ds_with_content.map_batches(decode_bytes)
|
|
results = ds_decoded.take_all()
|
|
|
|
assert len(results) == 1
|
|
assert results[0]["decoded_text"] == "Hello, World!"
|
|
assert results[0]["raw_bytes"] == test_content
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import sys
|
|
|
|
sys.exit(pytest.main(["-v", __file__]))
|