feat: Add RaySource, to_ray_dataset first-class method, docs, and tests by ntkathole · Pull Request #6343 · feast-dev/feast

ntkathole · 2026-04-28T06:48:48Z

What this PR does / why we need it:

Introduces RaySource, a data source descriptor that lets Feast load any Ray Data-readable format (Parquet, CSV, JSON, HuggingFace datasets, MongoDB, binary files, images, TFRecords, WebDataset, SQL, and text) without requiring an intermediate Parquet file. Makes to_ray_dataset() a first-class method on RetrievalJob and FeatureStore. Wires up distributed materialization in the Ray compute engine. Adds reference documentation and unit tests.

Which issue(s) this PR fixes:

Fixes #5568

Introduces RaySource, a pure-metadata data source descriptor that lets Feast load any Ray Data-readable format (Parquet, CSV, JSON, HuggingFace datasets, MongoDB, binary files, images, TFRecords, WebDataset, SQL, and text) without requiring an intermediate Parquet file. Makes to_ray_dataset() a first-class method on RetrievalJob and FeatureStore. Wires up distributed materialization in the Ray compute engine. Adds reference documentation and unit tests. Signed-off-by: ntkathole <[email protected]>

devin-ai-integration

Devin Review found 1 new potential issue.

View 18 additional findings in Devin Review.

devin-ai-integration · 2026-04-30T16:59:49Z

+                ray_ds = retrieval_job.to_ray_dataset()
+
+                # Distributed online store write — each Ray worker writes its shard
+                write_to_online_store_from_ray_ds(
+                    ray_ds=ray_ds,
                    feature_view=feature_view,
-                    table=arrow_table,
-                    progress=lambda x: None,
+                    online_store=self.online_store,
+                    repo_config=self.repo_config,
                )

-            # For derived views, also ensure data is written to sink_source if it exists
-            # This is critical for feature view chaining to work properly
-            sink_source = getattr(feature_view, "sink_source", None)
-            if sink_source is not None:
-                logger.debug(
-                    f"Writing derived view {feature_view.name} to sink_source: {sink_source.path}"
+                # offline_write_batch and sink_source are independent — both can
+                # apply when a feature view has offline=True AND a sink_source.
+                if getattr(feature_view, "offline", False):
+                    import pyarrow as pa
+                    import ray as _ray
+
+                    arrow_table = pa.concat_tables(_ray.get(ray_ds.to_arrow_refs()))
+                    self.offline_store.offline_write_batch(
+                        config=self.repo_config,
+                        feature_view=feature_view,
+                        table=arrow_table,
+                        progress=lambda x: None,
+                    )
+
+                sink_source = getattr(feature_view, "sink_source", None)
+                if sink_source is not None:
+                    logger.debug(
+                        f"Writing derived view {feature_view.name} to sink_source: {sink_source.path}"
+                    )
+                    try:
+                        ray_ds.write_parquet(sink_source.path)
+                    except Exception as e:
+                        logger.error(
+                            f"Failed to write to sink_source {sink_source.path}: {e}"
+                        )


🔴 Triple re-execution of lazy Ray pipeline causes data inconsistency and performance regression in legacy materialization

In the _materialize_from_offline_store RayRetrievalJob branch, ray_ds is a lazy pipeline that is never materialized (.materialize() is only called on a derived dataset inside write_to_online_store_from_ray_ds). Each subsequent use of ray_ds — to_arrow_refs() at line 202 and write_parquet() at line 216 — re-executes the entire pipeline from scratch, re-reading from the data source each time. This introduces a data consistency risk: if the source changes between executions (e.g., rows appended to Parquet), the online store, offline store, and sink_source could contain different data snapshots. The old code collected data once via retrieval_job.to_arrow() and reused the single arrow_table for all three writes, which was always consistent.

Comparison of old vs new behavior

Old code (removed in this PR):

arrow_table = retrieval_job.to_arrow() # single read write_to_online_store(arrow_table=arrow_table, ...) # reuses collected data self.offline_store.offline_write_batch(table=arrow_table, ...) # reuses ray_wrapper.from_arrow(arrow_table).write_parquet(...) # reuses

New code:

ray_ds = retrieval_job.to_ray_dataset() # lazy plan, not materialized write_to_online_store_from_ray_ds(ray_ds, ...) # execution #1 arrow_table = pa.concat_tables(_ray.get(ray_ds.to_arrow_refs())) # execution #2 ray_ds.write_parquet(sink_source.path) # execution #3

The fix is to call ray_ds = ray_ds.materialize() before any writes when the offline or sink paths are needed, so all downstream operations read from the in-memory materialized dataset.

Prompt for agents

In _materialize_from_offline_store (compute.py), after obtaining `ray_ds = retrieval_job.to_ray_dataset()` at line 186, the lazy dataset is used three separate times without being materialized: once inside write_to_online_store_from_ray_ds (which only materializes a derived dataset internally), once via to_arrow_refs() for offline_write_batch, and once via write_parquet() for sink_source. Each use triggers a separate execution of the entire lazy pipeline, re-reading from the data source. Fix: After line 186, check whether the offline write or sink_source write paths will be triggered. If so, materialize the dataset once so all downstream operations read from the cached in-memory result: ray_ds = retrieval_job.to_ray_dataset() # Materialize once if we need the data for multiple writes needs_offline = getattr(feature_view, 'offline', False) has_sink = getattr(feature_view, 'sink_source', None) is not None if needs_offline or has_sink: ray_ds = ray_ds.materialize() This ensures consistency across all three write targets and avoids redundant I/O. The write_to_online_store_from_ray_ds function will also benefit because map_batches on a materialized dataset reads from memory.

Was this helpful? React with 👍 or 👎 to provide feedback.

ntkathole self-assigned this Apr 28, 2026

ntkathole requested review from a team as code owners April 28, 2026 06:48

ntkathole requested review from haoxu0, shuchu and tokoko and removed request for a team April 28, 2026 06:48

ntkathole changed the title ~~feat(ray): Add RaySource, to_ray_dataset first-class method, docs, and tests~~ feat: Add RaySource, to_ray_dataset first-class method, docs, and tests Apr 28, 2026