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

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

        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(
            [

Review comment (refactor):
This test assumes that the dataset outputs row in a specific order, but that ordering isn't guaranteed by the Ray Data interface. Historically, this has been a major source of flakiness for Ray Data's test suite.

If you've updated the test to use 1 file, this shouldn't be an issue; alternatively, consider using a utility function like `rows_same`

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