diff --git a/sdk/python/feast/infra/key_encoding_utils.py b/sdk/python/feast/infra/key_encoding_utils.py index 3e9ba70d3ba..10a9934ad6a 100644 --- a/sdk/python/feast/infra/key_encoding_utils.py +++ b/sdk/python/feast/infra/key_encoding_utils.py @@ -57,15 +57,20 @@ def serialize_entity_key_prefix( This encoding is a partial implementation of serialize_entity_key, only operating on the keys of entities, and not the values. """ - sorted_keys = sorted(entity_keys) + # Fast path optimization for single entity + if len(entity_keys) == 1: + sorted_keys = [entity_keys[0]] + else: + sorted_keys = sorted(entity_keys) output: List[bytes] = [] if entity_key_serialization_version > 2: output.append(struct.pack(" 2: - output.append(struct.pack(" 2: output.append(struct.pack(" 2: - output.append(struct.pack(" 2: + output.append(struct.pack(" 0 + + +@pytest.mark.benchmark(group="serialize_single") +@pytest.mark.parametrize("entity_key_serialization_version", [3]) +def test_serialize_single_entity_int( + benchmark, single_entity_key_int, entity_key_serialization_version +): + """Benchmark single entity key serialization (int64 value).""" + result = benchmark( + serialize_entity_key, single_entity_key_int, entity_key_serialization_version + ) + assert len(result) > 0 + + +@pytest.mark.benchmark(group="serialize_multi") +@pytest.mark.parametrize("entity_key_serialization_version", [3]) +def test_serialize_multi_entity_small( + benchmark, multi_entity_key_small, entity_key_serialization_version +): + """Benchmark small multi-entity key serialization.""" + result = benchmark( + serialize_entity_key, multi_entity_key_small, entity_key_serialization_version + ) + assert len(result) > 0 + + +@pytest.mark.benchmark(group="serialize_multi") +@pytest.mark.parametrize("entity_key_serialization_version", [3]) +def test_serialize_multi_entity_large( + benchmark, multi_entity_key_large, entity_key_serialization_version +): + """Benchmark large multi-entity key serialization.""" + result = benchmark( + serialize_entity_key, multi_entity_key_large, entity_key_serialization_version + ) + assert len(result) > 0 + + +@pytest.mark.benchmark(group="serialize_mixed") +@pytest.mark.parametrize("entity_key_serialization_version", [3]) +def test_serialize_mixed_value_types( + benchmark, mixed_value_types_key, entity_key_serialization_version +): + """Benchmark serialization with mixed value types.""" + result = benchmark( + serialize_entity_key, mixed_value_types_key, entity_key_serialization_version + ) + assert len(result) > 0 + + +# Deserialization Benchmarks + + +@pytest.mark.benchmark(group="deserialize_single") +@pytest.mark.parametrize("entity_key_serialization_version", [3]) +def test_deserialize_single_entity_string( + benchmark, single_entity_key, entity_key_serialization_version +): + """Benchmark single entity key deserialization (string value).""" + serialized = serialize_entity_key( + single_entity_key, entity_key_serialization_version + ) + result = benchmark( + deserialize_entity_key, serialized, entity_key_serialization_version + ) + assert result == single_entity_key + + +@pytest.mark.benchmark(group="deserialize_single") +@pytest.mark.parametrize("entity_key_serialization_version", [3]) +def test_deserialize_single_entity_int( + benchmark, single_entity_key_int, entity_key_serialization_version +): + """Benchmark single entity key deserialization (int64 value).""" + serialized = serialize_entity_key( + single_entity_key_int, entity_key_serialization_version + ) + result = benchmark( + deserialize_entity_key, serialized, entity_key_serialization_version + ) + assert result == single_entity_key_int + + +@pytest.mark.benchmark(group="deserialize_multi") +@pytest.mark.parametrize("entity_key_serialization_version", [3]) +def test_deserialize_multi_entity_small( + benchmark, multi_entity_key_small, entity_key_serialization_version +): + """Benchmark small multi-entity key deserialization.""" + serialized = serialize_entity_key( + multi_entity_key_small, entity_key_serialization_version + ) + result = benchmark( + deserialize_entity_key, serialized, entity_key_serialization_version + ) + assert result == multi_entity_key_small + + +@pytest.mark.benchmark(group="deserialize_multi") +@pytest.mark.parametrize("entity_key_serialization_version", [3]) +def test_deserialize_multi_entity_large( + benchmark, multi_entity_key_large, entity_key_serialization_version +): + """Benchmark large multi-entity key deserialization.""" + serialized = serialize_entity_key( + multi_entity_key_large, entity_key_serialization_version + ) + result = benchmark( + deserialize_entity_key, serialized, entity_key_serialization_version + ) + assert result == multi_entity_key_large + + +@pytest.mark.benchmark(group="deserialize_mixed") +@pytest.mark.parametrize("entity_key_serialization_version", [3]) +def test_deserialize_mixed_value_types( + benchmark, mixed_value_types_key, entity_key_serialization_version +): + """Benchmark deserialization with mixed value types.""" + serialized = serialize_entity_key( + mixed_value_types_key, entity_key_serialization_version + ) + result = benchmark( + deserialize_entity_key, serialized, entity_key_serialization_version + ) + assert result == mixed_value_types_key + + +# Round-trip Benchmarks + + +@pytest.mark.benchmark(group="roundtrip_single") +def test_roundtrip_single_entity(benchmark, single_entity_key): + """Benchmark complete serialize + deserialize round-trip for single entity.""" + + def roundtrip(): + serialized = serialize_entity_key(single_entity_key, 3) + return deserialize_entity_key(serialized, 3) + + result = benchmark(roundtrip) + assert result == single_entity_key + + +@pytest.mark.benchmark(group="roundtrip_multi") +def test_roundtrip_multi_entity(benchmark, multi_entity_key_small): + """Benchmark complete serialize + deserialize round-trip for multi-entity.""" + + def roundtrip(): + serialized = serialize_entity_key(multi_entity_key_small, 3) + return deserialize_entity_key(serialized, 3) + + result = benchmark(roundtrip) + assert result == multi_entity_key_small + + +# Prefix Serialization Benchmarks + + +@pytest.mark.benchmark(group="prefix") +def test_serialize_entity_key_prefix_single(benchmark): + """Benchmark entity key prefix serialization for single key.""" + result = benchmark(serialize_entity_key_prefix, ["user_id"], 3) + assert len(result) > 0 + + +@pytest.mark.benchmark(group="prefix") +def test_serialize_entity_key_prefix_multi(benchmark): + """Benchmark entity key prefix serialization for multiple keys.""" + keys = ["user_id", "session_id", "device_id"] + result = benchmark(serialize_entity_key_prefix, keys, 3) + assert len(result) > 0 + + +# Bulk Operations Benchmarks + + +@pytest.mark.benchmark(group="bulk_serialize") +def test_bulk_serialize_batch(benchmark, batch_entity_keys): + """Benchmark batch serialization of 100 mixed entity keys.""" + + def bulk_serialize(): + results = [] + for entity_key in batch_entity_keys: + serialized = serialize_entity_key(entity_key, 3) + results.append(serialized) + return results + + results = benchmark(bulk_serialize) + assert len(results) == 100 + + +@pytest.mark.benchmark(group="bulk_deserialize") +def test_bulk_deserialize_batch(benchmark, batch_entity_keys): + """Benchmark batch deserialization of 100 mixed entity keys.""" + # Pre-serialize all keys + serialized_keys = [serialize_entity_key(key, 3) for key in batch_entity_keys] + + def bulk_deserialize(): + results = [] + for serialized in serialized_keys: + deserialized = deserialize_entity_key(serialized, 3) + results.append(deserialized) + return results + + results = benchmark(bulk_deserialize) + assert len(results) == 100 + + +@pytest.mark.benchmark(group="bulk_roundtrip") +def test_bulk_roundtrip_batch(benchmark, batch_entity_keys): + """Benchmark bulk serialize + deserialize for realistic workload.""" + + def bulk_roundtrip(): + results = [] + for entity_key in batch_entity_keys: + serialized = serialize_entity_key(entity_key, 3) + deserialized = deserialize_entity_key(serialized, 3) + results.append(deserialized) + return results + + results = benchmark(bulk_roundtrip) + assert len(results) == 100 + + +# Memory Efficiency Tests + + +def test_memory_efficiency_serialization(single_entity_key): + """Test memory usage during serialization (not a benchmark, just validation).""" + import os + + import psutil + + process = psutil.Process(os.getpid()) + initial_memory = process.memory_info().rss + + # Perform many serializations + for i in range(10000): + entity_key = EntityKeyProto( + join_keys=["user_id"], entity_values=[ValueProto(string_val=f"user{i}")] + ) + serialize_entity_key(entity_key, 3) + + final_memory = process.memory_info().rss + memory_increase = final_memory - initial_memory + + # Memory increase should be minimal (< 10MB for 10k operations) + # This validates that we're not leaking memory in the optimized version + assert memory_increase < 10 * 1024 * 1024, ( + f"Memory usage increased by {memory_increase / 1024 / 1024:.2f} MB" + ) + + +# Performance Regression Tests + + +def test_performance_regression_single_entity(): + """Regression test: single entity serialization should be faster than baseline.""" + entity_key = EntityKeyProto( + join_keys=["user_id"], entity_values=[ValueProto(string_val="user123")] + ) + + # Warm up + for _ in range(100): + serialize_entity_key(entity_key, 3) + + # Time 1000 operations + start_time = time.perf_counter() + for _ in range(1000): + serialize_entity_key(entity_key, 3) + elapsed = time.perf_counter() - start_time + + # Should be able to do 1000 single entity serializations in < 10ms + # This is a conservative regression test + assert elapsed < 0.01, ( + f"Single entity serialization too slow: {elapsed:.4f}s for 1000 operations" + ) + + +def test_performance_regression_deserialization(): + """Regression test: deserialization should be fast with memoryview optimization.""" + entity_key = EntityKeyProto( + join_keys=["user_id", "session_id"], + entity_values=[ + ValueProto(string_val="user123"), + ValueProto(string_val="sess456"), + ], + ) + + serialized = serialize_entity_key(entity_key, 3) + + # Warm up + for _ in range(100): + deserialize_entity_key(serialized, 3) + + # Time 1000 operations + start_time = time.perf_counter() + for _ in range(1000): + deserialize_entity_key(serialized, 3) + elapsed = time.perf_counter() - start_time + + # Should be able to do 1000 deserializations in < 15ms + assert elapsed < 0.015, ( + f"Deserialization too slow: {elapsed:.4f}s for 1000 operations" + ) + + +# Binary Compatibility Tests + + +def test_binary_format_consistency_single(): + """Ensure optimizations don't change binary format for single entities.""" + entity_key = EntityKeyProto( + join_keys=["user_id"], entity_values=[ValueProto(string_val="test")] + ) + + # Serialize multiple times - results should be identical + results = [] + for _ in range(10): + serialized = serialize_entity_key(entity_key, 3) + results.append(serialized) + + # All results should be identical + for result in results[1:]: + assert result == results[0], "Binary format inconsistency detected" + + +def test_binary_format_consistency_multi(): + """Ensure optimizations don't change binary format for multi-entity keys.""" + entity_key = EntityKeyProto( + join_keys=["user", "session", "device"], + entity_values=[ + ValueProto(string_val="u1"), + ValueProto(string_val="s1"), + ValueProto(string_val="d1"), + ], + ) + + # Serialize multiple times - results should be identical + results = [] + for _ in range(10): + serialized = serialize_entity_key(entity_key, 3) + results.append(serialized) + + # All results should be identical + for result in results[1:]: + assert result == results[0], "Binary format inconsistency detected" diff --git a/sdk/python/tests/unit/infra/test_key_encoding_utils.py b/sdk/python/tests/unit/infra/test_key_encoding_utils.py index 14433a41e65..b6c8a024e2d 100644 --- a/sdk/python/tests/unit/infra/test_key_encoding_utils.py +++ b/sdk/python/tests/unit/infra/test_key_encoding_utils.py @@ -4,6 +4,7 @@ deserialize_entity_key, reserialize_entity_v2_key_to_v3, serialize_entity_key, + serialize_entity_key_prefix, ) from feast.protos.feast.types.EntityKey_pb2 import EntityKey as EntityKeyProto from feast.protos.feast.types.Value_pb2 import Value as ValueProto @@ -151,3 +152,187 @@ def test_reserialize_entity_v2_key_to_v3(): join_keys=["user"], entity_values=[ValueProto(int64_val=int(2**15))], ) + + +def test_single_entity_fast_path(): + """Test that single entity optimization works correctly.""" + entity_key_proto = EntityKeyProto( + join_keys=["user_id"], + entity_values=[ValueProto(string_val="test_user")], + ) + + serialized_key = serialize_entity_key( + entity_key_proto, entity_key_serialization_version=3 + ) + deserialized_key = deserialize_entity_key( + serialized_key, entity_key_serialization_version=3 + ) + + assert deserialized_key == entity_key_proto + + +def test_empty_entity_key(): + """Test handling of empty entity keys.""" + entity_key_proto = EntityKeyProto(join_keys=[], entity_values=[]) + + serialized_key = serialize_entity_key( + entity_key_proto, entity_key_serialization_version=3 + ) + deserialized_key = deserialize_entity_key( + serialized_key, entity_key_serialization_version=3 + ) + + assert deserialized_key == entity_key_proto + + +def test_binary_format_deterministic(): + """Test that serialization is deterministic (same input produces same output).""" + entity_key_proto = EntityKeyProto( + join_keys=["customer", "user", "session"], + entity_values=[ + ValueProto(string_val="cust1"), + ValueProto(string_val="user1"), + ValueProto(string_val="sess1"), + ], + ) + + # Serialize the same entity multiple times + serializations = [] + for _ in range(5): + serialized = serialize_entity_key( + entity_key_proto, entity_key_serialization_version=3 + ) + serializations.append(serialized) + + # All serializations should be identical + for s in serializations[1:]: + assert s == serializations[0], "Serialization is not deterministic" + + +def test_optimization_preserves_sorting(): + """Test that optimizations preserve the sorting behavior for multi-entity keys.""" + # Create entity key with unsorted keys + entity_key_proto = EntityKeyProto( + join_keys=["zebra", "alpha", "beta"], + entity_values=[ + ValueProto(string_val="z_val"), + ValueProto(string_val="a_val"), + ValueProto(string_val="b_val"), + ], + ) + + serialized = serialize_entity_key( + entity_key_proto, entity_key_serialization_version=3 + ) + deserialized = deserialize_entity_key( + serialized, entity_key_serialization_version=3 + ) + + # Keys should be sorted in the result + expected_sorted_keys = ["alpha", "beta", "zebra"] + expected_sorted_values = ["a_val", "b_val", "z_val"] + + assert deserialized.join_keys == expected_sorted_keys + assert [v.string_val for v in deserialized.entity_values] == expected_sorted_values + + +def test_performance_bounds_single_entity(): + """Regression test to ensure single entity performance meets minimum bounds.""" + import time + + entity_key = EntityKeyProto( + join_keys=["user_id"], entity_values=[ValueProto(string_val="user123")] + ) + + # Measure serialization time for 1000 operations + start = time.perf_counter() + for _ in range(1000): + serialize_entity_key(entity_key, entity_key_serialization_version=3) + serialize_time = time.perf_counter() - start + + # Measure deserialization time + serialized = serialize_entity_key(entity_key, entity_key_serialization_version=3) + start = time.perf_counter() + for _ in range(1000): + deserialize_entity_key(serialized, entity_key_serialization_version=3) + deserialize_time = time.perf_counter() - start + + # Conservative performance bounds (should be much faster with optimizations) + # 1000 operations should complete in < 20ms each for serialization and deserialization + assert serialize_time < 0.02, f"Serialization too slow: {serialize_time:.4f}s" + assert deserialize_time < 0.02, f"Deserialization too slow: {deserialize_time:.4f}s" + + +def test_non_ascii_prefix_compatibility(): + """Critical test: ensure prefix serialization matches full entity key serialization for non-ASCII keys.""" + # Test with non-ASCII characters that have different byte vs character lengths + non_ascii_keys = ["用户ID", "사용자ID", "идентификатор", "مُعرِّف"] + + for key in non_ascii_keys: + # Test single key prefix + prefix_result = serialize_entity_key_prefix( + [key], entity_key_serialization_version=3 + ) + + # Create full entity key and serialize it + entity_key = EntityKeyProto( + join_keys=[key], entity_values=[ValueProto(string_val="test_value")] + ) + full_result = serialize_entity_key( + entity_key, entity_key_serialization_version=3 + ) + + # The prefix should match the beginning of the full serialization + # Extract just the key portion (skip entity count, but include key metadata) + prefix_len = len(prefix_result) + assert full_result[:prefix_len] == prefix_result, ( + f"Prefix mismatch for non-ASCII key '{key}': " + f"Character length: {len(key)}, " + f"UTF-8 byte length: {len(key.encode('utf8'))}" + ) + + +def test_ascii_prefix_compatibility(): + """Verify prefix compatibility still works for ASCII keys.""" + ascii_keys = ["user_id", "session_id", "device_id"] + + for key in ascii_keys: + prefix_result = serialize_entity_key_prefix( + [key], entity_key_serialization_version=3 + ) + + entity_key = EntityKeyProto( + join_keys=[key], entity_values=[ValueProto(string_val="test_value")] + ) + full_result = serialize_entity_key( + entity_key, entity_key_serialization_version=3 + ) + + prefix_len = len(prefix_result) + assert full_result[:prefix_len] == prefix_result, ( + f"Prefix mismatch for ASCII key '{key}'" + ) + + +def test_multi_key_non_ascii_prefix_compatibility(): + """Test multi-key prefix compatibility with non-ASCII characters.""" + mixed_keys = ["user_id", "用户会话", "session_id"] # Mix ASCII and non-ASCII + + prefix_result = serialize_entity_key_prefix( + mixed_keys, entity_key_serialization_version=3 + ) + + entity_key = EntityKeyProto( + join_keys=mixed_keys, + entity_values=[ + ValueProto(string_val="test1"), + ValueProto(string_val="test2"), + ValueProto(string_val="test3"), + ], + ) + full_result = serialize_entity_key(entity_key, entity_key_serialization_version=3) + + prefix_len = len(prefix_result) + assert full_result[:prefix_len] == prefix_result, ( + "Multi-key prefix mismatch with non-ASCII" + )