Skip to content

Commit 2eb0cd7

Browse files
Jwredentkathole
Jwrede
authored and
ntkathole
committed
feat(dynamodb): Use ProjectionExpression when requested_features is set
The requested_features parameter was accepted by online_read and online_read_async but never used -- DynamoDB always fetched all features stored in the values map regardless. Add a ProjectionExpression to BatchGetItem requests when requested_features is provided, reducing data transfer, latency, and read costs. Fixes feast-dev#6058 Signed-off-by: Jonathan Wrede <[email protected]>
1 parent 4541690 commit 2eb0cd7

2 files changed

Lines changed: 153 additions & 14 deletions

File tree

‎sdk/python/feast/infra/online_stores/dynamodb.py‎

Lines changed: 49 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -504,7 +504,7 @@ def online_read(
504504
# For single batch, no parallelization overhead needed
505505
if len(batches) == 1:
506506
batch_entity_ids = self._to_resource_batch_get_payload(
507-
online_config, table_name, batches[0]
507+
online_config, table_name, batches[0], requested_features
508508
)
509509
response = dynamodb_resource.batch_get_item(RequestItems=batch_entity_ids)
510510
return self._process_batch_get_response(table_name, response, batches[0])
@@ -520,7 +520,7 @@ def online_read(
520520

521521
def fetch_batch(batch: List[str]) -> Dict[str, Any]:
522522
batch_entity_ids = self._to_client_batch_get_payload(
523-
online_config, table_name, batch
523+
online_config, table_name, batch, requested_features
524524
)
525525
return dynamodb_client.batch_get_item(RequestItems=batch_entity_ids)
526526

@@ -599,7 +599,7 @@ def to_tbl_resp(raw_client_response):
599599
if not batch:
600600
break
601601
entity_id_batch = self._to_client_batch_get_payload(
602-
online_config, table_name, batch
602+
online_config, table_name, batch, requested_features
603603
)
604604
batches.append(batch)
605605
entity_id_batches.append(entity_id_batch)
@@ -760,21 +760,56 @@ def _to_entity_ids(config: RepoConfig, entity_keys: List[EntityKeyProto]):
760760
]
761761

762762
@staticmethod
763-
def _to_resource_batch_get_payload(online_config, table_name, batch):
764-
return {
765-
table_name: {
766-
"Keys": [{"entity_id": entity_id} for entity_id in batch],
767-
"ConsistentRead": online_config.consistent_reads,
768-
}
763+
def _to_resource_batch_get_payload(
764+
online_config, table_name, batch, requested_features=None
765+
):
766+
payload: Dict[str, Any] = {
767+
"Keys": [{"entity_id": entity_id} for entity_id in batch],
768+
"ConsistentRead": online_config.consistent_reads,
769+
}
770+
projection = DynamoDBOnlineStore._build_projection_expression(
771+
requested_features
772+
)
773+
if projection:
774+
payload["ProjectionExpression"] = projection["ProjectionExpression"]
775+
payload["ExpressionAttributeNames"] = projection["ExpressionAttributeNames"]
776+
return {table_name: payload}
777+
778+
@staticmethod
779+
def _to_client_batch_get_payload(
780+
online_config, table_name, batch, requested_features=None
781+
):
782+
payload: Dict[str, Any] = {
783+
"Keys": [{"entity_id": {"S": entity_id}} for entity_id in batch],
784+
"ConsistentRead": online_config.consistent_reads,
769785
}
786+
projection = DynamoDBOnlineStore._build_projection_expression(
787+
requested_features
788+
)
789+
if projection:
790+
payload["ProjectionExpression"] = projection["ProjectionExpression"]
791+
payload["ExpressionAttributeNames"] = projection["ExpressionAttributeNames"]
792+
return {table_name: payload}
770793

771794
@staticmethod
772-
def _to_client_batch_get_payload(online_config, table_name, batch):
795+
def _build_projection_expression(
796+
requested_features: Optional[List[str]],
797+
) -> Optional[Dict[str, str]]:
798+
if not requested_features:
799+
return None
800+
attr_names: Dict[str, str] = {
801+
"#entity_id": "entity_id",
802+
"#event_ts": "event_ts",
803+
"#vals": "values",
804+
}
805+
projections = ["#entity_id", "#event_ts"]
806+
for i, feat in enumerate(requested_features):
807+
alias = f"#feat{i}"
808+
attr_names[alias] = feat
809+
projections.append(f"#vals.{alias}")
773810
return {
774-
table_name: {
775-
"Keys": [{"entity_id": {"S": entity_id}} for entity_id in batch],
776-
"ConsistentRead": online_config.consistent_reads,
777-
}
811+
"ProjectionExpression": ", ".join(projections),
812+
"ExpressionAttributeNames": attr_names,
778813
}
779814

780815
def update_online_store(

‎sdk/python/tests/unit/infra/online_store/test_dynamodb_online_store.py‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1050,3 +1050,107 @@ def tracking_client(*args, **kwargs):
10501050
f"Expected 1 shared client for thread-safety, "
10511051
f"got {len(set(dynamodb_clients))} unique clients"
10521052
)
1053+
1054+
1055+
@mock_dynamodb
1056+
def test_dynamodb_online_store_online_read_with_requested_features(
1057+
repo_config, dynamodb_online_store
1058+
):
1059+
"""Test that requested_features filters returned features."""
1060+
n_samples = 5
1061+
db_table_name = f"{TABLE_NAME}_requested_features"
1062+
create_test_table(PROJECT, db_table_name, REGION)
1063+
data = create_n_customer_test_samples(n=n_samples)
1064+
insert_data_test_table(data, PROJECT, db_table_name, REGION)
1065+
1066+
entity_keys, features, *rest = zip(*data)
1067+
returned_items = dynamodb_online_store.online_read(
1068+
config=repo_config,
1069+
table=MockFeatureView(name=db_table_name),
1070+
entity_keys=entity_keys,
1071+
requested_features=["name", "age"],
1072+
)
1073+
assert len(returned_items) == n_samples
1074+
for _, feat_dict in returned_items:
1075+
assert feat_dict is not None
1076+
assert "name" in feat_dict
1077+
assert "age" in feat_dict
1078+
assert "avg_orders_day" not in feat_dict
1079+
1080+
1081+
@mock_dynamodb
1082+
def test_dynamodb_online_store_online_read_without_requested_features(
1083+
repo_config, dynamodb_online_store
1084+
):
1085+
"""Test that omitting requested_features returns all features."""
1086+
n_samples = 5
1087+
db_table_name = f"{TABLE_NAME}_all_features"
1088+
create_test_table(PROJECT, db_table_name, REGION)
1089+
data = create_n_customer_test_samples(n=n_samples)
1090+
insert_data_test_table(data, PROJECT, db_table_name, REGION)
1091+
1092+
entity_keys, features, *rest = zip(*data)
1093+
returned_items = dynamodb_online_store.online_read(
1094+
config=repo_config,
1095+
table=MockFeatureView(name=db_table_name),
1096+
entity_keys=entity_keys,
1097+
requested_features=None,
1098+
)
1099+
assert len(returned_items) == n_samples
1100+
for _, feat_dict in returned_items:
1101+
assert feat_dict is not None
1102+
assert set(feat_dict.keys()) == {"avg_orders_day", "name", "age"}
1103+
1104+
1105+
def test_build_projection_expression():
1106+
"""Test that _build_projection_expression generates correct DynamoDB expressions."""
1107+
result = DynamoDBOnlineStore._build_projection_expression(["feat_a", "feat_b"])
1108+
assert result is not None
1109+
assert "#entity_id" in result["ProjectionExpression"]
1110+
assert "#event_ts" in result["ProjectionExpression"]
1111+
assert "#vals.#feat0" in result["ProjectionExpression"]
1112+
assert "#vals.#feat1" in result["ProjectionExpression"]
1113+
attr_names = result["ExpressionAttributeNames"]
1114+
assert attr_names["#vals"] == "values"
1115+
assert attr_names["#feat0"] == "feat_a"
1116+
assert attr_names["#feat1"] == "feat_b"
1117+
1118+
1119+
def test_build_projection_expression_none():
1120+
"""Test that _build_projection_expression returns None for empty input."""
1121+
assert DynamoDBOnlineStore._build_projection_expression(None) is None
1122+
assert DynamoDBOnlineStore._build_projection_expression([]) is None
1123+
1124+
1125+
@mock_dynamodb
1126+
def test_dynamodb_online_store_online_read_requested_features_parallel(
1127+
dynamodb_online_store,
1128+
):
1129+
"""Test that requested_features works across parallel batches."""
1130+
small_batch_config = RepoConfig(
1131+
registry=REGISTRY,
1132+
project=PROJECT,
1133+
provider=PROVIDER,
1134+
online_store=DynamoDBOnlineStoreConfig(region=REGION, batch_size=5),
1135+
offline_store=DaskOfflineStoreConfig(),
1136+
entity_key_serialization_version=3,
1137+
)
1138+
n_samples = 15
1139+
db_table_name = f"{TABLE_NAME}_requested_parallel"
1140+
create_test_table(PROJECT, db_table_name, REGION)
1141+
data = create_n_customer_test_samples(n=n_samples)
1142+
insert_data_test_table(data, PROJECT, db_table_name, REGION)
1143+
1144+
entity_keys, features, *rest = zip(*data)
1145+
returned_items = dynamodb_online_store.online_read(
1146+
config=small_batch_config,
1147+
table=MockFeatureView(name=db_table_name),
1148+
entity_keys=entity_keys,
1149+
requested_features=["age"],
1150+
)
1151+
assert len(returned_items) == n_samples
1152+
for _, feat_dict in returned_items:
1153+
assert feat_dict is not None
1154+
assert "age" in feat_dict
1155+
assert "name" not in feat_dict
1156+
assert "avg_orders_day" not in feat_dict

0 commit comments

Comments
 (0)