[data] Fix high memory usage with FileBasedDatasource & ParquetDatasource when using a large number of files

[JackGammack]

Why are these changes needed?

Using FileBasedDatasource or ParquetDatasource with a very large number of files causes OOM when creating read tasks. The full list of file paths is stored in self, causing it to persist to every read task, leading to this warning:

The serialized size of your read function named 'read_task_fn' is 49.8MB. This size relatively large. As a result, Ray might excessively spill objects during execution. To fix this issue, avoid accessing `self` or other large objects in 'read_task_fn'.

When using a small number of blocks, OOM does not occur because the large file list is not repeated so many times. But when setting high parallelism with override_num_blocks, OOM occurs.

This is because the full list of paths is added to self._unresolved_paths. This attribute isn't currently used anywhere in ray. This PR removes self._unresolved_paths to alleviate unexpected high memory usage with very large numbers of files.

Related issue number

Similar to this issue with Iceberg: #49054

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[gvspraveen]

Thank you for the contribution. And good catch on high mem for large number of blocks.

On closer look, _unresolved_paths appears to be not used anywhere. @bveeramani @omatthew98 Do we still need it or should we delete it?
cc @JDarDagran who may have more thoughts on this (this was introduced in #50332).

[omatthew98]

Do we still need it or should we delete it?

Yeah it is not used anywhere and I don't see any reason to keep it. The only argument for keeping it would be to ensure that we don't break any assumptions for downstream consumers (e.g. users who extend FileBasedDatasource or ParquetDatasource), but I doubt there are any users doing that and if they are they aren't following a public API to do so.

@gvspraveen

[JackGammack]

Yes, I'm happy to update this to just remove it. The unresolved paths aren't used anywhere, and aren't accessible by users using the read_parquet etc API since the datasource is created internally and not passed outside of that function. The PRs that added the unresolved paths were relatively recent, so I thought they might want to be kept.

[JDarDagran]

@gvspraveen thanks for pinging me. I'm using this for lineage tracking purpose which was not yet open-sourced. The relation is not that clear - I'm looking at the object through ray.data._internal.execution.execution_callback.ExecutionCallback.

Maybe those unresolved_paths could be put into GCS, similarly to resolved paths?

        # Read tasks serialize `FileBasedDatasource` instances, and the list of paths
        # can be large. To avoid slow serialization speeds, we store a reference to
        # the paths rather than the paths themselves.
        self._paths_ref = ray.put(paths)

[JackGammack]

I think keeping it in the object store is okay, the size of the many file paths isn't so bad if it's just stored once. I only run into issues when I increase the number of blocks over 2,000 or so

[JDarDagran]

I think keeping it in the object store is okay, the size of the many file paths isn't so bad if it's just stored once. I only run into issues when I increase the number of blocks over 2,000 or so

Thanks! I'd really love to avoid losing that metadata.

[omatthew98]

Thanks @JDarDagran for the input, thanks @JackGammack for the change and iteration. LGTM.

[JackGammack]

A couple things to note after updating:

1. I didn't add the helper to ray.get the actual unresolved paths as is done here, but could if you wanted. The test case shows getting the unresolved paths from the object store
2. I noticed that the ParquetDatasource doesn't store the resolved paths in the object store the way FileBasedDatasource does, but instead in self._pq_paths. So it could also have memory issues for large numbers of files and blocks. But the ParquetDatasourcealready has issues with large numbers of files because of the metadata collection

[alexeykudinkin]

Thank you for your contribution @JackGammack!

[alexeykudinkin]

Let's actually just delete this field

[JackGammack]

Are you saying to remove self._unresolved_paths_ref for ParquetDatasource, but keep it for FileBasedDatasource? Or is this an older comment

[alexeykudinkin]

I see no good reason for us not to nuke it in both places

[JackGammack]

I do agree if it's not used anywhere currently especially with the big performance issue, but I'm not aware of the other ongoing work. It should be very easy to add back.

I have updated to just remove self._unresolved_paths now