# Copyright 2021 The Feast Authors

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# https://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

import asyncio

import contextlib

import itertools

import json

import logging

from collections import OrderedDict, defaultdict

from concurrent.futures import ThreadPoolExecutor

from datetime import datetime

from typing import (

Any,

Callable,

Dict,

List,

Literal,

NamedTuple,

Optional,

Sequence,

Tuple,

Union,

)

from aiobotocore.config import AioConfig

from pydantic import NonNegativeInt, StrictBool, StrictStr

from feast import Entity, FeatureView, utils

from feast.infra.infra_object import InfraObject

from feast.infra.online_stores.helpers import compute_entity_id, compute_versioned_name

from feast.infra.online_stores.online_store import OnlineStore

from feast.infra.supported_async_methods import SupportedAsyncMethods

from feast.infra.utils.aws_utils import dynamo_write_items_async

from feast.protos.feast.core.InfraObject_pb2 import InfraObject as InfraObjectProto

from feast.protos.feast.core.Registry_pb2 import Registry as RegistryProto

from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto

from feast.protos.feast.types.Value_pb2 import Value as ValueProto

from feast.repo_config import FeastConfigBaseModel, RepoConfig

from feast.stream_feature_view import StreamFeatureView

from feast.utils import get_user_agent

try:

import boto3

from aiobotocore import session

from boto3.dynamodb.types import TypeDeserializer

from botocore.config import Config

from botocore.exceptions import ClientError

except ImportError as e:

from feast.errors import FeastExtrasDependencyImportError

raise FeastExtrasDependencyImportError("aws", str(e))

logger = logging.getLogger(__name__)

# Default for ``max_batched_result_rows``: the largest request, measured as entity

# rows x feature views, whose feature views are read through shared BatchGetItem

# requests. Below it, packing the keys of every feature view into the same requests

# replaces one or more requests per view with as few as one in total. Past it, each

# view already fills several requests of its own, so packing saves at most one

# partially filled request per view, while holding every view's items at once makes

# peak memory grow with the number of views. Past the limit the generic path is used

# instead, which reads and releases a single view at a time.

MAX_BATCHED_RESULT_ROWS = 2048

class _BatchedRead(NamedTuple):

"""One feature view's share of a batched online read."""

table: FeatureView

requested_features: List[str]

table_name: str

entity_ids: List[str]

idxs: Tuple[List[int], ...]

output_len: int

class DynamoDBOnlineStoreConfig(FeastConfigBaseModel):

"""Online store config for DynamoDB store"""

type: Literal["dynamodb"] = "dynamodb"

"""Online store type selector"""

batch_size: int = 100

"""Number of items to retrieve in a DynamoDB BatchGetItem call.

DynamoDB supports up to 100 items per BatchGetItem request."""

endpoint_url: Union[str, None] = None

"""DynamoDB endpoint URL. Use for local development (e.g., http://localhost:8000)

or VPC endpoints for improved latency."""

region: StrictStr

"""AWS Region Name"""

table_name_template: StrictStr = "{project}.{table_name}"

"""DynamoDB table name template"""

consistent_reads: StrictBool = False

"""Whether to read from Dynamodb by forcing consistent reads"""

warmup_connections: StrictBool = False

"""Whether to warm up the connection pool on initialization, with

max_pool_connections concurrent lightweight calls"""

tags: Union[Dict[str, str], None] = None

"""AWS resource tags added to each table"""

session_based_auth: bool = False

"""AWS session based client authentication"""

max_read_workers: int = 10

"""Maximum number of parallel threads for batch read operations.

Higher values improve throughput for large batch reads but increase resource usage."""

max_batched_result_rows: NonNegativeInt = MAX_BATCHED_RESULT_ROWS

"""Largest request, as entity rows x feature views, whose feature views are read

through shared BatchGetItem requests. Set to 0 to always read one feature view

at a time."""

max_pool_connections: int = 50

"""Max number of connections for async Dynamodb operations.

Increase for high-throughput workloads."""

keepalive_timeout: float = 30.0

"""Keep-alive timeout in seconds for async Dynamodb connections.

Higher values help reuse connections under sustained load."""

connect_timeout: Union[int, float] = 5

"""The time in seconds until a timeout exception is thrown when attempting to make

an async connection. Lower values enable faster failure detection."""

read_timeout: Union[int, float] = 10

"""The time in seconds until a timeout exception is thrown when attempting to read

from an async connection. Lower values enable faster failure detection."""

total_max_retry_attempts: Union[int, None] = 3

"""Maximum number of total attempts that will be made on a single request.

Maps to `retries.total_max_attempts` in botocore.config.Config.

"""

retry_mode: Union[Literal["legacy", "standard", "adaptive"], None] = "adaptive"

"""The type of retry mode (aio)botocore should use.

Maps to `retries.mode` in botocore.config.Config.

'adaptive' mode provides intelligent retry with client-side rate limiting.

"""

class DynamoDBOnlineStore(OnlineStore):

"""

AWS DynamoDB implementation of the online store interface.

Attributes:

_dynamodb_client: Boto3 DynamoDB client.

_dynamodb_resource: Boto3 DynamoDB resource.

_aioboto_session: Async boto session.

_aioboto_client: Async boto client.

_aioboto_context_stack: Async context stack.

_type_deserializer: Cached TypeDeserializer instance for performance.

"""

_dynamodb_client = None

_dynamodb_resource = None

# Class-level cached TypeDeserializer to avoid per-request instantiation

_type_deserializer: Optional[TypeDeserializer] = None

def __init__(self):

super().__init__()

self._aioboto_session = None

self._aioboto_client = None

self._aioboto_context_stack = None

# Initialize cached TypeDeserializer if not already done

if DynamoDBOnlineStore._type_deserializer is None:

DynamoDBOnlineStore._type_deserializer = TypeDeserializer()

async def initialize(self, config: RepoConfig):

online_config = config.online_store

client = await self._get_aiodynamodb_client(

online_config.region,

online_config.max_pool_connections,

online_config.keepalive_timeout,

online_config.connect_timeout,

online_config.read_timeout,

online_config.total_max_retry_attempts,

online_config.retry_mode,

online_config.endpoint_url,

)

if online_config.warmup_connections:

# Concurrent calls, so that each one needs its own connection; a

# sequential call would reuse the connection of the previous one.

# describe_endpoints is not tied to a table and, unlike

# describe_limits, has no documented once-a-minute call limit.

results = await asyncio.gather(

*(

client.describe_endpoints()

for _ in range(online_config.max_pool_connections)

),

return_exceptions=True,

)

errors = [r for r in results if isinstance(r, BaseException)]

if errors:

logger.warning(

"Failed to warmup %d of %d DynamoDB connections",

len(errors),

len(results),

exc_info=errors[0],

)

async def close(self):

await self._aiodynamodb_close()

def _get_aioboto_session(self):

if self._aioboto_session is None:

logger.debug("initializing the aiobotocore session")

self._aioboto_session = session.get_session()

return self._aioboto_session

async def _get_aiodynamodb_client(

self,

region: str,

max_pool_connections: int,

keepalive_timeout: float,

connect_timeout: Union[int, float],

read_timeout: Union[int, float],

total_max_retry_attempts: Union[int, None],

retry_mode: Union[Literal["legacy", "standard", "adaptive"], None],

endpoint_url: Optional[str] = None,

):

if self._aioboto_client is None:

logger.debug("initializing the aiobotocore dynamodb client")

retries: Dict[str, Any] = {}

if total_max_retry_attempts is not None:

retries["total_max_attempts"] = total_max_retry_attempts

if retry_mode is not None:

retries["mode"] = retry_mode

# Build client kwargs, including endpoint_url for VPC endpoints or local testing

client_kwargs: Dict[str, Any] = {

"region_name": region,

"config": AioConfig(

max_pool_connections=max_pool_connections,

connect_timeout=connect_timeout,

read_timeout=read_timeout,

retries=retries if retries else None,

connector_args={"keepalive_timeout": keepalive_timeout},

),

}

if endpoint_url:

client_kwargs["endpoint_url"] = endpoint_url

client_context = self._get_aioboto_session().create_client(

"dynamodb",

**client_kwargs,

)

self._aioboto_context_stack = contextlib.AsyncExitStack()

self._aioboto_client = (

await self._aioboto_context_stack.enter_async_context(client_context)

)

return self._aioboto_client

async def _aiodynamodb_close(self):

if self._aioboto_context_stack:

await self._aioboto_context_stack.aclose()

self._aioboto_context_stack = None

self._aioboto_client = None

if self._aioboto_session:

self._aioboto_session = None

@property

def async_supported(self) -> SupportedAsyncMethods:

return SupportedAsyncMethods(read=True, write=True)

@staticmethod

def _table_tags(online_config, table_instance) -> list[dict[str, str]]:

table_instance_tags = table_instance.tags or {}

online_tags = online_config.tags or {}

common_tags = [

{"Key": key, "Value": table_instance_tags.get(key) or value}

for key, value in online_tags.items()

]

table_tags = [

{"Key": key, "Value": value}

for key, value in table_instance_tags.items()

if key not in online_tags

]

return common_tags + table_tags

@staticmethod

def _update_tags(dynamodb_client, table_name: str, new_tags: list[dict[str, str]]):

"""Update DynamoDB table tags using a diff-based approach.

Instead of removing all tags and re-adding them (which is vulnerable to

the eventual-consistency window between UntagResource and TagResource),

this method computes the minimal set of changes needed:

- Only removes tags that are no longer present in new_tags.

- Only adds/updates tags whose value has changed or that are new.

This avoids the race condition described in

https://github.com/feast-dev/feast/issues/6418 where calling

TagResource immediately after UntagResource can leave a table with no

tags due to DynamoDB's asynchronous tag operations.

"""

table_arn = dynamodb_client.describe_table(TableName=table_name)["Table"][

"TableArn"

]

current_tags = dynamodb_client.list_tags_of_resource(ResourceArn=table_arn)[

"Tags"

]

current_tag_map = {tag["Key"]: tag["Value"] for tag in current_tags}

new_tag_map = {tag["Key"]: tag["Value"] for tag in new_tags}

# Remove only tags that are no longer in new_tags

keys_to_remove = [k for k in current_tag_map if k not in new_tag_map]

# Add / update only tags whose value is new or has changed

tags_to_add = [

{"Key": k, "Value": v}

for k, v in new_tag_map.items()

if current_tag_map.get(k) != v

]

if keys_to_remove:

dynamodb_client.untag_resource(

ResourceArn=table_arn, TagKeys=keys_to_remove

)

if tags_to_add:

dynamodb_client.tag_resource(ResourceArn=table_arn, Tags=tags_to_add)

def update(

self,

config: RepoConfig,

tables_to_delete: Sequence[FeatureView],

tables_to_keep: Sequence[FeatureView],

entities_to_delete: Sequence[Entity],

entities_to_keep: Sequence[Entity],

partial: bool,

):

"""

Update tables from the DynamoDB Online Store.

Args:

config: The RepoConfig for the current FeatureStore.

tables_to_delete: Tables to delete from the DynamoDB Online Store.

tables_to_keep: Tables to keep in the DynamoDB Online Store.

"""

online_config = config.online_store

assert isinstance(online_config, DynamoDBOnlineStoreConfig)

dynamodb_client = self._get_dynamodb_client(

online_config.region,

online_config.endpoint_url,

online_config.session_based_auth,

)

dynamodb_resource = self._get_dynamodb_resource(

online_config.region,

online_config.endpoint_url,

online_config.session_based_auth,

)

do_tag_updates = defaultdict(bool)

for table_instance in tables_to_keep:

# Add Tags attribute to creation request only if configured to prevent

# TagResource permission issues, even with an empty Tags array.

table_tags = self._table_tags(online_config, table_instance)

kwargs = {"Tags": table_tags} if table_tags else {}

table_name = _get_table_name(online_config, config, table_instance)

# Check if table already exists before attempting to create

# This is required for environments where IAM roles don't have

# dynamodb:CreateTable permissions (e.g., Terraform-managed tables)

table_exists = False

try:

dynamodb_client.describe_table(TableName=table_name)

table_exists = True

do_tag_updates[table_name] = True

logger.info(

f"DynamoDB table {table_name} already exists, skipping creation"

)

except ClientError as ce:

if ce.response["Error"]["Code"] != "ResourceNotFoundException":

# If it's not a "table not found" error, re-raise

raise

# Only attempt to create table if it doesn't exist

if not table_exists:

try:

dynamodb_resource.create_table(

TableName=table_name,

KeySchema=[{"AttributeName": "entity_id", "KeyType": "HASH"}],

AttributeDefinitions=[

{"AttributeName": "entity_id", "AttributeType": "S"}

],

BillingMode="PAY_PER_REQUEST",

**kwargs,

)

logger.info(f"Created DynamoDB table {table_name}")

except ClientError as ce:

do_tag_updates[table_name] = True

# If the table creation fails with ResourceInUseException,

# it means the table already exists or is being created.

# Otherwise, re-raise the exception

if ce.response["Error"]["Code"] != "ResourceInUseException":

raise

for table_instance in tables_to_keep:

table_name = _get_table_name(online_config, config, table_instance)

dynamodb_client.get_waiter("table_exists").wait(TableName=table_name)

# once table is confirmed to exist, update the tags.

# tags won't be updated in the create_table call if the table already exists

if do_tag_updates[table_name]:

tags = self._table_tags(online_config, table_instance)

try:

self._update_tags(dynamodb_client, table_name, tags)

except ClientError as ce:

# If tag update fails with AccessDeniedException, log warning and continue

# This allows Feast to work in environments where IAM roles don't have

# dynamodb:TagResource and dynamodb:UntagResource permissions

if ce.response["Error"]["Code"] == "AccessDeniedException":

logger.warning(

f"Unable to update tags for table {table_name} due to insufficient permissions."

)

else:

raise

for table_to_delete in tables_to_delete:

_delete_table_idempotent(

dynamodb_resource,

_get_table_name(online_config, config, table_to_delete),

)

def plan(

self, config: RepoConfig, desired_registry_proto: RegistryProto

) -> List[InfraObject]:

online_config = config.online_store

assert isinstance(online_config, DynamoDBOnlineStoreConfig)

# feature_views and stream_feature_views are distinct proto types

# (FeatureViewProto vs StreamFeatureViewProto); each needs its

# matching from_proto(), not one applied to both indiscriminately.

views = [

FeatureView.from_proto(view)

for view in desired_registry_proto.feature_views

] + [

StreamFeatureView.from_proto(view)

for view in desired_registry_proto.stream_feature_views

]

return [

DynamoDBTable(

name=_get_table_name(online_config, config, view),

region=online_config.region,

endpoint_url=online_config.endpoint_url,

)

for view in views

]

def teardown(

self,

config: RepoConfig,

tables: Sequence[FeatureView],

entities: Sequence[Entity],

):

"""

Delete tables from the DynamoDB Online Store.

Args:

config: The RepoConfig for the current FeatureStore.

tables: Tables to delete from the feature repo.

"""

online_config = config.online_store

assert isinstance(online_config, DynamoDBOnlineStoreConfig)

dynamodb_resource = self._get_dynamodb_resource(

online_config.region,

online_config.endpoint_url,

online_config.session_based_auth,

)

for table in tables:

_delete_table_idempotent(

dynamodb_resource, _get_table_name(online_config, config, table)

)

def online_write_batch(

self,

config: RepoConfig,

table: FeatureView,

data: List[

Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]

],

progress: Optional[Callable[[int], Any]],

) -> None:

"""

Write a batch of feature rows to online DynamoDB store.

Note: This method applies a ``batch_writer`` to automatically handle any unprocessed items

and resend them as needed, this is useful if you're loading a lot of data at a time.

Args:

config: The RepoConfig for the current FeatureStore.

table: Feast FeatureView.

data: a list of quadruplets containing Feature data. Each quadruplet contains an Entity Key,

a dict containing feature values, an event timestamp for the row, and

the created timestamp for the row if it exists.

progress: Optional function to be called once every mini-batch of rows is written to

the online store. Can be used to display progress.

"""

online_config = config.online_store

assert isinstance(online_config, DynamoDBOnlineStoreConfig)

dynamodb_resource = self._get_dynamodb_resource(

online_config.region,

online_config.endpoint_url,

online_config.session_based_auth,

)

table_instance = dynamodb_resource.Table(

_get_table_name(online_config, config, table)

)

self._write_batch_non_duplicates(table_instance, data, progress, config)

async def online_write_batch_async(

self,

config: RepoConfig,

table: FeatureView,

data: List[

Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]

],

progress: Optional[Callable[[int], Any]],

) -> None:

"""

Writes a batch of feature rows to the online store asynchronously.

If a tz-naive timestamp is passed to this method, it is assumed to be UTC.

Args:

config: The config for the current feature store.

table: Feature view to which these feature rows correspond.

data: A list of quadruplets containing feature data. Each quadruplet contains an entity

key, a dict containing feature values, an event timestamp for the row, and the created

timestamp for the row if it exists.

progress: Function to be called once a batch of rows is written to the online store, used

to show progress.

"""

online_config = config.online_store

assert isinstance(online_config, DynamoDBOnlineStoreConfig)

table_name = _get_table_name(online_config, config, table)

items = [

_to_client_write_item(config, entity_key, features, timestamp)

for entity_key, features, timestamp, _ in _latest_data_to_write(data)

]

client = await self._get_aiodynamodb_client(

online_config.region,

online_config.max_pool_connections,

online_config.keepalive_timeout,

online_config.connect_timeout,

online_config.read_timeout,

online_config.total_max_retry_attempts,

online_config.retry_mode,

online_config.endpoint_url,

)

await dynamo_write_items_async(client, table_name, items)

def online_read(

self,

config: RepoConfig,

table: FeatureView,

entity_keys: List[EntityKeyProto],

requested_features: Optional[List[str]] = None,

) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:

"""

Retrieve feature values from the online DynamoDB store.

Args:

config: The RepoConfig for the current FeatureStore.

table: Feast FeatureView.

entity_keys: a list of entity keys that should be read from the FeatureStore.

requested_features: Optional list of feature names to retrieve.

"""

online_config = config.online_store

assert isinstance(online_config, DynamoDBOnlineStoreConfig)

dynamodb_resource = self._get_dynamodb_resource(

online_config.region,

online_config.endpoint_url,

online_config.session_based_auth,

)

table_name = _get_table_name(online_config, config, table)

batch_size = online_config.batch_size

entity_ids = self._to_entity_ids(config, entity_keys)

# Split entity_ids into batches upfront

batches: List[List[str]] = []

entity_ids_iter = iter(entity_ids)

while True:

batch = list(itertools.islice(entity_ids_iter, batch_size))

if not batch:

break

batches.append(batch)

if not batches:

return []

# For single batch, no parallelization overhead needed

if len(batches) == 1:

batch_entity_ids = self._to_resource_batch_get_payload(

online_config, table_name, batches[0], requested_features

)

response = dynamodb_resource.batch_get_item(RequestItems=batch_entity_ids)

return self._process_batch_get_response(table_name, response, batches[0])

# Execute batch requests in parallel for multiple batches

# Note: boto3 clients ARE thread-safe, so we can share a single client

# https://docs.aws.amazon.com/boto3/latest/guide/clients.html#multithreading-or-multiprocessing-with-clients

dynamodb_client = self._get_dynamodb_client(

online_config.region,

online_config.endpoint_url,

online_config.session_based_auth,

)

def fetch_batch(batch: List[str]) -> Dict[str, Any]:

batch_entity_ids = self._to_client_batch_get_payload(

online_config, table_name, batch, requested_features

)

return dynamodb_client.batch_get_item(RequestItems=batch_entity_ids)

# Use ThreadPoolExecutor for parallel I/O

max_workers = min(len(batches), online_config.max_read_workers)

with ThreadPoolExecutor(max_workers=max_workers) as executor:

responses = list(executor.map(fetch_batch, batches))

# Process responses and merge results in order

# Client responses need deserialization (unlike resource responses)

if self._type_deserializer is None:

self._type_deserializer = TypeDeserializer()

deserialize = self._type_deserializer.deserialize

def to_tbl_resp(raw_client_response):

return {

"entity_id": deserialize(raw_client_response["entity_id"]),

"event_ts": deserialize(raw_client_response["event_ts"]),

"values": deserialize(raw_client_response["values"]),

}

result: List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]] = []

for batch, response in zip(batches, responses):

batch_result = self._process_batch_get_response(

table_name, response, batch, to_tbl_response=to_tbl_resp

)

result.extend(batch_result)

return result

async def online_read_async(

self,

config: RepoConfig,

table: FeatureView,

entity_keys: List[EntityKeyProto],

requested_features: Optional[List[str]] = None,

) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:

"""

Reads features values for the given entity keys asynchronously.

Args:

config: The config for the current feature store.

table: The feature view whose feature values should be read.

entity_keys: The list of entity keys for which feature values should be read.

requested_features: The list of features that should be read.

Returns:

A list of the same length as entity_keys. Each item in the list is a tuple where the first

item is the event timestamp for the row, and the second item is a dict mapping feature names

to values, which are returned in proto format.

"""

online_config = config.online_store

assert isinstance(online_config, DynamoDBOnlineStoreConfig)

batch_size = online_config.batch_size

entity_ids = self._to_entity_ids(config, entity_keys)

entity_ids_iter = iter(entity_ids)

table_name = _get_table_name(online_config, config, table)

# Use cached TypeDeserializer for better performance

if self._type_deserializer is None:

self._type_deserializer = TypeDeserializer()

deserialize = self._type_deserializer.deserialize

def to_tbl_resp(raw_client_response):

return {

"entity_id": deserialize(raw_client_response["entity_id"]),

"event_ts": deserialize(raw_client_response["event_ts"]),

"values": deserialize(raw_client_response["values"]),

}

batches = []

entity_id_batches = []

while True:

batch = list(itertools.islice(entity_ids_iter, batch_size))

if not batch:

break

entity_id_batch = self._to_client_batch_get_payload(

online_config, table_name, batch, requested_features

)

batches.append(batch)

entity_id_batches.append(entity_id_batch)

client = await self._get_aiodynamodb_client(

online_config.region,

online_config.max_pool_connections,

online_config.keepalive_timeout,

online_config.connect_timeout,

online_config.read_timeout,

online_config.total_max_retry_attempts,

online_config.retry_mode,

online_config.endpoint_url,

)

response_batches = await asyncio.gather(

*[

client.batch_get_item(

RequestItems=entity_id_batch,

)

for entity_id_batch in entity_id_batches

]

)

result_batches = []

for batch, response in zip(batches, response_batches):

result_batch = self._process_batch_get_response(

table_name,

response,

batch,

to_tbl_response=to_tbl_resp,

)

result_batches.append(result_batch)

return list(itertools.chain(*result_batches))

def _read_features_per_fv(

self,

config: RepoConfig,

grouped_refs: List,

join_key_values: Dict,

entity_name_to_join_key_map: Dict,

online_features_response,

full_feature_names: bool,

include_feature_view_version_metadata: bool,

) -> None:

"""Read every requested feature view through shared BatchGetItem requests.

The generic path reads one feature view after another, with at least one

BatchGetItem request each. Each feature view lives in its own table here, and

a BatchGetItem request can read from several tables at once, so the keys of

every view are packed into the same requests, which are sent in parallel.

"""

if not self._should_batch(config, grouped_refs, join_key_values):

return super()._read_features_per_fv(

config,

grouped_refs,

join_key_values,

entity_name_to_join_key_map,

online_features_response,

full_feature_names,

include_feature_view_version_metadata,

)

online_config = config.online_store

assert isinstance(online_config, DynamoDBOnlineStoreConfig)

reads = self._prepare_batched_reads(

config, grouped_refs, join_key_values, entity_name_to_join_key_map

)

request_items = self._to_batched_request_items(online_config, reads)

# boto3 clients are thread-safe, so the requests can share one client.

dynamodb_client = self._get_dynamodb_client(

online_config.region,

online_config.endpoint_url,

online_config.session_based_auth,

)

def fetch_batch(batch_request_items: Dict[str, Any]) -> Dict[str, Any]:

return dynamodb_client.batch_get_item(RequestItems=batch_request_items)

if len(request_items) <= 1:

responses = [fetch_batch(items) for items in request_items]

else:

max_workers = min(len(request_items), online_config.max_read_workers)

with ThreadPoolExecutor(max_workers=max_workers) as executor:

responses = list(executor.map(fetch_batch, request_items))

self._populate_from_batched_responses(

reads,

responses,

online_features_response,

full_feature_names,

include_feature_view_version_metadata,

)

async def _read_features_per_fv_async(

self,

config: RepoConfig,

grouped_refs: List,

join_key_values: Dict,

entity_name_to_join_key_map: Dict,

online_features_response,

full_feature_names: bool,

include_feature_view_version_metadata: bool,

) -> None:

"""Async version of :meth:`_read_features_per_fv`.

The generic async path already sends the per-view requests concurrently.

Packing them still replaces at least one request per view with as few as one

in total.

"""

if not self._should_batch(config, grouped_refs, join_key_values):

return await super()._read_features_per_fv_async(

config,

grouped_refs,

join_key_values,

entity_name_to_join_key_map,

online_features_response,

full_feature_names,

include_feature_view_version_metadata,

)

online_config = config.online_store

assert isinstance(online_config, DynamoDBOnlineStoreConfig)

reads = self._prepare_batched_reads(

config, grouped_refs, join_key_values, entity_name_to_join_key_map

)

request_items = self._to_batched_request_items(online_config, reads)

client = await self._get_aiodynamodb_client(

online_config.region,

online_config.max_pool_connections,

online_config.keepalive_timeout,

online_config.connect_timeout,

online_config.read_timeout,

online_config.total_max_retry_attempts,

online_config.retry_mode,

online_config.endpoint_url,

)

responses = await asyncio.gather(

*[

client.batch_get_item(RequestItems=batch_request_items)

for batch_request_items in request_items

]

)

self._populate_from_batched_responses(

reads,

responses,

online_features_response,

full_feature_names,

include_feature_view_version_metadata,

)

@staticmethod

def _should_batch(

config: RepoConfig, grouped_refs: List, join_key_values: Dict

) -> bool:

"""Whether to read this request's feature views through shared requests.

A single view has nothing to share requests with, 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

entity_rows = max(

(len(values) for values in join_key_values.values()), default=0

)

return (

entity_rows * len(grouped_refs)

<= config.online_store.max_batched_result_rows

)

def _prepare_batched_reads(

self,

config: RepoConfig,

grouped_refs: List,

join_key_values: Dict,

entity_name_to_join_key_map: Dict,

) -> List[_BatchedRead]:

"""Resolve the table and entity ids each feature view needs before reading.

Feature views can be keyed on different entities, so the entity ids are

computed per view rather than shared.

"""

online_config = config.online_store

reads = []

for table, requested_features in grouped_refs:

table_entity_values, idxs, output_len = utils._get_unique_entities(

table, join_key_values, entity_name_to_join_key_map

)

entity_key_protos = utils._get_entity_key_protos(table_entity_values)

reads.append(

_BatchedRead(

table=table,

requested_features=requested_features,

table_name=_get_table_name(online_config, config, table),

entity_ids=self._to_entity_ids(config, entity_key_protos),

idxs=idxs,

output_len=output_len,

)

)

return reads

@staticmethod

def _to_batched_request_items(

online_config: DynamoDBOnlineStoreConfig, reads: List[_BatchedRead]

) -> List[Dict[str, Any]]:

"""Pack the keys of every feature view into as few requests as possible.

``batch_size`` caps the keys of one BatchGetItem request across all of its

tables. A request must not repeat a key within a table, so views that share a

table, such as one feature view requested under two aliases, are read once,

projecting the union of their requested features.

"""

features_by_table: Dict[str, Optional[List[str]]] = {}

entity_ids_by_table: Dict[str, Dict[str, None]] = {}

for read in reads:

if read.table_name not in features_by_table:

features_by_table[read.table_name] = read.requested_features or None

entity_ids_by_table[read.table_name] = {}

else:

features = features_by_table[read.table_name]

features_by_table[read.table_name] = (

list(dict.fromkeys(features + read.requested_features))

if features and read.requested_features

else None

)

entity_ids_by_table[read.table_name].update(dict.fromkeys(read.entity_ids))

keys = [

(table_name, entity_id)

for table_name, entity_ids in entity_ids_by_table.items()

for entity_id in entity_ids

]

batch_size = online_config.batch_size

request_items: List[Dict[str, Any]] = []

for start in range(0, len(keys), batch_size):

batch_by_table: Dict[str, List[str]] = defaultdict(list)

for table_name, entity_id in keys[start : start + batch_size]:

batch_by_table[table_name].append(entity_id)

batch_request_items: Dict[str, Any] = {}

for table_name, batch in batch_by_table.items():

batch_request_items.update(

DynamoDBOnlineStore._to_client_batch_get_payload(

online_config, table_name, batch, features_by_table[table_name]

)

)

request_items.append(batch_request_items)

return request_items

def _populate_from_batched_responses(

self,

reads: List[_BatchedRead],

responses: Sequence[Dict[str, Any]],

online_features_response,

full_feature_names: bool,

include_feature_view_version_metadata: bool,

) -> None:

"""Split the responses by table and hand each view its rows in key order.

BatchGetItem returns items in no particular order, and a request can mix

tables, so the items are indexed by table and entity id first.

"""

if self._type_deserializer is None:

self._type_deserializer = TypeDeserializer()

deserialize = self._type_deserializer.deserialize

rows_by_table: Dict[

str, Dict[str, Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]

] = defaultdict(dict)

for response in responses:

for table_name, items in response.get("Responses", {}).items():

table_rows = rows_by_table[table_name]

for item in items:

table_rows[deserialize(item["entity_id"])] = self._to_read_row(