Repository: ray-project/ray
File: python/ray/data/tests/test_download_expression.py (Python)

Code before:
        ]

        # Create files in the subdirectory
        file_names = []
        for i, data in enumerate(sample_data):
            file_path = subdir / f"file_{i}.txt"
            file_path.write_bytes(data)
            file_names.append(f"file_{i}.txt")

        # Create a SubTreeFileSystem rooted at the subdirectory
        # This means we can use relative paths without the full path
        base_fs = pafs.LocalFileSystem()
        custom_fs = pafs.SubTreeFileSystem(str(subdir), base_fs)

        # Create dataset with just filenames (not full paths)
        # This only works because our custom filesystem is rooted at subdir
        table = pa.Table.from_arrays(
            [
                pa.array(file_names),  # Just filenames, not full paths
                pa.array([f"id_{i}" for i in range(len(file_names))]),
            ],
            names=["file_uri", "file_id"],
        )

        ds = ray.data.from_arrow(table)

        # Download with custom filesystem - this should work because
        # the SubTreeFileSystem resolves relative paths from subdir
        ds_with_downloads = ds.with_column(
            "content", download("file_uri", filesystem=custom_fs)
        )

        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}"


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"],
        )

Review comment (suggestion):
Consider using `from_items` instead. It's easier to read

Code after:
        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)