[Data] Simplify and remove the ordering dependency of download expression error handling tests (ray-project#58518)

## Description

This PR refactors tests in `test_download_expression.py` to make them
easier to maintain and less prone to brittle failures. Some of the
previous tests were more complex than necessary and relied on
assumptions that could occasionally cause false negatives.

### Key updates:
* **Reduce flaky behavior**: Added explicit sorting by ID in
`test_download_expression_handles_failed_downloads` to avoid relying on
a specific output order, which isn’t guaranteed and could sometimes
cause intermittent failures.
* **Simplify test logic**: Reduced
`test_download_expression_failed_size_estimation` from 30 URIs to just
1. A single failing URI is sufficient to confirm that failed downloads
don’t trigger divide-by-zero errors, and this change makes the test
easier to understand and faster to run.
* **Improve readability**: Replaced `pa.Table.from_arrays()` with
`ray.data.from_items()`, which makes the test setup more straightforward
for future maintainers.
* **Remove redundancy**: Deleted
`test_download_expression_mixed_valid_and_invalid_size_estimation`,
since its behavior is already covered by the other tests.

Overall, these updates streamline the test suite, making it faster,
clearer, and more robust while keeping the key behaviors fully verified.

## Related issue

ray-project#58464 (comment)

---------

Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu>
Signed-off-by: Xinyu Zhang <60529799+xyuzh@users.noreply.github.com>
Signed-off-by: Robert Nishihara <robertnishihara@gmail.com>
Co-authored-by: Xinyu Zhang <60529799+xyuzh@users.noreply.github.com>
Co-authored-by: Robert Nishihara <robertnishihara@gmail.com>
Signed-off-by: justinyeh1995 <justinyeh1995@gmail.com>

Author: 3 people

python/ray/data/tests/test_download_expression.py

@@ -299,6 +299,9 @@ def test_download_expression_with_malformed_uris(self, tmp_path):
 
         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
@@ -324,78 +327,30 @@ def test_download_expression_with_malformed_uris(self, tmp_path):
         for result in results:
             assert result["bytes"] is None
 
-    def test_download_expression_all_size_estimations_fail(self):
-        """Test download expression when all URI size estimations fail.
-
-        This tests the failed download does not cause division by zero error.
-        """
-        # Create URIs that will fail size estimation (non-existent files)
-        # Using enough URIs to trigger size estimation sampling
-        invalid_uris = [
-            f"local:///nonexistent/path/file_{i}.txt"
-            for i in range(30)  # More than INIT_SAMPLE_BATCH_SIZE (25)
-        ]
+    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")
 
-        table = pa.Table.from_arrays(
-            [pa.array(invalid_uris)],
-            names=["uri"],
+        # 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 = ray.data.from_arrow(table)
         ds_with_downloads = ds.with_column("bytes", download("uri"))
 
-        # Should not crash with divide-by-zero error
-        # The PartitionActor should handle all failed size estimations gracefully
-        # and fall back to using the number of rows in the block as partition size
-        results = ds_with_downloads.take_all()
-
-        # All downloads should fail gracefully (return None)
-        assert len(results) == 30
-        for result in results:
-            assert result["bytes"] is None
-
-    def test_download_expression_mixed_valid_and_invalid_size_estimation(
-        self, tmp_path
-    ):
-        """Test download expression with mix of valid and invalid URIs for size estimation.
-
-        This tests that size estimation handles partial failures correctly.
-        """
-        # Create some valid files
-        valid_files = []
-        for i in range(10):
-            file_path = tmp_path / f"valid_{i}.txt"
-            file_path.write_bytes(b"x" * 100)  # 100 bytes each
-            valid_files.append(str(file_path))
-
-        # Mix valid and invalid URIs
-        mixed_uris = []
-        for i in range(30):
-            if i % 3 == 0 and i // 3 < len(valid_files):
-                # Every 3rd URI is valid (for first 10)
-                mixed_uris.append(f"local://{valid_files[i // 3]}")
-            else:
-                # Others are invalid
-                mixed_uris.append(f"local:///nonexistent/file_{i}.txt")
+        # Should not crash - failed downloads return None
+        results = sorted(ds_with_downloads.take_all(), key=lambda row: row["id"])
+        assert len(results) == 2
 
-        table = pa.Table.from_arrays(
-            [pa.array(mixed_uris)],
-            names=["uri"],
-        )
+        # First URI should succeed
+        assert results[0]["bytes"] == b"valid content"
 
-        ds = ray.data.from_arrow(table)
-        ds_with_downloads = ds.with_column("bytes", download("uri"))
-
-        # Should not crash - should handle mixed valid/invalid gracefully
-        results = ds_with_downloads.take_all()
-        assert len(results) == 30
-
-        # Verify valid URIs downloaded successfully
-        for i, result in enumerate(results):
-            if i % 3 == 0 and i // 3 < len(valid_files):
-                assert result["bytes"] == b"x" * 100
-            else:
-                assert result["bytes"] is None
+        # Second URI should fail gracefully (return None)
+        assert results[1]["bytes"] is None
 
 
 class TestDownloadExpressionIntegration: