feast/sdk/python/feast/infra/online_stores/redis.py at stable · feast-dev/feast

757 lines (654 loc) · 28.9 KB
# 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 json
import logging
from datetime import datetime, timezone
from enum import Enum
from typing import (
    ByteString,
    Callable,
    Literal,
    Optional,
    Sequence,
from google.protobuf.timestamp_pb2 import Timestamp
from pydantic import StrictStr
from feast import Entity, FeatureView, RepoConfig, utils
from feast.infra.key_encoding_utils import serialize_entity_key
from feast.infra.online_stores.helpers import (
    _redis_key,
    _redis_key_prefix,
    compute_versioned_name,
from feast.infra.online_stores.online_store import OnlineStore
from feast.infra.supported_async_methods import SupportedAsyncMethods
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
    from redis import Redis
    from redis import asyncio as redis_asyncio
    from redis.cluster import ClusterNode, RedisCluster
    from redis.sentinel import Sentinel
except ImportError as e:
    from feast.errors import FeastExtrasDependencyImportError
    raise FeastExtrasDependencyImportError("redis", str(e))
logger = logging.getLogger(__name__)
def _versioned_fv_name(table: FeatureView, config: RepoConfig) -> str:
    """Return the feature view name with version suffix when versioning is enabled."""
    return compute_versioned_name(
        table, config.registry.enable_online_feature_view_versioning
class RedisType(str, Enum):
    redis = "redis"
    redis_cluster = "redis_cluster"
    redis_sentinel = "redis_sentinel"
class RedisOnlineStoreConfig(FeastConfigBaseModel):
    """Online store config for Redis store"""
    type: Literal["redis"] = "redis"
    """Online store type selector"""
    redis_type: RedisType = RedisType.redis
    """Redis type: redis or redis_cluster"""
    sentinel_master: StrictStr = "mymaster"
    """Sentinel's master name"""
    connection_string: StrictStr = "localhost:6379"
    """Connection string containing the host, port, and configuration parameters for Redis
     format: host:port,parameter1,parameter2 eg. redis:6379,db=0 """
    key_ttl_seconds: Optional[int] = None
    """(Optional) redis key bin ttl (in seconds) for expiring entities"""
    full_scan_for_deletion: Optional[bool] = True
    """(Optional) whether to scan for deletion of features"""
    skip_dedup: bool = False
    """(Optional) Skip timestamp deduplication check on writes for higher throughput.
    When True, writes proceed in a single pipeline without reading existing timestamps first.
    This may cause older feature values to overwrite newer ones under concurrent writers."""
class RedisOnlineStore(OnlineStore):
    Redis implementation of the online store interface.
    See https://github.com/feast-dev/feast/blob/master/docs/specs/online_store_format.md#redis-online-store-format
    for more details about the data model for this implementation.
    Attributes:
        _client: Redis connection.
    _client: Optional[Union[Redis, RedisCluster]] = None
    _client_async: Optional[Union[redis_asyncio.Redis, redis_asyncio.RedisCluster]] = (
        None
    @property
    def async_supported(self) -> SupportedAsyncMethods:
        return SupportedAsyncMethods(read=True, write=True)
    def delete_entity_values(self, config: RepoConfig, join_keys: List[str]):
        client = self._get_client(config.online_store)
        deleted_count = 0
        prefix = _redis_key_prefix(join_keys)
        with client.pipeline(transaction=False) as pipe:
            for _k in client.scan_iter(
                b"".join([prefix, b"*", config.project.encode("utf8")])
                pipe.delete(_k)
                deleted_count += 1
            pipe.execute()
        logger.debug(f"Deleted {deleted_count} rows for entity {', '.join(join_keys)}")
    def delete_table(self, config: RepoConfig, table: FeatureView):
        """
        Delete all rows in Redis for a specific feature view.
        Uses a two-phase pipelined approach to avoid an O(K) synchronous hkeys
        call inside the scan loop (N+1 pattern).
        Args:
            config: Feast config
            table: Feature view to delete
        """
        client = self._get_client(config.online_store)
        prefix = _redis_key_prefix(table.join_keys)
        fv_name = _versioned_fv_name(table, config)
        fv_name_bytes = bytes(fv_name, "utf8")
        redis_hash_keys = [_mmh3(f"{fv_name}:{f.name}") for f in table.features]
        redis_hash_keys.append(bytes(f"_ts:{fv_name}", "utf8"))
        # Phase 1: collect all matching entity keys from SCAN (no per-key round trips)
        scan_pattern = b"".join([prefix, b"*", config.project.encode("utf8")])
        all_keys = list(client.scan_iter(scan_pattern))
        if not all_keys:
            logger.debug(f"Deleted 0 rows for feature view {fv_name}")
            return
        # Phase 2: pipeline hkeys for all collected entity keys (1 round trip)
        with client.pipeline(transaction=False) as pipe:
            for _k in all_keys:
                pipe.hkeys(_k)
            all_hkeys = pipe.execute()
        # Phase 3: pipeline all deletions based on phase 2 results (1 round trip)
        deleted_count = 0
        with client.pipeline(transaction=False) as pipe:
            for _k, field_names in zip(all_keys, all_hkeys):
                _tables = {_hk[4:] for _hk in field_names if _hk.startswith(b"_ts:")}
                if fv_name_bytes not in _tables:
                    continue
                if len(_tables) == 1:
                    pipe.delete(_k)
                else:
                    pipe.hdel(_k, *redis_hash_keys)
                deleted_count += 1
            pipe.execute()
        logger.debug(f"Deleted {deleted_count} rows for feature view {fv_name}")
    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,
        """
        Delete data from feature views that are no longer in use.
        Args:
            config: Feast config
            tables_to_delete: Feature views to delete
            tables_to_keep: Feature views to keep
            entities_to_delete: Entities to delete
            entities_to_keep: Entities to keep
            partial: Whether to do a partial update
        """
        online_store_config = config.online_store
        assert isinstance(online_store_config, RedisOnlineStoreConfig)
        if online_store_config.full_scan_for_deletion:
            for table in tables_to_delete:
                self.delete_table(config, table)
    def teardown(
        self,
        config: RepoConfig,
        tables: Sequence[FeatureView],
        entities: Sequence[Entity],
        """
        We delete the keys in redis for tables/views being removed.
        """
        join_keys_to_delete = set(tuple(table.join_keys) for table in tables)
        for join_keys in join_keys_to_delete:
            self.delete_entity_values(config, list(join_keys))
    @staticmethod
    def _parse_connection_string(connection_string: str):
        """
        Reads Redis connections string using format
        for RedisCluster:
            redis1:6379,redis2:6379,skip_full_coverage_check=true,ssl=true,password=...
        for Redis:
            redis_master:6379,db=0,ssl=true,password=...
        """
        startup_nodes = [
            dict(zip(["host", "port"], c.split(":")))
            for c in connection_string.split(",")
            if "=" not in c
        params = {}
        for c in connection_string.split(","):
            if "=" in c:
                kv = c.split("=", 1)
                    kv[1] = json.loads(kv[1])
                except json.JSONDecodeError:
                it = iter(kv)
                params.update(dict(zip(it, it)))
        return startup_nodes, params
    def _get_client(self, online_store_config: RedisOnlineStoreConfig):
        """
        Creates the Redis client RedisCluster or Redis depending on configuration
        """
        if not self._client:
            startup_nodes, kwargs = self._parse_connection_string(
                online_store_config.connection_string
            if online_store_config.redis_type == RedisType.redis_cluster:
                kwargs["startup_nodes"] = [
                    ClusterNode(**node) for node in startup_nodes
                self._client = RedisCluster(**kwargs)
            elif online_store_config.redis_type == RedisType.redis_sentinel:
                sentinel_hosts = []
                for item in startup_nodes:
                    sentinel_hosts.append((item["host"], int(item["port"])))
                sentinel = Sentinel(sentinel_hosts, **kwargs)
                master = sentinel.master_for(online_store_config.sentinel_master)
                self._client = master
            else:
                kwargs["host"] = startup_nodes[0]["host"]
                kwargs["port"] = startup_nodes[0]["port"]
                self._client = Redis(**kwargs)
        return self._client
    async def _get_client_async(self, online_store_config: RedisOnlineStoreConfig):
        if not self._client_async:
            startup_nodes, kwargs = self._parse_connection_string(
                online_store_config.connection_string
            if online_store_config.redis_type == RedisType.redis_cluster:
                kwargs["startup_nodes"] = [
                    redis_asyncio.cluster.ClusterNode(**node) for node in startup_nodes
                self._client_async = redis_asyncio.RedisCluster(**kwargs)
            elif online_store_config.redis_type == RedisType.redis_sentinel:
                sentinel_hosts = []
                for item in startup_nodes:
                    sentinel_hosts.append((item["host"], int(item["port"])))
                sentinel = redis_asyncio.Sentinel(sentinel_hosts, **kwargs)
                master = sentinel.master_for(online_store_config.sentinel_master)
                self._client_async = master
            else:
                kwargs["host"] = startup_nodes[0]["host"]
                kwargs["port"] = startup_nodes[0]["port"]
                self._client_async = redis_asyncio.Redis(**kwargs)
        return self._client_async
    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:
        online_store_config = config.online_store
        assert isinstance(online_store_config, RedisOnlineStoreConfig)
        client = self._get_client(online_store_config)
        project = config.project
        feature_view = _versioned_fv_name(table, config)
        ts_key = f"_ts:{feature_view}"
        if online_store_config.skip_dedup:
            # Single-pipeline fast path: no timestamp read, directly write all rows.
            # Reduces round trips from 2 to 1. Suitable for initial loads or
            # append-only pipelines where out-of-order writes are not a concern.
            with client.pipeline(transaction=False) as pipe:
                for entity_key, values, timestamp, _ in data:
                    redis_key_bin = _redis_key(
                        project,
                        entity_key,
                        entity_key_serialization_version=config.entity_key_serialization_version,
                    aware_ts = utils.make_tzaware(timestamp)
                    ts = Timestamp()
                    ts.FromDatetime(aware_ts)
                    entity_hset: Dict[Any, Any] = {ts_key: ts.SerializeToString()}
                    for feature_name, val in values.items():
                        f_key = _mmh3(f"{feature_view}:{feature_name}")
                        entity_hset[f_key] = val.SerializeToString()
                    pipe.hset(redis_key_bin, mapping=entity_hset)
                    if online_store_config.key_ttl_seconds:
                        pipe.expire(
                            name=redis_key_bin,
                            time=online_store_config.key_ttl_seconds,
                results = pipe.execute()
            if progress:
                progress(len(results))
            return
        keys = []
        # redis pipelining optimization: send multiple commands to redis server without waiting for every reply
        with client.pipeline(transaction=False) as pipe:
            # check if a previous record under the key bin exists
            # TODO: investigate if check and set is a better approach rather than pulling all entity ts and then setting
            # it may be significantly slower but avoids potential (rare) race conditions
            for entity_key, _, _, _ in data:
                redis_key_bin = _redis_key(
                    project,
                    entity_key,
                    entity_key_serialization_version=config.entity_key_serialization_version,
                keys.append(redis_key_bin)
                pipe.hmget(redis_key_bin, ts_key)
            prev_event_timestamps = pipe.execute()
            # flattening the list of lists. `hmget` does the lookup assuming a list of keys in the key bin
            prev_event_timestamps = [i[0] for i in prev_event_timestamps]
            for redis_key_bin, prev_event_time, (_, values, timestamp, _) in zip(
                keys, prev_event_timestamps, data
                # Convert incoming timestamp to millisecond-aware datetime
                aware_ts = utils.make_tzaware(timestamp)
                # Build protobuf timestamp with nanos
                ts = Timestamp()
                ts.FromDatetime(aware_ts)
                # New timestamp in nanoseconds
                new_total_nanos = ts.seconds * 1_000_000_000 + ts.nanos
                # Compare against existing timestamp (nanosecond precision)
                if prev_event_time:
                    prev_ts = Timestamp()
                    prev_ts.ParseFromString(prev_event_time)
                    prev_total_nanos = prev_ts.seconds * 1_000_000_000 + prev_ts.nanos
                    # Skip only if older OR exact same instant
                    if prev_total_nanos and new_total_nanos <= prev_total_nanos:
                        if progress:
                            progress(1)
                        continue
                # Store full timestamp (seconds + nanos)
                entity_hset = {ts_key: ts.SerializeToString()}
                for feature_name, val in values.items():
                    f_key = _mmh3(f"{feature_view}:{feature_name}")
                    entity_hset[f_key] = val.SerializeToString()
                pipe.hset(redis_key_bin, mapping=entity_hset)
                if online_store_config.key_ttl_seconds:
                    pipe.expire(
                        name=redis_key_bin, time=online_store_config.key_ttl_seconds
            results = pipe.execute()
            if progress:
                progress(len(results))
    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:
        """Async version of online_write_batch using the async Redis client."""
        online_store_config = config.online_store
        assert isinstance(online_store_config, RedisOnlineStoreConfig)
        client = await self._get_client_async(online_store_config)
        project = config.project
        feature_view = _versioned_fv_name(table, config)
        ts_key = f"_ts:{feature_view}"
        if online_store_config.skip_dedup:
            async with client.pipeline(transaction=False) as pipe:
                for entity_key, values, timestamp, _ in data:
                    redis_key_bin = _redis_key(
                        project,
                        entity_key,
                        entity_key_serialization_version=config.entity_key_serialization_version,
                    aware_ts = utils.make_tzaware(timestamp)
                    ts = Timestamp()
                    ts.FromDatetime(aware_ts)
                    entity_hset: Dict[Any, Any] = {ts_key: ts.SerializeToString()}
                    for feature_name, val in values.items():
                        f_key = _mmh3(f"{feature_view}:{feature_name}")
                        entity_hset[f_key] = val.SerializeToString()
                    pipe.hset(redis_key_bin, mapping=entity_hset)
                    if online_store_config.key_ttl_seconds:
                        pipe.expire(
                            name=redis_key_bin,
                            time=online_store_config.key_ttl_seconds,
                results = await pipe.execute()
            if progress:
                progress(len(results))
            return
        keys = []
        async with client.pipeline(transaction=False) as pipe:
            for entity_key, _, _, _ in data:
                redis_key_bin = _redis_key(
                    project,
                    entity_key,
                    entity_key_serialization_version=config.entity_key_serialization_version,
                keys.append(redis_key_bin)
                pipe.hmget(redis_key_bin, ts_key)
            prev_event_timestamps = await pipe.execute()
        prev_event_timestamps = [i[0] for i in prev_event_timestamps]
        async with client.pipeline(transaction=False) as pipe:
            for redis_key_bin, prev_event_time, (_, values, timestamp, _) in zip(
                keys, prev_event_timestamps, data
                aware_ts = utils.make_tzaware(timestamp)
                ts = Timestamp()
                ts.FromDatetime(aware_ts)
                new_total_nanos = ts.seconds * 1_000_000_000 + ts.nanos
                if prev_event_time:
                    prev_ts = Timestamp()
                    prev_ts.ParseFromString(prev_event_time)
                    prev_total_nanos = prev_ts.seconds * 1_000_000_000 + prev_ts.nanos
                    if prev_total_nanos and new_total_nanos <= prev_total_nanos:
                        if progress:
                            progress(1)
                        continue
                entity_hset = {ts_key: ts.SerializeToString()}
                for feature_name, val in values.items():
                    f_key = _mmh3(f"{feature_view}:{feature_name}")
                    entity_hset[f_key] = val.SerializeToString()
                pipe.hset(redis_key_bin, mapping=entity_hset)
                if online_store_config.key_ttl_seconds:
                    pipe.expire(
                        name=redis_key_bin, time=online_store_config.key_ttl_seconds
            results = await pipe.execute()
            if progress:
                progress(len(results))
    def _generate_redis_keys_for_entities(
        self, config: RepoConfig, entity_keys: List[EntityKeyProto]
    ) -> List[bytes]:
        project = config.project
        version = config.entity_key_serialization_version
        project_bytes = project.encode("utf-8")
        return [
            serialize_entity_key(ek, entity_key_serialization_version=version)
            + project_bytes
            for ek in entity_keys
    def _generate_hset_keys_for_features(
        self,
        feature_view: FeatureView,
        requested_features: Optional[List[str]] = None,
        fv_name_override: Optional[str] = None,
    ) -> Tuple[List[str], List[str]]:
        if not requested_features:
            requested_features = [f.name for f in feature_view.features]
        fv_name = fv_name_override or feature_view.name
        hset_keys = [_mmh3(f"{fv_name}:{k}") for k in requested_features]
        ts_key = f"_ts:{fv_name}"
        hset_keys.append(ts_key)
        requested_features.append(ts_key)
        return requested_features, hset_keys
    def _convert_redis_values_to_protobuf(
        self,
        redis_values: List[List[ByteString]],
        feature_view: str,
        requested_features: List[str],
    ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
        return [
            self._get_features_for_entity(values, feature_view, requested_features)
            for values in redis_values
    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]]]]:
        online_store_config = config.online_store
        assert isinstance(online_store_config, RedisOnlineStoreConfig)
        client = self._get_client(online_store_config)
        feature_view = table
        fv_name = _versioned_fv_name(table, config)
        requested_features, hset_keys = self._generate_hset_keys_for_features(
            feature_view, requested_features, fv_name_override=fv_name
        keys = self._generate_redis_keys_for_entities(config, entity_keys)
        with client.pipeline(transaction=False) as pipe:
            for redis_key_bin in keys:
                pipe.hmget(redis_key_bin, hset_keys)
            redis_values = pipe.execute()
        return self._convert_redis_values_to_protobuf(
            redis_values, fv_name, requested_features
    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]]]]:
        online_store_config = config.online_store
        assert isinstance(online_store_config, RedisOnlineStoreConfig)
        client = await self._get_client_async(online_store_config)
        feature_view = table
        fv_name = _versioned_fv_name(table, config)
        requested_features, hset_keys = self._generate_hset_keys_for_features(
            feature_view, requested_features, fv_name_override=fv_name
        keys = self._generate_redis_keys_for_entities(config, entity_keys)
        async with client.pipeline(transaction=False) as pipe:
            for redis_key_bin in keys:
                pipe.hmget(redis_key_bin, hset_keys)
            redis_values = await pipe.execute()
        return self._convert_redis_values_to_protobuf(
            redis_values, fv_name, requested_features
    def _read_features_per_fv(
        self,
        config: RepoConfig,
        grouped_refs,
        join_key_values,
        entity_name_to_join_key_map,
        online_features_response,
        full_feature_names: bool,
        include_feature_view_version_metadata: bool,
    ) -> None:
        """Batch all per-FV HMGET commands into a single Redis pipeline."""
        work_items = []
        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)
            fv_name = _versioned_fv_name(table, config)
            redis_keys = self._generate_redis_keys_for_entities(
                config, entity_key_protos
            req_features, hset_keys = self._generate_hset_keys_for_features(
                table, requested_features, fv_name_override=fv_name
            work_items.append(
                (table, req_features, fv_name, hset_keys, redis_keys, idxs, output_len)
        if work_items:
            client = self._get_client(config.online_store)
            with client.pipeline(transaction=False) as pipe:
                for _, _, _, hset_keys, redis_keys, _, _ in work_items:
                    for redis_key in redis_keys:
                        pipe.hmget(redis_key, hset_keys)
                all_results = pipe.execute()
            offset = 0
            for (
                table,
                req_features,
                fv_name,
                redis_keys,
                idxs,
                output_len,
            ) in work_items:
                n = len(redis_keys)
                redis_values = all_results[offset : offset + n]
                offset += n
                read_rows = self._convert_redis_values_to_protobuf(
                    redis_values, fv_name, req_features
                utils._populate_response_from_feature_data(
                    req_features,
                    read_rows,
                    online_features_response,
                    full_feature_names,
                    table,
                    output_len,
                    include_feature_view_version_metadata,
    async def _read_features_per_fv_async(
        self,
        config: RepoConfig,
        grouped_refs,
        join_key_values,
        entity_name_to_join_key_map,
        online_features_response,
        full_feature_names: bool,
        include_feature_view_version_metadata: bool,
    ) -> None:
        """Async version: batch all per-FV HMGET into a single async pipeline."""
        work_items = []
        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)
            fv_name = _versioned_fv_name(table, config)
            redis_keys = self._generate_redis_keys_for_entities(
                config, entity_key_protos
            req_features, hset_keys = self._generate_hset_keys_for_features(
                table, requested_features, fv_name_override=fv_name
            work_items.append(
                (table, req_features, fv_name, hset_keys, redis_keys, idxs, output_len)
        if work_items:
            client = await self._get_client_async(config.online_store)
            async with client.pipeline(transaction=False) as pipe:
                for _, _, _, hset_keys, redis_keys, _, _ in work_items:
                    for redis_key in redis_keys:
                        pipe.hmget(redis_key, hset_keys)
                all_results = await pipe.execute()
            offset = 0
            for (
                table,
                req_features,
                fv_name,
                redis_keys,
                idxs,
                output_len,
            ) in work_items:
                n = len(redis_keys)
                redis_values = all_results[offset : offset + n]
                offset += n
                read_rows = self._convert_redis_values_to_protobuf(
                    redis_values, fv_name, req_features
                utils._populate_response_from_feature_data(
                    req_features,
                    read_rows,
                    online_features_response,
                    full_feature_names,
                    table,
                    output_len,
                    include_feature_view_version_metadata,
    def _get_features_for_entity(
        self,
        values: List[ByteString],
        feature_view: str,
        requested_features: List[str],
    ) -> Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]:
        res_val = dict(zip(requested_features, values))
        res_ts = Timestamp()
        ts_key = f"_ts:{feature_view}"
        ts_val = res_val.pop(ts_key)
        if ts_val:
            res_ts.ParseFromString(ts_val)
        res: Dict[str, ValueProto] = {}
        for feature_name, val_bin in res_val.items():
            val = ValueProto()
            if val_bin:
                val.ParseFromString(val_bin)
            res[feature_name] = val
        if not res:
            return None, None
        total_seconds = res_ts.seconds + res_ts.nanos / 1_000_000_000.0
        timestamp = datetime.fromtimestamp(total_seconds, tz=timezone.utc)
        return timestamp, res
Provide feedback

Saved searches

Use saved searches to filter your results more quickly

FilesExpand file tree

redis.py

Latest commit

History

redis.py

File metadata and controls
