import json

from abc import abstractmethod

from datetime import datetime

from typing import TYPE_CHECKING, Dict, List, Optional, Type, cast

import pandas as pd

import pyarrow

from google.protobuf.json_format import MessageToJson

from feast.data_source import DataSource

from feast.dqm.profilers.profiler import Profile, Profiler

from feast.importer import import_class

from feast.protos.feast.core.DataSource_pb2 import DataSource as DataSourceProto

from feast.protos.feast.core.SavedDataset_pb2 import SavedDataset as SavedDatasetProto

from feast.protos.feast.core.SavedDataset_pb2 import SavedDatasetMeta, SavedDatasetSpec

from feast.protos.feast.core.SavedDataset_pb2 import (

SavedDatasetStorage as SavedDatasetStorageProto,

)

from feast.protos.feast.core.ValidationProfile_pb2 import (

ValidationReference as ValidationReferenceProto,

)

if TYPE_CHECKING:

from feast.infra.offline_stores.offline_store import RetrievalJob

class _StorageRegistry(type):

classes_by_proto_attr_name: Dict[str, Type["SavedDatasetStorage"]] = {}

# Several storage classes share the ``custom_storage`` oneof field, so the

# mapping above cannot identify them. Every claimant is recorded here in

# registration order so that dispatch can tell "one obvious class" apart

# from "several, pick deliberately".

all_classes_by_proto_attr_name: Dict[str, List[Type["SavedDatasetStorage"]]] = {}

def __new__(cls, name, bases, dct):

kls = type.__new__(cls, name, bases, dct)

proto_attr_name = dct.get("_proto_attr_name")

if proto_attr_name:

cls.classes_by_proto_attr_name[proto_attr_name] = kls

claimants = cls.all_classes_by_proto_attr_name.setdefault(

proto_attr_name, []

)

if kls not in claimants:

claimants.append(kls)

return kls

# Key under which a custom_storage payload records the concrete storage class

# that wrote it. CustomSourceOptions.configuration is an opaque byte string the

# implementer owns, and every claimant already reads it as JSON by explicit key,

# so an extra key is ignored by older readers.

CUSTOM_STORAGE_CLASS_KEY = "saved_dataset_storage_class"

def _class_path(cls: type) -> str:

return f"{cls.__module__}.{cls.__qualname__}"

def _read_custom_storage_class(

options_proto: DataSourceProto.CustomSourceOptions,

) -> Optional[Type["SavedDatasetStorage"]]:

"""Resolve the class recorded in a custom_storage payload, if it has one.

Returns ``None`` for a payload written before this key existed, or one whose

configuration is not the JSON object the claimants write, so the caller can

fall back instead of failing on data it could previously read.

"""

try:

config = json.loads(options_proto.configuration.decode("utf8"))

except (UnicodeDecodeError, ValueError):

return None

if not isinstance(config, dict):

return None

class_path = config.get(CUSTOM_STORAGE_CLASS_KEY)

if not isinstance(class_path, str) or not class_path:

return None

try:

return get_saved_dataset_storage_class_from_path(class_path)

except Exception:

# A recorded path can be unresolvable here: the class may live in an

# extra that is not installed, or may not be importable at all. Fall

# back rather than turn a readable payload into an import failure.

return None

_DATA_SOURCE_TO_SAVED_DATASET_STORAGE = {

"FileSource": "feast.infra.offline_stores.file_source.SavedDatasetFileStorage",

"SparkSource": "feast.infra.offline_stores.contrib.spark_offline_store.spark_source.SavedDatasetSparkStorage",

"LanceSource": "feast.infra.data_sources.contrib.lance.lance_source.SavedDatasetLanceStorage",

}

def get_saved_dataset_storage_class_from_path(saved_dataset_storage_path: str):

module_name, class_name = saved_dataset_storage_path.rsplit(".", 1)

return import_class(module_name, class_name, "SavedDatasetStorage")

class SavedDatasetStorage(metaclass=_StorageRegistry):

_proto_attr_name: str

@staticmethod

def from_proto(storage_proto: SavedDatasetStorageProto) -> "SavedDatasetStorage":

proto_attr_name = cast(str, storage_proto.WhichOneof("kind"))

if proto_attr_name == "custom_storage":

recorded = _read_custom_storage_class(storage_proto.custom_storage)

if recorded is not None:

return recorded.from_proto(storage_proto)

claimants = _StorageRegistry.all_classes_by_proto_attr_name.get(

proto_attr_name, []

)

if len(claimants) > 1:

raise ValueError(

f"{len(claimants)} storage classes claim the "

f"'{proto_attr_name}' field and this payload does not record "

f"which one wrote it, so it cannot be read back unambiguously: "

f"{', '.join(sorted(_class_path(c) for c in claimants))}. "

f"Saved datasets written after this version record their class "

f"under the '{CUSTOM_STORAGE_CLASS_KEY}' key; this one predates "

f"that, and the payloads of these classes are not otherwise "

f"distinguishable. Call the intended class's from_proto directly."

)

return _StorageRegistry.classes_by_proto_attr_name[proto_attr_name].from_proto(

storage_proto

)

def _custom_storage_proto(

self, options_proto: DataSourceProto.CustomSourceOptions

) -> SavedDatasetStorageProto:

"""Wrap custom source options, recording which class wrote them.

Classes that share the ``custom_storage`` field must use this instead of

constructing the proto directly, otherwise ``from_proto`` has nothing to

dispatch on and has to guess from import order.

"""

config = json.loads(options_proto.configuration.decode("utf8"))

config[CUSTOM_STORAGE_CLASS_KEY] = _class_path(type(self))

return SavedDatasetStorageProto(

custom_storage=DataSourceProto.CustomSourceOptions(

configuration=json.dumps(config).encode()

)

)

@abstractmethod

def to_proto(self) -> SavedDatasetStorageProto:

pass

@abstractmethod

def to_data_source(self) -> DataSource:

pass

@staticmethod

def from_data_source(data_source: DataSource) -> "SavedDatasetStorage":

data_source_type = type(data_source).__name__

if data_source_type in _DATA_SOURCE_TO_SAVED_DATASET_STORAGE:

cls = get_saved_dataset_storage_class_from_path(

_DATA_SOURCE_TO_SAVED_DATASET_STORAGE[data_source_type]

)

return cls.from_data_source(data_source)

else:

raise ValueError(

f"This method currently does not support {data_source_type}."

)

class SavedDataset:

name: str

features: List[str]

join_keys: List[str]

full_feature_names: bool

storage: SavedDatasetStorage

tags: Dict[str, str]

feature_service_name: Optional[str] = None

namespace: str = ""

collection: str = ""

description: str = ""

created_timestamp: Optional[datetime] = None

last_updated_timestamp: Optional[datetime] = None

min_event_timestamp: Optional[datetime] = None

max_event_timestamp: Optional[datetime] = None

_retrieval_job: Optional["RetrievalJob"] = None

def __init__(

self,

name: str,

features: List[str],

join_keys: List[str],

storage: SavedDatasetStorage,

full_feature_names: bool = False,

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

feature_service_name: Optional[str] = None,

namespace: str = "",

collection: str = "",

description: str = "",

):

self.name = name

self.features = features

self.join_keys = join_keys

self.storage = storage

self.full_feature_names = full_feature_names

self.tags = tags or {}

self.feature_service_name = feature_service_name

self.namespace = namespace

self.collection = collection

self.description = description

self._retrieval_job = None

def __repr__(self):

items = (f"{k} = {v}" for k, v in self.__dict__.items())

return f"<{self.__class__.__name__}({', '.join(items)})>"

def __str__(self):

return str(MessageToJson(self.to_proto()))

def __hash__(self):

return hash((self.name))

def __eq__(self, other):

if not isinstance(other, SavedDataset):

return False

if (

self.name != other.name

or sorted(self.features) != sorted(other.features)

or sorted(self.join_keys) != sorted(other.join_keys)

or self.storage != other.storage

or self.full_feature_names != other.full_feature_names

or self.tags != other.tags

or self.feature_service_name != other.feature_service_name

or self.namespace != other.namespace

or self.collection != other.collection

or self.description != other.description

):

return False

return True

@staticmethod

def from_proto(saved_dataset_proto: SavedDatasetProto):

"""

Converts a SavedDatasetProto to a SavedDataset object.

Args:

saved_dataset_proto: A protobuf representation of a SavedDataset.

"""

ds = SavedDataset(

name=saved_dataset_proto.spec.name,

features=list(saved_dataset_proto.spec.features),

join_keys=list(saved_dataset_proto.spec.join_keys),

full_feature_names=saved_dataset_proto.spec.full_feature_names,

storage=SavedDatasetStorage.from_proto(saved_dataset_proto.spec.storage),

tags=dict(saved_dataset_proto.spec.tags.items()),

namespace=saved_dataset_proto.spec.namespace,

collection=saved_dataset_proto.spec.collection,

description=saved_dataset_proto.spec.description,

)

if saved_dataset_proto.spec.feature_service_name:

ds.feature_service_name = saved_dataset_proto.spec.feature_service_name

if saved_dataset_proto.meta.HasField("created_timestamp"):

ds.created_timestamp = (

saved_dataset_proto.meta.created_timestamp.ToDatetime()

)

if saved_dataset_proto.meta.HasField("last_updated_timestamp"):

ds.last_updated_timestamp = (

saved_dataset_proto.meta.last_updated_timestamp.ToDatetime()

)

if saved_dataset_proto.meta.HasField("min_event_timestamp"):

ds.min_event_timestamp = (

saved_dataset_proto.meta.min_event_timestamp.ToDatetime()

)

if saved_dataset_proto.meta.HasField("max_event_timestamp"):

ds.max_event_timestamp = (

saved_dataset_proto.meta.max_event_timestamp.ToDatetime()

)

return ds

def to_proto(self) -> SavedDatasetProto:

"""

Converts a SavedDataset to its protobuf representation.

Returns:

A SavedDatasetProto protobuf.

"""

meta = SavedDatasetMeta()

if self.created_timestamp:

meta.created_timestamp.FromDatetime(self.created_timestamp)

if self.min_event_timestamp:

meta.min_event_timestamp.FromDatetime(self.min_event_timestamp)

if self.max_event_timestamp:

meta.max_event_timestamp.FromDatetime(self.max_event_timestamp)

spec = SavedDatasetSpec(

name=self.name,

features=self.features,

join_keys=self.join_keys,

full_feature_names=self.full_feature_names,

storage=self.storage.to_proto(),

tags=self.tags,

namespace=self.namespace,

collection=self.collection,

description=self.description,

)

if self.feature_service_name:

spec.feature_service_name = self.feature_service_name

saved_dataset_proto = SavedDatasetProto(spec=spec, meta=meta)

return saved_dataset_proto

def with_retrieval_job(self, retrieval_job: "RetrievalJob") -> "SavedDataset":

self._retrieval_job = retrieval_job

return self

def to_df(self) -> pd.DataFrame:

if not self._retrieval_job:

raise RuntimeError(

"To load this dataset use FeatureStore.get_saved_dataset() "

"instead of instantiating it directly."

)

return self._retrieval_job.to_df()

def to_arrow(self) -> pyarrow.Table:

if not self._retrieval_job:

raise RuntimeError(

"To load this dataset use FeatureStore.get_saved_dataset() "

"instead of instantiating it directly."

)

return self._retrieval_job.to_arrow()

def as_reference(self, name: str, profiler: "Profiler") -> "ValidationReference":

return ValidationReference.from_saved_dataset(

name=name, profiler=profiler, dataset=self

)

def get_profile(self, profiler: Profiler) -> Profile:

return profiler.analyze_dataset(self.to_df())

class ValidationReference:

name: str

dataset_name: str

description: str

tags: Dict[str, str]

profiler: Optional[Profiler]

_profile: Optional[Profile] = None

_dataset: Optional[SavedDataset] = None

def __init__(

self,

name: str,

dataset_name: str,

profiler: Optional[Profiler] = None,

description: str = "",

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

):

"""

Validation reference combines a reference dataset (currently only a saved dataset object can be used as

a reference) and a profiler function to generate a validation profile.

The validation profile can be cached in this object, and in this case

the saved dataset retrieval and the profiler call will happen only once.

Validation reference is being stored in the Feast registry and can be retrieved by its name, which

must be unique within one project.

Args:

name: the unique name for validation reference

dataset_name: the name of the saved dataset used as a reference

description: a human-readable description

tags: a dictionary of key-value pairs to store arbitrary metadata

profiler: the profiler function used to generate profile from the saved dataset

"""

self.name = name

self.dataset_name = dataset_name

self.profiler = profiler

self.description = description

self.tags = tags or {}

@classmethod

def from_saved_dataset(cls, name: str, dataset: SavedDataset, profiler: Profiler):

"""

Internal constructor to create validation reference object with actual saved dataset object

(regular constructor requires only its name).

"""

ref = ValidationReference(name, dataset.name, profiler)

ref._dataset = dataset

return ref

@property

def profile(self) -> Profile:

if not self._profile:

if not self._dataset:

raise RuntimeError(

"In order to calculate a profile validation reference must be instantiated from a saved dataset. "

"Use ValidationReference.from_saved_dataset constructor or FeatureStore.get_validation_reference "

"to get validation reference object."

)

if not self.profiler:

raise RuntimeError(

"Failed to calculate profile: profiler is not set on this "

"ValidationReference."

)

self._profile = self.profiler.analyze_dataset(self._dataset.to_df())

return self._profile

@classmethod

def from_proto(

cls, proto: ValidationReferenceProto, skip_udf: bool = False

) -> "ValidationReference":

"""Build a ValidationReference from its proto.

Args:

proto: the serialized validation reference.

skip_udf: when True, leave ``profiler`` and the cached profile unset

instead of deserializing them. Both are dill payloads, so callers

that only need identity metadata (notably authorization checks in

the registry server) must not deserialize them.

"""

profiler: Optional[Profiler] = None

profile: Optional[Profile] = None

if not skip_udf:

profiler_attr = proto.WhichOneof("profiler")

if profiler_attr == "ge_profiler":

from feast.dqm.profilers.ge_profiler import GEProfiler

profiler = GEProfiler.from_proto(proto.ge_profiler)

else:

raise RuntimeError("Unrecognized profiler")

profile_attr = proto.WhichOneof("cached_profile")

if profile_attr == "ge_profile":

from feast.dqm.profilers.ge_profiler import GEProfile

profile = GEProfile.from_proto(proto.ge_profile)

elif not profile_attr:

profile = None

else:

raise RuntimeError("Unrecognized profile")

ref = ValidationReference(

name=proto.name,

dataset_name=proto.reference_dataset_name,

profiler=profiler,

description=proto.description,

tags=dict(proto.tags),

)

ref._profile = profile

return ref

def to_proto(self) -> ValidationReferenceProto:

from feast.dqm.profilers.ge_profiler import GEProfile, GEProfiler

proto = ValidationReferenceProto(

name=self.name,

reference_dataset_name=self.dataset_name,

tags=self.tags,

description=self.description,

ge_profiler=self.profiler.to_proto()

if isinstance(self.profiler, GEProfiler)

else None,

ge_profile=self._profile.to_proto()

if isinstance(self._profile, GEProfile)

else None,

)

return proto