feat: Make the Postgres batched read limit configurable

Address review feedback: rename MAX_BATCHED_READ_KEYS to MAX_BATCHED_RESULT_ROWS, expose it as the max_batched_result_rows online store setting (0 disables batching), and keep the generic path for single-view requests. Cover the gate on both the sync and async paths, and check that the fallback reads each view. Signed-off-by: DKilkenny <danny.kilkenny18@gmail.com>
DKilkenny · Sep 11, 2026 · Sep 16, 2026 · Sep 16, 2026 · Sep 17, 2026 · Sep 30, 2026
commit 4efa7daf8fa48e8f860bbcf70ada48df401a450f
@@ -34,7 +34,7 @@ When the server processes a `get_online_features()` call, it groups the requeste
 {% hint style="info" %}
 **Redis exception:** The Redis online store overrides `_read_features_per_fv()` to batch all `HMGET` commands across every feature view into a **single pipeline execution**. Because all feature views for the same entity share one Redis hash key, the number of Redis round trips is always **1**, regardless of how many feature views the request touches. This means the "fewer feature views" guideline is less critical for Redis than for other stores — but consolidating feature views still reduces serialization and protobuf overhead at the application layer.
 
-**PostgreSQL exception:** PostgreSQL stores each feature view in its own table and overrides `_read_features_per_fv()` to combine the per-view reads into a **single `UNION ALL` statement**, so a request touching any number of feature views costs **1** query on both the sync and async paths. This helps most where round trips dominate — a handful of entities against a remote database. Past `MAX_BATCHED_READ_KEYS` entity keys in a request the store reverts to one query per view, because at that size the saved round trips no longer offset holding every view's rows at once. As with Redis, consolidating feature views still reduces application-layer overhead.
+**PostgreSQL exception:** PostgreSQL stores each feature view in its own table and overrides `_read_features_per_fv()` to combine the per-view reads into a **single `UNION ALL` statement**, so a request touching any number of feature views costs **1** query on both the sync and async paths. This helps most where round trips dominate — a handful of entities against a remote database. Requests larger than `max_batched_result_rows` (entity rows × feature views, default 2048) fall back to one query per view, because at that size the saved round trips no longer offset holding every view's rows at once. A single-view request is never batched, since there is nothing to combine. Set `max_batched_result_rows: 0` in the online store config to turn batching off. As with Redis, consolidating feature views still reduces application-layer overhead.
 {% endhint %}
 
 ### Feature services are free (and can be faster)

@@ -8,6 +8,8 @@ The PostgreSQL online store provides support for materializing feature values in
 
 * `sslmode` defaults to `require`, which encrypts the connection without certificate verification. To disable SSL (e.g. for local development), set `sslmode: disable`. For certificate verification, set `sslmode` to `verify-ca` or `verify-full` and provide the corresponding `sslrootcert_path` (and optionally `sslcert_path` and `sslkey_path` for mutual TLS)
 
+* Reads that span several feature views are combined into one query. `max_batched_result_rows` (default 2048) caps the request size, in entity rows × feature views, that is batched this way; set it to `0` to always issue one query per feature view
+
 ## Getting started
 In order to use this online store, you'll need to run `pip install 'feast[postgres]'`. You can get started by then running `feast init -t postgres`.
 

@@ -20,6 +20,7 @@
 from psycopg import AsyncConnection, sql
 from psycopg.connection import Connection
 from psycopg_pool import AsyncConnectionPool, ConnectionPool
+from pydantic import NonNegativeInt
 
 from feast import Entity, FeatureView, ValueType, utils
 from feast.filter_models import (
@@ -56,14 +57,17 @@
     "inner_product": "<#>",
 }
 
-# Above this many entity keys in a request, batching stops paying for itself: the
-# saved round trips are amortized away by the payload, while holding every feature
-# view's rows at once makes peak memory grow with the number of views. Measured
-# against Postgres 16, a 10-view request is ~3-5x faster at 1-100 entities and a
-# wash at 5000, where the combined result costs tens of MB per in-flight request.
-# Past the threshold the generic one-query-per-view path is used instead, which
-# processes and releases a single view at a time.
-MAX_BATCHED_READ_KEYS = 2048
+# Default for ``max_batched_result_rows``: the largest request, measured as entity
+# rows x feature views, that is read in one batched query. Past that, batching stops
+# paying for itself: the saved round trips are amortized away by the payload, while
+# holding every feature view's rows at once makes peak memory grow with the number of
+# views. The measure is a proxy taken from the request shape; the rows actually
+# returned also scale with the number of requested features. Measured against
+# Postgres 16, a 10-view request is ~3-5x faster at 1-100 entities and a wash at
+# 5000, where the combined result costs tens of MB per in-flight request. Past the
+# limit the generic one-query-per-view path is used instead, which processes and
+# releases a single view at a time.
+MAX_BATCHED_RESULT_ROWS = 2048
 
 _PG_COMPARISON_OPS: Dict[str, str] = {
     "eq": "=",
@@ -179,6 +183,9 @@ def _pg_filter_col_and_val(value: Any) -> Tuple[str, Any]:
 class PostgreSQLOnlineStoreConfig(PostgreSQLConfig, VectorStoreConfig):
     type: Literal["postgres"] = "postgres"
     enable_openai_compatible_store: Optional[bool] = False
+    max_batched_result_rows: NonNegativeInt = MAX_BATCHED_RESULT_ROWS
+    """Largest request, as entity rows x feature views, read in one batched query.
+    Set to 0 to always read one feature view per query."""
 
 
 class PostgreSQLOnlineStore(OnlineStore):
@@ -481,7 +488,7 @@ def _read_features_per_fv(
         in its own table here, so the per-view reads can be combined with UNION ALL
         and split apart again afterwards.
         """
-        if self._too_large_to_batch(grouped_refs, join_key_values):
+        if not self._should_batch(config, grouped_refs, join_key_values):
             return super()._read_features_per_fv(
                 config,
                 grouped_refs,
@@ -495,8 +502,6 @@ def _read_features_per_fv(
         reads = self._prepare_batched_reads(
             config, grouped_refs, join_key_values, entity_name_to_join_key_map
         )
-        if not reads:
-            return
 
         query, params = self._construct_batched_query_and_params(config, reads)
         with self._get_conn(config, autocommit=True) as conn, conn.cursor() as cur:
@@ -526,7 +531,7 @@ async def _read_features_per_fv_async(
         The generic async path issues the per-view queries concurrently. Combining
         them still replaces those N queries with one.
         """
-        if self._too_large_to_batch(grouped_refs, join_key_values):
+        if not self._should_batch(config, grouped_refs, join_key_values):
             return await super()._read_features_per_fv_async(
                 config,
                 grouped_refs,
@@ -540,8 +545,6 @@ async def _read_features_per_fv_async(
         reads = self._prepare_batched_reads(
             config, grouped_refs, join_key_values, entity_name_to_join_key_map
         )
-        if not reads:
-            return
 
         query, params = self._construct_batched_query_and_params(config, reads)
         async with self._get_conn_async(config, autocommit=True) as conn:
@@ -622,18 +625,26 @@ def _construct_batched_query_and_params(
         return sql.SQL(" UNION ALL ").join(branches), params
 
     @staticmethod
-    def _too_large_to_batch(grouped_refs: List, join_key_values: Dict) -> bool:
-        """Whether this request is big enough that batching would cost more than it saves.
+    def _should_batch(
+        config: RepoConfig, grouped_refs: List, join_key_values: Dict
+    ) -> bool:
+        """Whether to read this request's feature views in one batched query.
 
-        Estimated from the request shape rather than from the resolved keys, so that
-        deciding against batching costs nothing. The number of entity rows in the
-        request is an upper bound on the unique keys any one view will read.
-        See :data:`MAX_BATCHED_READ_KEYS`.
+        A single view has no round trips to save, so it keeps the generic path. The
+        size check is estimated from the request shape rather than from the resolved
+        keys, so that deciding against batching costs nothing. The number of entity
+        rows in the request is an upper bound on the unique keys any one view will
+        read. See :data:`MAX_BATCHED_RESULT_ROWS`.
         """
+        if len(grouped_refs) < 2:
+            return False
+        limit = getattr(
+            config.online_store, "max_batched_result_rows", MAX_BATCHED_RESULT_ROWS
+        )
         entity_rows = max(
             (len(values) for values in join_key_values.values()), default=0
         )
-        return entity_rows * len(grouped_refs) > MAX_BATCHED_READ_KEYS
+        return entity_rows * len(grouped_refs) <= limit
 
     @staticmethod
     def _bucket_row(buckets: List[Dict[bytes, List[Tuple]]], row: Tuple) -> None:

@@ -6,11 +6,12 @@
 from unittest.mock import patch
 
 import pytest
+from pydantic import ValidationError
 
 from feast import Entity, FeatureView
 from feast.field import Field
 from feast.infra.online_stores.postgres_online_store.postgres import (
-    MAX_BATCHED_READ_KEYS,
+    MAX_BATCHED_RESULT_ROWS,
     PostgreSQLOnlineStore,
 )
 from feast.protos.feast.serving.ServingService_pb2 import GetOnlineFeaturesResponse
@@ -36,7 +37,7 @@ def _feature_view(name: str, feature: str, entity: Entity = DRIVER) -> FeatureVi
     )
 
 
-def _config() -> RepoConfig:
+def _config(**online_store) -> RepoConfig:
     return RepoConfig(
         project="proj",
         provider="local",
@@ -49,6 +50,7 @@ def _config() -> RepoConfig:
             "db_schema": "public",
             "user": "root",
             "password": "test",  # pragma: allowlist secret
+            **online_store,
         },
         entity_key_serialization_version=3,
     )
@@ -120,9 +122,9 @@ async def _fake_conn_async(rows: List[Tuple], executed: List):
     yield _AsyncFakeConnection(rows, executed)
 
 
-def _args(grouped_refs, join_key_values, response):
+def _args(grouped_refs, join_key_values, response, config=None):
     return dict(
-        config=_config(),
+        config=config or _config(),
         grouped_refs=grouped_refs,
         join_key_values=join_key_values,
         entity_name_to_join_key_map=JOIN_KEY_MAP,
@@ -136,7 +138,7 @@ def _drivers(*ids) -> dict:
     return {"driver_id": [ValueProto(int64_val=i) for i in ids]}
 
 
-def _read(store, grouped_refs, rows, join_key_values=None):
+def _read(store, grouped_refs, rows, join_key_values=None, config=None):
     """Drive the sync override against a fake connection."""
     executed: List = []
     response = GetOnlineFeaturesResponse()
@@ -146,12 +148,12 @@ def _read(store, grouped_refs, rows, join_key_values=None):
         lambda self, cfg, autocommit=False: _fake_conn(rows, executed),
     ):
         store._read_features_per_fv(
-            **_args(grouped_refs, join_key_values or _drivers(1, 2), response)
+            **_args(grouped_refs, join_key_values or _drivers(1, 2), response, config)
         )
     return executed, response
 
 
-async def _read_async(store, grouped_refs, rows, join_key_values=None):
+async def _read_async(store, grouped_refs, rows, join_key_values=None, config=None):
     """Drive the async override against a fake connection."""
     executed: List = []
     response = GetOnlineFeaturesResponse()
@@ -161,25 +163,27 @@ async def _read_async(store, grouped_refs, rows, join_key_values=None):
         lambda self, cfg, autocommit=False: _fake_conn_async(rows, executed),
     ):
         await store._read_features_per_fv_async(
-            **_args(grouped_refs, join_key_values or _drivers(1, 2), response)
+            **_args(grouped_refs, join_key_values or _drivers(1, 2), response, config)
         )
     return executed, response
 
 
-async def _read_sync(store, grouped_refs, rows, join_key_values=None):
+async def _read_sync(store, grouped_refs, rows, join_key_values=None, config=None):
     """Await-compatible wrapper so both overrides share one test body."""
-    return _read(store, grouped_refs, rows, join_key_values)
+    return _read(store, grouped_refs, rows, join_key_values, config)
 
 
 READERS = [pytest.param(_read_sync, id="sync"), pytest.param(_read_async, id="async")]
 
 
-def _two_views_same_feature_name(store):
-    """Two views exposing the SAME feature name, both keyed on the same entity.
+@pytest.mark.parametrize("read", READERS)
+async def test_views_read_in_one_query_and_demultiplexed_by_tag(read):
+    """Two views sharing an entity key AND a feature name come back in one query.
 
     Distinct feature names make a demux failure invisible: _process_rows keys values
     by feature name, so a cross-contaminated bucket still yields the right value.
     """
+    store = PostgreSQLOnlineStore()
     grouped_refs = [
         (_feature_view("fv_a", "feat"), ["feat"]),
         (_feature_view("fv_b", "feat"), ["feat"]),
@@ -193,36 +197,14 @@ def _two_views_same_feature_name(store):
         (0, shared_key, "feat", _int(11), TS),
         (1, shared_key, "feat", _int(22), TS),
     ]
-    return grouped_refs, rows
 
-
-@pytest.mark.parametrize("read", READERS)
-async def test_single_round_trip_for_multiple_feature_views(read):
-    """Three feature views must produce one execute, not three."""
-    store = PostgreSQLOnlineStore()
-    grouped_refs = [
-        (_feature_view("fv_a", "feat_a"), ["feat_a"]),
-        (_feature_view("fv_b", "feat_b"), ["feat_b"]),
-        (_feature_view("fv_c", "feat_c"), ["feat_c"]),
-    ]
-
-    executed, _ = await read(store, grouped_refs, rows=[])
+    executed, response = await read(store, grouped_refs, rows)
 
     assert len(executed) == 1
     statement = executed[0][0].as_string(None)
-    assert statement.count("UNION ALL") == 2
-    for table in ("proj_fv_a", "proj_fv_b", "proj_fv_c"):
+    assert statement.count("UNION ALL") == 1
+    for table in ("proj_fv_a", "proj_fv_b"):
         assert table in statement
-
-
-@pytest.mark.parametrize("read", READERS)
-async def test_rows_are_demultiplexed_by_tag(read):
-    """Two views sharing an entity key AND a feature name must not cross."""
-    store = PostgreSQLOnlineStore()
-    grouped_refs, rows = _two_views_same_feature_name(store)
-
-    _, response = await read(store, grouped_refs, rows)
-
     assert response.results[0].values[0].int64_val == 11
     assert response.results[1].values[0].int64_val == 22
 
@@ -260,38 +242,67 @@ def test_per_view_entity_bookkeeping_is_not_shared():
     assert [v.int64_val for v in response.results[1].values] == [70, 70, 71]
 
 
+@pytest.mark.parametrize("read", READERS)
 @pytest.mark.parametrize(
-    ("entity_rows", "exp_executes"),
-    [(2, 1), (MAX_BATCHED_READ_KEYS, 0)],
-    ids=["under_threshold", "over_threshold"],
+    ("views", "entity_rows", "online_store", "exp_batched"),
+    [
+        (2, 2, {}, True),
+        (2, MAX_BATCHED_RESULT_ROWS, {}, False),
+        (1, 2, {}, False),
+        (2, 2, {"max_batched_result_rows": 4}, True),
+        (2, 2, {"max_batched_result_rows": 3}, False),
+        (2, 2, {"max_batched_result_rows": 0}, False),
+    ],
+    ids=[
+        "under_default_limit",
+        "over_default_limit",
+        "single_view",
+        "at_configured_limit",
+        "over_configured_limit",
+        "zero_disables_batching",
+    ],
 )
-def test_batching_is_gated_by_the_key_threshold(entity_rows, exp_executes):
-    """Past the key threshold, batching is skipped rather than ballooning memory.
+async def test_when_batching_applies(
+    read, views, entity_rows, online_store, exp_batched
+):
+    """Batching needs at least two views and a request within the row limit.
 
-    Two views, so the row count only has to reach the limit to exceed it. A batched
-    request issues exactly one statement; the generic fallback issues none, because
-    it reads each view through online_read instead.
+    A batched request issues one statement and never calls online_read; the generic
+    path issues none itself and reads each view through online_read instead.
     """
     store = PostgreSQLOnlineStore()
-    grouped_refs = [
-        (_feature_view("fv_a", "feat_a"), ["feat_a"]),
-        (_feature_view("fv_b", "feat_b"), ["feat_b"]),
-    ]
-    join_key_values = _drivers(*range(entity_rows))
+    names = [f"fv_{i}" for i in range(views)]
+    grouped_refs = [(_feature_view(n, f"feat_{n}"), [f"feat_{n}"]) for n in names]
+    per_view_reads: List[str] = []
 
     def one_row_per_key(self, config, table, entity_keys, requested_features=None):
+        per_view_reads.append(table.name)
         return [(None, None)] * len(entity_keys)
 
-    with patch.object(PostgreSQLOnlineStore, "online_read", one_row_per_key):
-        executed, _ = _read(store, grouped_refs, [], join_key_values)
+    async def one_row_per_key_async(self, *args, **kwargs):
+        return one_row_per_key(self, *args, **kwargs)
 
-    assert len(executed) == exp_executes
+    with (
+        patch.object(PostgreSQLOnlineStore, "online_read", one_row_per_key),
+        patch.object(PostgreSQLOnlineStore, "online_read_async", one_row_per_key_async),
+    ):
+        executed, _ = await read(
+            store,
+            grouped_refs,
+            [],
+            _drivers(*range(entity_rows)),
+            _config(**online_store),
+        )
 
+    assert len(executed) == (1 if exp_batched else 0)
+    assert sorted(per_view_reads) == ([] if exp_batched else names)
 
-def test_empty_grouped_refs_issues_no_query():
-    store = PostgreSQLOnlineStore()
-    executed, _ = _read(store, grouped_refs=[], rows=[])
-    assert executed == []
+
+def test_max_batched_result_rows_config():
+    """Defaults to 2048, and rejects negatives rather than silently never batching."""
+    assert _config().online_store.max_batched_result_rows == 2048
+    with pytest.raises(ValidationError):
+        _config(max_batched_result_rows=-1)
 
 
 @pytest.mark.parametrize(