[Data] Update document embedding benchmark to use canonical Ray Data API

[bveeramani]

Summary

This PR updates the document embedding benchmark to use the canonical Ray Data implementation pattern, following best practices for the framework.

Key Changes

Use download() expression instead of separate materialization

Before:

file_paths = (
    ray.data.read_parquet(INPUT_PATH)
    .filter(lambda row: row["file_name"].endswith(".pdf"))
    .take_all()
)
file_paths = [row["uploaded_pdf_path"] for row in file_paths]
ds = ray.data.read_binary_files(file_paths, include_paths=True)

After:

(
    ray.data.read_parquet(INPUT_PATH)
    .filter(lambda row: row["file_name"].endswith(".pdf"))
    .with_column("bytes", download("uploaded_pdf_path"))

This change:

• Eliminates the intermediate materialization with take_all(), which loads all data into memory
• Uses the download() expression to lazily fetch file contents as part of the pipeline
• Removes the need for a separate read_binary_files() call

Method chaining for cleaner code

All operations are now chained in a single pipeline, making the data flow more clear and idiomatic.

Consistent column naming

Updated references from path to uploaded_pdf_path throughout the code for consistency with the source data schema.

[gemini-code-assist]

Code Review

This pull request is a great improvement, refactoring the document embedding benchmark to use the idiomatic Ray Data API. The switch from take_all() and read_binary_files to the lazy download() expression is a significant enhancement for performance and memory efficiency, especially for large datasets. The code is now much cleaner and easier to follow with the single chained pipeline.

I have a suggestion to further improve the code's adherence to Ray Data best practices by using a col expression for filtering instead of a lambda function. This can make the code more declarative and potentially unlock further optimizations.

[gemini-code-assist]

To make the code even more idiomatic and potentially more performant through expression optimization, consider importing col here. It can be used to replace the lambda in the filter operation below.

Suggested change

	from ray.data.expressions import download
	from ray.data.expressions import col, download

[gemini-code-assist]

Using a col expression is more idiomatic in Ray Data than using a lambda for simple filtering operations. This declarative style can also allow Ray Data to perform more optimizations on the execution plan.

Suggested change

	.filter(lambda row: row["file_name"].endswith(".pdf"))
	.filter(col("file_name").str.endswith(".pdf"))

[omatthew98]

🔥