feat: Add RetrievalJob.to_arrow_reader and shared offline metrics helper

Extract to_arrow()'s metrics/audit finally-block into a reusable _emit_offline_store_request_metrics helper, and give every RetrievalJob a default to_arrow_reader() that wraps to_arrow(). This keeps every existing offline store working unchanged while giving stores that can page results natively (Trino, next) a hook to override without duplicating metrics code. Signed-off-by: dbbvitor <vitor.diniz@gympass.com>
feast-dev · dbbvitor · Sep 24, 2026 · Sep 24, 2026 · Sep 25, 2026 · Sep 25, 2026
commit e7cadc345a546cc61ac5cc690d659efb890bfe82
@@ -89,6 +89,48 @@ def _extract_retrieval_metadata(job: "RetrievalJob") -> tuple:
     return [], 0
 
 
+def _emit_offline_store_request_metrics(
+    job: "RetrievalJob",
+    method: str,
+    status_label: str,
+    row_count: int,
+    elapsed: float,
+) -> None:
+    """Record offline-store request metrics and the audit log entry. Never raises."""
+    try:
+        from feast import metrics as feast_metrics
+
+        if feast_metrics._config.offline_features:
+            feast_metrics.offline_store_request_total.labels(
+                method=method, status=status_label
+            ).inc()
+            feast_metrics.offline_store_request_latency_seconds.labels(
+                method=method
+            ).observe(elapsed)
+            feast_metrics.offline_store_row_count.labels(method=method).observe(
+                row_count
+            )
+
+        if feast_metrics._config.audit_logging:
+            feature_views, feature_count = _extract_retrieval_metadata(job)
+            end_dt = datetime.now(tz=timezone.utc)
+            start_dt = end_dt - timedelta(seconds=elapsed)
+            feast_metrics.emit_offline_audit_log(
+                method=method,
+                feature_views=feature_views,
+                feature_count=feature_count,
+                row_count=row_count,
+                status=status_label,
+                start_time=start_dt.isoformat(),
+                end_time=end_dt.isoformat(),
+                duration_ms=elapsed * 1000,
+            )
+    except Exception:
+        logging.getLogger(__name__).debug(
+            "Failed to record offline store metrics", exc_info=True
+        )
+
+
 class RetrievalJob(ABC):
     """A RetrievalJob manages the execution of a query to retrieve data from the offline store."""
 
@@ -181,40 +223,13 @@ def to_arrow(
             status_label = "error"
             raise
         finally:
-            try:
-                from feast import metrics as feast_metrics
-
-                elapsed = time.monotonic() - start_wall
-
-                if feast_metrics._config.offline_features:
-                    feast_metrics.offline_store_request_total.labels(
-                        method="to_arrow", status=status_label
-                    ).inc()
-                    feast_metrics.offline_store_request_latency_seconds.labels(
-                        method="to_arrow"
-                    ).observe(elapsed)
-                    feast_metrics.offline_store_row_count.labels(
-                        method="to_arrow"
-                    ).observe(row_count)
-
-                if feast_metrics._config.audit_logging:
-                    feature_views, feature_count = _extract_retrieval_metadata(self)
-                    end_dt = datetime.now(tz=timezone.utc)
-                    start_dt = end_dt - timedelta(seconds=elapsed)
-                    feast_metrics.emit_offline_audit_log(
-                        method="to_arrow",
-                        feature_views=feature_views,
-                        feature_count=feature_count,
-                        row_count=row_count,
-                        status=status_label,
-                        start_time=start_dt.isoformat(),
-                        end_time=end_dt.isoformat(),
-                        duration_ms=elapsed * 1000,
-                    )
-            except Exception:
-                logging.getLogger(__name__).debug(
-                    "Failed to record offline store metrics", exc_info=True
-                )
+            _emit_offline_store_request_metrics(
+                job=self,
+                method="to_arrow",
+                status_label=status_label,
+                row_count=row_count,
+                elapsed=time.monotonic() - start_wall,
+            )
 
         if self.on_demand_feature_views:
             # Build a mapping of ODFV name to requested feature names
@@ -288,6 +303,18 @@ def to_arrow(
 
         return features_table
 
+    def to_arrow_reader(
+        self, timeout: Optional[int] = None
+    ) -> pyarrow.RecordBatchReader:
+        """
+        Returns the result as a stream of record batches.
+
+        The default materializes the full result via ``to_arrow()``; offline stores that
+        can page results natively override this to bound memory.
+        """
+        table = self.to_arrow(timeout=timeout)
+        return pyarrow.RecordBatchReader.from_batches(table.schema, table.to_batches())
+
     def to_tensor(
         self,
         kind: str = "torch",

@@ -249,6 +249,96 @@ def test_to_arrow_timeout(retrieval_job, timeout: Optional[int]):
         mock_to_arrow_internal.assert_called_once_with(timeout=timeout)
 
 
+def test_default_to_arrow_reader_matches_to_arrow():
+    table = pyarrow.table({"a": [1, 2, 3]})
+
+    class _Job(RetrievalJob):
+        def _to_df_internal(self, timeout=None):
+            return table.to_pandas()
+
+        def _to_arrow_internal(self, timeout=None):
+            return table
+
+        @property
+        def full_feature_names(self):
+            return False
+
+        @property
+        def on_demand_feature_views(self):
+            return []
+
+        @property
+        def metadata(self):
+            return None
+
+    reader = _Job().to_arrow_reader()
+
+    assert isinstance(reader, pyarrow.RecordBatchReader)
+    assert reader.read_all().equals(table)
+
+
+def test_default_to_arrow_reader_passes_through_timeout():
+    """The default to_arrow_reader() must forward `timeout` to to_arrow(),
+    not silently drop it (e.g. under a mutant that hardcodes None)."""
+    table = pyarrow.table({"a": [1]})
+
+    class _Job(RetrievalJob):
+        def _to_df_internal(self, timeout=None):
+            return table.to_pandas()
+
+        def _to_arrow_internal(self, timeout=None):
+            return table
+
+        @property
+        def full_feature_names(self):
+            return False
+
+        @property
+        def on_demand_feature_views(self):
+            return []
+
+        @property
+        def metadata(self):
+            return None
+
+    job = _Job()
+    with patch.object(job, "to_arrow", wraps=job.to_arrow) as mock_to_arrow:
+        job.to_arrow_reader(timeout=42)
+
+    mock_to_arrow.assert_called_once_with(timeout=42)
+
+
+def test_default_to_arrow_reader_yields_empty_table_as_one_empty_batch():
+    """An empty result must still produce a valid (zero-row) reader rather
+    than raising, so `RecordBatchReader.from_batches` gets a non-empty
+    batches iterable even for `table.to_batches()` on 0 rows."""
+    table = pyarrow.table({"a": pyarrow.array([], type=pyarrow.int64())})
+
+    class _Job(RetrievalJob):
+        def _to_df_internal(self, timeout=None):
+            return table.to_pandas()
+
+        def _to_arrow_internal(self, timeout=None):
+            return table
+
+        @property
+        def full_feature_names(self):
+            return False
+
+        @property
+        def on_demand_feature_views(self):
+            return []
+
+        @property
+        def metadata(self):
+            return None
+
+    reader = _Job().to_arrow_reader()
+
+    assert reader.schema == table.schema
+    assert reader.read_all().num_rows == 0
+
+
 @pytest.mark.parametrize(
     "repo_path, uri, expected",
     [

@@ -1352,6 +1352,146 @@ def test_different_methods_tracked_independently(self):
         )
 
 
+class TestEmitOfflineStoreRequestMetrics:
+    """Tests for the shared `_emit_offline_store_request_metrics` helper."""
+
+    def test_records_method_label_when_offline_features_enabled(self):
+        import feast.metrics as feast_metrics
+        from feast.infra.offline_stores.offline_store import (
+            _emit_offline_store_request_metrics,
+        )
+
+        with patch.object(
+            feast_metrics,
+            "_config",
+            feast_metrics._MetricsFlags(offline_features=True, audit_logging=False),
+        ):
+            before = offline_store_request_total.labels(
+                method="to_arrow_reader", status="success"
+            )._value.get()
+
+            _emit_offline_store_request_metrics(
+                job=MagicMock(),
+                method="to_arrow_reader",
+                status_label="success",
+                row_count=5,
+                elapsed=0.1,
+            )
+
+            after = offline_store_request_total.labels(
+                method="to_arrow_reader", status="success"
+            )._value.get()
+
+        assert after == before + 1
+
+    def test_skips_metrics_when_offline_features_disabled(self):
+        import feast.metrics as feast_metrics
+        from feast.infra.offline_stores.offline_store import (
+            _emit_offline_store_request_metrics,
+        )
+
+        with patch.object(
+            feast_metrics,
+            "_config",
+            feast_metrics._MetricsFlags(offline_features=False, audit_logging=False),
+        ):
+            before = offline_store_request_total.labels(
+                method="to_arrow_reader", status="success"
+            )._value.get()
+
+            _emit_offline_store_request_metrics(
+                job=MagicMock(),
+                method="to_arrow_reader",
+                status_label="success",
+                row_count=5,
+                elapsed=0.1,
+            )
+
+            after = offline_store_request_total.labels(
+                method="to_arrow_reader", status="success"
+            )._value.get()
+
+        assert after == before
+
+    def test_emits_audit_log_with_correct_fields_when_enabled(self):
+        import feast.metrics as feast_metrics
+        from feast.infra.offline_stores.offline_store import (
+            _emit_offline_store_request_metrics,
+        )
+
+        job = MagicMock()
+        job.metadata.features = ["fv1:f1", "fv1:f2"]
+
+        with (
+            patch.object(
+                feast_metrics,
+                "_config",
+                feast_metrics._MetricsFlags(offline_features=False, audit_logging=True),
+            ),
+            patch.object(feast_metrics, "emit_offline_audit_log") as mock_emit,
+        ):
+            _emit_offline_store_request_metrics(
+                job=job,
+                method="to_arrow_reader",
+                status_label="success",
+                row_count=7,
+                elapsed=2.0,
+            )
+
+        mock_emit.assert_called_once()
+        kwargs = mock_emit.call_args.kwargs
+        assert kwargs["method"] == "to_arrow_reader"
+        assert kwargs["feature_views"] == ["fv1"]
+        assert kwargs["feature_count"] == 2
+        assert kwargs["row_count"] == 7
+        assert kwargs["status"] == "success"
+        assert kwargs["duration_ms"] == pytest.approx(2000.0)
+        start_dt = datetime.fromisoformat(kwargs["start_time"])
+        end_dt = datetime.fromisoformat(kwargs["end_time"])
+        assert (end_dt - start_dt).total_seconds() == pytest.approx(2.0)
+
+    def test_skips_audit_log_when_disabled(self):
+        import feast.metrics as feast_metrics
+        from feast.infra.offline_stores.offline_store import (
+            _emit_offline_store_request_metrics,
+        )
+
+        with (
+            patch.object(
+                feast_metrics,
+                "_config",
+                feast_metrics._MetricsFlags(
+                    offline_features=False, audit_logging=False
+                ),
+            ),
+            patch.object(feast_metrics, "emit_offline_audit_log") as mock_emit,
+        ):
+            _emit_offline_store_request_metrics(
+                job=MagicMock(),
+                method="to_arrow_reader",
+                status_label="success",
+                row_count=7,
+                elapsed=2.0,
+            )
+
+        mock_emit.assert_not_called()
+
+    def test_never_raises_when_metrics_recording_fails(self):
+        from feast.infra.offline_stores.offline_store import (
+            _emit_offline_store_request_metrics,
+        )
+
+        with patch("feast.metrics.offline_store_request_total") as counter:
+            counter.labels.side_effect = RuntimeError("boom")
+            _emit_offline_store_request_metrics(
+                job=MagicMock(),
+                method="to_arrow",
+                status_label="error",
+                row_count=0,
+                elapsed=0.0,
+            )
+
+
 class TestEmitAuditLogs:
     """Tests for structured JSON audit log emission."""
 
@@ -1657,6 +1797,31 @@ def test_success_increments_counter_and_records_latency(self):
             > before_latency
         )
 
+    def test_elapsed_is_end_minus_start_not_sum(self):
+        """`time.monotonic() - start_wall` must stay subtraction: kills a
+        mutant that turns it into addition, which the "> before" latency
+        assertion above can't distinguish from a real (small) elapsed time."""
+        import pyarrow as pa
+
+        from feast.infra.offline_stores import offline_store as offline_store_module
+
+        table = pa.table({"col": [1]})
+        job = self._make_job(table)
+
+        before_sum = offline_store_request_latency_seconds.labels(
+            method="to_arrow"
+        )._sum.get()
+
+        with patch.object(
+            offline_store_module.time, "monotonic", side_effect=[100.0, 100.25]
+        ):
+            job.to_arrow()
+
+        after_sum = offline_store_request_latency_seconds.labels(
+            method="to_arrow"
+        )._sum.get()
+        assert after_sum - before_sum == pytest.approx(0.25)
+
     def test_error_increments_error_counter(self):
         job = self._make_job(None, raise_on_internal=RuntimeError("query failed"))