|
1 | 1 | import asyncio |
2 | | -from datetime import datetime, timezone |
| 2 | +from datetime import datetime, timedelta, timezone |
3 | 3 | from unittest.mock import AsyncMock, MagicMock, patch |
4 | 4 |
|
5 | 5 | import pytest |
6 | 6 | from google.protobuf.timestamp_pb2 import Timestamp |
7 | 7 |
|
8 | 8 | from feast import Entity, FeatureView, Field, FileSource, RepoConfig |
| 9 | +from feast.infra.online_stores.helpers import _mmh3 |
9 | 10 | from feast.infra.online_stores.redis import RedisOnlineStore, RedisOnlineStoreConfig |
10 | 11 | from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto |
11 | 12 | from feast.protos.feast.types.Value_pb2 import Value as ValueProto |
@@ -426,6 +427,145 @@ def test_online_write_batch_with_dedup_uses_two_pipelines( |
426 | 427 | pipe.hmget.assert_called_once() |
427 | 428 |
|
428 | 429 |
|
| 430 | +def _dedup_config(): |
| 431 | + return RepoConfig( |
| 432 | + provider="local", |
| 433 | + project="test", |
| 434 | + entity_key_serialization_version=3, |
| 435 | + registry="dummy_registry.db", |
| 436 | + online_store=RedisOnlineStoreConfig(), # default: skip_dedup=False |
| 437 | + ) |
| 438 | + |
| 439 | + |
| 440 | +def _single_entity_batch(timestamps_and_values): |
| 441 | + """Build a write batch for one entity key, one row per (timestamp, value) pair.""" |
| 442 | + entity_key = EntityKeyProto( |
| 443 | + join_keys=["entity"], entity_values=[ValueProto(int32_val=1)] |
| 444 | + ) |
| 445 | + return [ |
| 446 | + (entity_key, {"feature_10": ValueProto(int32_val=value)}, timestamp, None) |
| 447 | + for timestamp, value in timestamps_and_values |
| 448 | + ] |
| 449 | + |
| 450 | + |
| 451 | +def _written_feature_values(hset_calls, feature_view_name, feature_name): |
| 452 | + """Extract the serialized feature value from each queued hset call.""" |
| 453 | + f_key = _mmh3(f"{feature_view_name}:{feature_name}") |
| 454 | + return [call.kwargs["mapping"][f_key] for call in hset_calls] |
| 455 | + |
| 456 | + |
| 457 | +@pytest.mark.parametrize( |
| 458 | + "order", |
| 459 | + ["descending", "ascending", "unordered"], |
| 460 | + ids=["descending", "ascending", "unordered"], |
| 461 | +) |
| 462 | +def test_online_write_batch_keeps_latest_event_time_within_batch( |
| 463 | + redis_online_store: RedisOnlineStore, feature_view, order |
| 464 | +): |
| 465 | + """A batch containing several rows for one entity key must leave the value |
| 466 | + belonging to the latest event timestamp in the store, whatever order the rows |
| 467 | + arrive in. |
| 468 | +
|
| 469 | + Regression test for #5163: previous timestamps were read once up front, so rows |
| 470 | + sharing an entity key all compared against the same pre-batch snapshot and could |
| 471 | + not see each other. Every row passed the staleness guard and the last row queued |
| 472 | + won, so a reverse-chronological batch stored the *oldest* value. |
| 473 | + """ |
| 474 | + t1 = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc) |
| 475 | + t2 = t1 + timedelta(seconds=5) |
| 476 | + t3 = t2 + timedelta(seconds=5) |
| 477 | + |
| 478 | + rows = { |
| 479 | + "descending": [(t3, 30), (t2, 20), (t1, 10)], |
| 480 | + "ascending": [(t1, 10), (t2, 20), (t3, 30)], |
| 481 | + "unordered": [(t2, 20), (t3, 30), (t1, 10)], |
| 482 | + }[order] |
| 483 | + data = _single_entity_batch(rows) |
| 484 | + |
| 485 | + mock_client = MagicMock() |
| 486 | + pipe = MagicMock() |
| 487 | + pipe.__enter__ = MagicMock(return_value=pipe) |
| 488 | + pipe.__exit__ = MagicMock(return_value=False) |
| 489 | + # No pre-existing value: one hmget result per row, each holding None. |
| 490 | + pipe.execute.side_effect = [[[None]] * len(data), []] |
| 491 | + mock_client.pipeline.return_value = pipe |
| 492 | + |
| 493 | + with patch.object(redis_online_store, "_get_client", return_value=mock_client): |
| 494 | + redis_online_store.online_write_batch( |
| 495 | + _dedup_config(), feature_view, data, progress=None |
| 496 | + ) |
| 497 | + |
| 498 | + written = _written_feature_values( |
| 499 | + pipe.hset.call_args_list, feature_view.name, "feature_10" |
| 500 | + ) |
| 501 | + assert written, "no write was queued for the batch" |
| 502 | + # Redis applies queued commands in order, so the surviving value is the last one. |
| 503 | + assert written[-1] == ValueProto(int32_val=30).SerializeToString() |
| 504 | + |
| 505 | + |
| 506 | +def test_online_write_batch_async_keeps_latest_event_time_within_batch( |
| 507 | + redis_online_store: RedisOnlineStore, feature_view |
| 508 | +): |
| 509 | + """online_write_batch_async must honour the same intra-batch ordering guarantee |
| 510 | + as the sync path (#5163).""" |
| 511 | + t1 = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc) |
| 512 | + t2 = t1 + timedelta(seconds=5) |
| 513 | + t3 = t2 + timedelta(seconds=5) |
| 514 | + data = _single_entity_batch([(t3, 30), (t2, 20), (t1, 10)]) |
| 515 | + |
| 516 | + async_pipe = AsyncMock() |
| 517 | + async_pipe.__aenter__ = AsyncMock(return_value=async_pipe) |
| 518 | + async_pipe.__aexit__ = AsyncMock(return_value=False) |
| 519 | + async_pipe.execute = AsyncMock(side_effect=[[[None]] * len(data), []]) |
| 520 | + async_pipe.hset = MagicMock() |
| 521 | + |
| 522 | + mock_async_client = AsyncMock() |
| 523 | + mock_async_client.pipeline = MagicMock(return_value=async_pipe) |
| 524 | + |
| 525 | + async def _run(): |
| 526 | + with patch.object( |
| 527 | + redis_online_store, |
| 528 | + "_get_client_async", |
| 529 | + AsyncMock(return_value=mock_async_client), |
| 530 | + ): |
| 531 | + await redis_online_store.online_write_batch_async( |
| 532 | + _dedup_config(), feature_view, data, progress=None |
| 533 | + ) |
| 534 | + |
| 535 | + asyncio.run(_run()) |
| 536 | + |
| 537 | + written = _written_feature_values( |
| 538 | + async_pipe.hset.call_args_list, feature_view.name, "feature_10" |
| 539 | + ) |
| 540 | + assert written, "no write was queued for the batch" |
| 541 | + assert written[-1] == ValueProto(int32_val=30).SerializeToString() |
| 542 | + |
| 543 | + |
| 544 | +def test_online_write_batch_still_skips_rows_older_than_stored_value( |
| 545 | + redis_online_store: RedisOnlineStore, feature_view |
| 546 | +): |
| 547 | + """The pre-existing staleness guard must keep working: a row older than the value |
| 548 | + already in Redis is dropped rather than written.""" |
| 549 | + stored_ts = Timestamp() |
| 550 | + stored_ts.FromDatetime(datetime(2024, 1, 1, 12, 0, 10, tzinfo=timezone.utc)) |
| 551 | + older = datetime(2024, 1, 1, 12, 0, 0, tzinfo=timezone.utc) |
| 552 | + data = _single_entity_batch([(older, 10)]) |
| 553 | + |
| 554 | + mock_client = MagicMock() |
| 555 | + pipe = MagicMock() |
| 556 | + pipe.__enter__ = MagicMock(return_value=pipe) |
| 557 | + pipe.__exit__ = MagicMock(return_value=False) |
| 558 | + pipe.execute.side_effect = [[[stored_ts.SerializeToString()]], []] |
| 559 | + mock_client.pipeline.return_value = pipe |
| 560 | + |
| 561 | + with patch.object(redis_online_store, "_get_client", return_value=mock_client): |
| 562 | + redis_online_store.online_write_batch( |
| 563 | + _dedup_config(), feature_view, data, progress=None |
| 564 | + ) |
| 565 | + |
| 566 | + pipe.hset.assert_not_called() |
| 567 | + |
| 568 | + |
429 | 569 | def test_online_write_batch_async_skip_dedup_single_pipeline( |
430 | 570 | redis_online_store: RedisOnlineStore, feature_view |
431 | 571 | ): |
|
0 commit comments