import functools

import warnings

from datetime import datetime, timedelta

from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import dill

from feast import flags_helper

from feast.aggregation import Aggregation

from feast.data_source import DataSource

from feast.entity import Entity

from feast.feature_view import FeatureView

from feast.field import Field

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

from feast.transformation.base import Transformation

from feast.transformation.mode import TransformationMode

warnings.simplefilter("once", RuntimeWarning)

SUPPORTED_BATCH_SOURCES = {

"BigQuerySource",

"FileSource",

"RedshiftSource",

"SnowflakeSource",

"SparkSource",

"TrinoSource",

"AthenaSource",

}

class BatchFeatureView(FeatureView):

"""

A batch feature view defines a logical group of features that has only a batch data source.

Attributes:

name: The unique name of the batch feature view.

entities: List of entities or entity join keys.

ttl: The amount of time this group of features lives. A ttl of 0 indicates that

this group of features lives forever. Note that large ttl's or a ttl of 0

can result in extremely computationally intensive queries.

schema: The schema of the feature view, including feature, timestamp, and entity

columns. If not specified, can be inferred from the underlying data source.

source: The batch source of data where this group of features is stored.

online: A boolean indicating whether online retrieval and write to online store is enabled for this feature view.

offline: A boolean indicating whether offline retrieval and write to offline store is enabled for this feature view.

description: A human-readable description.

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

owner: The owner of the batch feature view, typically the email of the primary maintainer.

"""

name: str

mode: Union[TransformationMode, str]

entities: List[str]

ttl: Optional[timedelta]

source: DataSource

schema: List[Field]

entity_columns: List[Field]

features: List[Field]

online: bool

offline: bool

description: str

tags: Dict[str, str]

owner: str

timestamp_field: str

materialization_intervals: List[Tuple[datetime, datetime]]

udf: Optional[Callable[[Any], Any]]

udf_string: Optional[str]

feature_transformation: Transformation

batch_engine: Optional[Field]

aggregations: Optional[List[Aggregation]]

def __init__(

self,

*,

name: str,

mode: Union[TransformationMode, str] = TransformationMode.PYTHON,

source: DataSource,

entities: Optional[List[Entity]] = None,

ttl: Optional[timedelta] = None,

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

online: bool = False,

offline: bool = True,

description: str = "",

owner: str = "",

schema: Optional[List[Field]] = None,

udf: Optional[Callable[[Any], Any]],

udf_string: Optional[str] = "",

feature_transformation: Optional[Transformation] = None,

batch_engine: Optional[Field] = None,

aggregations: Optional[List[Aggregation]] = None,

):

if not flags_helper.is_test():

warnings.warn(

"Batch feature views are experimental features in alpha development. "

"Some functionality may still be unstable so functionality can change in the future.",

RuntimeWarning,

)

if (

type(source).__name__ not in SUPPORTED_BATCH_SOURCES

and source.to_proto().type != DataSourceProto.SourceType.CUSTOM_SOURCE

):

raise ValueError(

f"Batch feature views need a batch source, expected one of {SUPPORTED_BATCH_SOURCES} "

f"or CUSTOM_SOURCE, got {type(source).__name__}: {source.name} instead "

)

self.mode = mode

self.udf = udf

self.udf_string = udf_string

self.feature_transformation = (

feature_transformation or self.get_feature_transformation()

)

self.batch_engine = batch_engine

self.aggregations = aggregations or []

super().__init__(

name=name,

entities=entities,

ttl=ttl,

tags=tags,

online=online,

offline=offline,

description=description,

owner=owner,

schema=schema,

source=source,

)

def get_feature_transformation(self) -> Transformation:

if not self.udf:

raise ValueError(

"Either a UDF or a feature transformation must be provided for BatchFeatureView"

)

if self.mode in (

TransformationMode.PANDAS,

TransformationMode.PYTHON,

TransformationMode.SQL,

) or self.mode in ("pandas", "python", "sql"):

return Transformation(

mode=self.mode, udf=self.udf, udf_string=self.udf_string or ""

)

else:

raise ValueError(

f"Unsupported transformation mode: {self.mode} for StreamFeatureView"

)

def batch_feature_view(

*,

name: Optional[str] = None,

mode: Union[TransformationMode, str] = TransformationMode.PYTHON,

entities: Optional[List[str]] = None,

ttl: Optional[timedelta] = None,

source: Optional[DataSource] = None,

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

online: bool = True,

offline: bool = True,

description: str = "",

owner: str = "",

schema: Optional[List[Field]] = None,

):

"""

Args:

name:

mode:

entities:

ttl:

source:

tags:

online:

offline:

description:

owner:

schema:

Returns:

"""

def mainify(obj):

# Needed to allow dill to properly serialize the udf. Otherwise, clients will need to have a file with the same

# name as the original file defining the sfv.

if obj.__module__ != "__main__":

obj.__module__ = "__main__"

def decorator(user_function):

udf_string = dill.source.getsource(user_function)

mainify(user_function)

batch_feature_view_obj = BatchFeatureView(

name=name or user_function.__name__,

mode=mode,

entities=entities,

ttl=ttl,

source=source,

tags=tags,

online=online,

offline=offline,

description=description,

owner=owner,

schema=schema,

udf=user_function,

udf_string=udf_string,

)

functools.update_wrapper(wrapper=batch_feature_view_obj, wrapped=user_function)

return batch_feature_view_obj

return decorator