[Data] Fix file size ordering in download partitioning with multiple URI columns (#58517)

The `_sample_sizes` method was using `as_completed()` to collect file
sizes, which returns results in completion order rather than submission
order. This scrambled the file sizes list so it no longer corresponded
to the input URI order.

When multiple URI columns are used, `_estimate_nrows_per_partition`
calls `zip(*sampled_file_sizes_by_column.values())` on line 284, which
assumes file sizes from different columns align by row index. The
scrambled ordering caused file sizes from different rows to be
incorrectly combined, producing incorrect partition size estimates.

## Changes

- Pre-allocate the `file_sizes` list with the correct size
- Use a `future_to_file_index` mapping to track the original submission
order
- Place results at their correct positions regardless of completion
order
- Add assertion to verify list length matches expected size

## Related issues

#58464 (comment)

---------

Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Author: bveeramani

python/ray/data/_internal/planner/plan_download_op.py

@@ -326,20 +326,25 @@ def get_file_size(uri_path, fs):
         )
 
         # Use ThreadPoolExecutor for concurrent size fetching
-        file_sizes = []
+        file_sizes = [None] * len(paths)
         with ThreadPoolExecutor(max_workers=URI_DOWNLOAD_MAX_WORKERS) as executor:
             # Submit all size fetch tasks
-            futures = [
-                executor.submit(get_file_size, uri_path, fs) for uri_path in paths
-            ]
+            future_to_file_index = {
+                executor.submit(get_file_size, uri_path, fs): file_index
+                for file_index, uri_path in enumerate(paths)
+            }
 
             # Collect results as they complete (order doesn't matter)
-            for future in as_completed(futures):
+            for future in as_completed(future_to_file_index):
+                file_index = future_to_file_index[future]
                 try:
                     size = future.result()
-                    file_sizes.append(size if size is not None else 0)
+                    file_sizes[file_index] = size if size is not None else 0
                 except Exception as e:
                     logger.warning(f"Error fetching file size for download: {e}")
-                    file_sizes.append(0)
+                    file_sizes[file_index] = 0
 
+        assert all(
+            fs is not None for fs in file_sizes
+        ), "File size sampling did not complete for all paths"
         return file_sizes