# This is an example feature definition file

from datetime import timedelta

import pandas as pd

from feast import (
    Entity,
    FeatureService,
    FeatureView,
    Field,
    PushSource,
    RedshiftSource,
    RequestSource,
)
from feast.on_demand_feature_view import on_demand_feature_view
from feast.types import Float32, Float64, Int64

# Define an entity for the driver. You can think of an entity as a primary key used to
# fetch features.
driver = Entity(name="driver", join_keys=["driver_id"])

# Defines a data source from which feature values can be retrieved. Sources are queried when building training
# datasets or materializing features into an online store.
driver_stats_source = RedshiftSource(
    # The Redshift table where features can be found
    table="feast_driver_hourly_stats",
    # The event timestamp is used for point-in-time joins and for ensuring only
    # features within the TTL are returned
    timestamp_field="event_timestamp",
    # The (optional) created timestamp is used to ensure there are no duplicate
    # feature rows in the offline store or when building training datasets
    created_timestamp_column="created",
    # Database to redshift source.
    database="%REDSHIFT_DATABASE%",
)

# Our parquet files contain sample data that includes a driver_id column, timestamps and
# three feature column. Here we define a Feature View that will allow us to serve this
# data to our model online.
driver_stats_fv = FeatureView(
    # The unique name of this feature view. Two feature views in a single
    # project cannot have the same name
    name="driver_hourly_stats",
    entities=[driver],
    ttl=timedelta(days=1),
    # The list of features defined below act as a schema to both define features
    # for both materialization of features into a store, and are used as references
    # during retrieval for building a training dataset or serving features
    schema=[
        Field(name="conv_rate", dtype=Float32),
        Field(name="acc_rate", dtype=Float32),
        Field(name="avg_daily_trips", dtype=Int64),
    ],
    online=True,
    source=driver_stats_source,
    # Tags are user defined key/value pairs that are attached to each
    # feature view
    tags={"team": "driver_performance"},
    version="latest",
)

# Define a request data source which encodes features / information only
# available at request time (e.g. part of the user initiated HTTP request)
input_request = RequestSource(
    name="vals_to_add",
    schema=[
        Field(name="val_to_add", dtype=Int64),
        Field(name="val_to_add_2", dtype=Int64),
    ],
)

# Define an on demand feature view which can generate new features based on
# existing feature views and RequestSource features
@on_demand_feature_view(
    sources=[driver_stats_fv, input_request],
    schema=[
        Field(name="conv_rate_plus_val1", dtype=Float64),
        Field(name="conv_rate_plus_val2", dtype=Float64),
    ],
)
def transformed_conv_rate(inputs: pd.DataFrame) -> pd.DataFrame:
    df = pd.DataFrame()
    df["conv_rate_plus_val1"] = inputs["conv_rate"] + inputs["val_to_add"]
    df["conv_rate_plus_val2"] = inputs["conv_rate"] + inputs["val_to_add_2"]
    return df

# This groups features into a model version
driver_activity_v1 = FeatureService(
    name="driver_activity_v1",
    features=[
        driver_stats_fv[["conv_rate"]],  # Sub-selects a feature from a feature view
        transformed_conv_rate,  # Selects all features from the feature view
    ],
)
driver_activity_v2 = FeatureService(
    name="driver_activity_v2", features=[driver_stats_fv, transformed_conv_rate]
)

# Defines a way to push data (to be available offline, online or both) into Feast.
driver_stats_push_source = PushSource(
    name="driver_stats_push_source",
    batch_source=driver_stats_source,
)

# Defines a slightly modified version of the feature view from above, where the source
# has been changed to the push source. This allows fresh features to be directly pushed
# to the online store for this feature view.
driver_stats_fresh_fv = FeatureView(
    name="driver_hourly_stats_fresh",
    entities=[driver],
    ttl=timedelta(days=1),
    schema=[
        Field(name="conv_rate", dtype=Float32),
        Field(name="acc_rate", dtype=Float32),
        Field(name="avg_daily_trips", dtype=Int64),
    ],
    online=True,
    source=driver_stats_push_source,  # Changed from above
    tags={"team": "driver_performance"},
    version="latest",
)

# Define an on demand feature view which can generate new features based on
# existing feature views and RequestSource features
@on_demand_feature_view(
    sources=[driver_stats_fresh_fv, input_request],  # relies on fresh version of FV
    schema=[
        Field(name="conv_rate_plus_val1", dtype=Float64),
        Field(name="conv_rate_plus_val2", dtype=Float64),
    ],
)
def transformed_conv_rate_fresh(inputs: pd.DataFrame) -> pd.DataFrame:
    df = pd.DataFrame()
    df["conv_rate_plus_val1"] = inputs["conv_rate"] + inputs["val_to_add"]
    df["conv_rate_plus_val2"] = inputs["conv_rate"] + inputs["val_to_add_2"]
    return df

driver_activity_v3 = FeatureService(
    name="driver_activity_v3",
    features=[driver_stats_fresh_fv, transformed_conv_rate_fresh],
)


# This is an example feature definition file

from datetime import timedelta

import pandas as pd

from feast import (
    Entity,
    FeatureService,
    FeatureView,
    Field,
    FileSource,
    PushSource,
    RequestSource,
)
from feast.on_demand_feature_view import on_demand_feature_view
from feast.types import Float32, Float64, Int64

# Define an entity for the driver. You can think of an entity as a primary key used to
# fetch features.
driver = Entity(name="driver", join_keys=["driver_id"])

# Read data from parquet files. Parquet is convenient for local development mode. For
# production, you can use your favorite DWH, such as BigQuery. See Feast documentation
# for more info.
driver_stats_source = FileSource(
    name="driver_hourly_stats_source",
    path="%PARQUET_PATH%",
    timestamp_field="event_timestamp",
    created_timestamp_column="created",
)

# Our parquet files contain sample data that includes a driver_id column, timestamps and
# three feature column. Here we define a Feature View that will allow us to serve this
# data to our model online.
driver_stats_fv = FeatureView(
    # The unique name of this feature view. Two feature views in a single
    # project cannot have the same name
    name="driver_hourly_stats",
    entities=[driver],
    ttl=timedelta(days=1),
    # The list of features defined below act as a schema to both define features
    # for both materialization of features into a store, and are used as references
    # during retrieval for building a training dataset or serving features
    schema=[
        Field(name="conv_rate", dtype=Float32),
        Field(name="acc_rate", dtype=Float32),
        Field(name="avg_daily_trips", dtype=Int64),
    ],
    online=True,
    source=driver_stats_source,
    # Tags are user defined key/value pairs that are attached to each
    # feature view
    tags={"team": "driver_performance"},
    version="latest",
)

# Define a request data source which encodes features / information only
# available at request time (e.g. part of the user initiated HTTP request)
input_request = RequestSource(
    name="vals_to_add",
    schema=[
        Field(name="val_to_add", dtype=Int64),
        Field(name="val_to_add_2", dtype=Int64),
    ],
)

# Define an on demand feature view which can generate new features based on
# existing feature views and RequestSource features
@on_demand_feature_view(
    sources=[driver_stats_fv, input_request],
    schema=[
        Field(name="conv_rate_plus_val1", dtype=Float64),
        Field(name="conv_rate_plus_val2", dtype=Float64),
    ],
)
def transformed_conv_rate(inputs: pd.DataFrame) -> pd.DataFrame:
    df = pd.DataFrame()
    df["conv_rate_plus_val1"] = inputs["conv_rate"] + inputs["val_to_add"]
    df["conv_rate_plus_val2"] = inputs["conv_rate"] + inputs["val_to_add_2"]
    return df

# This groups features into a model version
driver_activity_v1 = FeatureService(
    name="driver_activity_v1",
    features=[
        driver_stats_fv[["conv_rate"]],  # Sub-selects a feature from a feature view
        transformed_conv_rate,  # Selects all features from the feature view
    ],
)
driver_activity_v2 = FeatureService(
    name="driver_activity_v2", features=[driver_stats_fv, transformed_conv_rate]
)

# Defines a way to push data (to be available offline, online or both) into Feast.
driver_stats_push_source = PushSource(
    name="driver_stats_push_source",
    batch_source=driver_stats_source,
)

# Defines a slightly modified version of the feature view from above, where the source
# has been changed to the push source. This allows fresh features to be directly pushed
# to the online store for this feature view.
driver_stats_fresh_fv = FeatureView(
    name="driver_hourly_stats_fresh",
    entities=[driver],
    ttl=timedelta(days=1),
    schema=[
        Field(name="conv_rate", dtype=Float32),
        Field(name="acc_rate", dtype=Float32),
        Field(name="avg_daily_trips", dtype=Int64),
    ],
    online=True,
    source=driver_stats_push_source,  # Changed from above
    tags={"team": "driver_performance"},
    version="latest",
)

# Define an on demand feature view which can generate new features based on
# existing feature views and RequestSource features
@on_demand_feature_view(
    sources=[driver_stats_fresh_fv, input_request],  # relies on fresh version of FV
    schema=[
        Field(name="conv_rate_plus_val1", dtype=Float64),
        Field(name="conv_rate_plus_val2", dtype=Float64),
    ],
)
def transformed_conv_rate_fresh(inputs: pd.DataFrame) -> pd.DataFrame:
    df = pd.DataFrame()
    df["conv_rate_plus_val1"] = inputs["conv_rate"] + inputs["val_to_add"]
    df["conv_rate_plus_val2"] = inputs["conv_rate"] + inputs["val_to_add_2"]
    return df

driver_activity_v3 = FeatureService(
    name="driver_activity_v3",
    features=[driver_stats_fresh_fv, transformed_conv_rate_fresh],
)


# This is an example feature definition file

from datetime import timedelta

import pandas as pd

from feast import Entity, FeatureService, FeatureView, Field, PushSource, RequestSource
from feast.infra.offline_stores.contrib.couchbase_offline_store.couchbase_source import (
    CouchbaseColumnarSource,
)
from feast.on_demand_feature_view import on_demand_feature_view
from feast.types import Float32, Float64, Int64

# Define an entity for the driver. You can think of an entity as a primary key used to
# fetch features.
driver = Entity(name="driver", join_keys=["driver_id"])

driver_stats_source = CouchbaseColumnarSource(
    name="driver_hourly_stats_source",
    query="SELECT * FROM Default.Default.`feast_driver_hourly_stats`",
    database="Default",
    scope="Default",
    collection="feast_driver_hourly_stats",
    timestamp_field="event_timestamp",
    created_timestamp_column="created",
)

# Our parquet files contain sample data that includes a driver_id column, timestamps and
# three feature column. Here we define a Feature View that will allow us to serve this
# data to our model online.
driver_stats_fv = FeatureView(
    # The unique name of this feature view. Two feature views in a single
    # project cannot have the same name
    name="driver_hourly_stats",
    entities=[driver],
    ttl=timedelta(days=1),
    # The list of features defined below act as a schema to both define features
    # for both materialization of features into a store, and are used as references
    # during retrieval for building a training dataset or serving features
    schema=[
        Field(name="conv_rate", dtype=Float32),
        Field(name="acc_rate", dtype=Float32),
        Field(name="avg_daily_trips", dtype=Int64),
    ],
    online=True,
    source=driver_stats_source,
    # Tags are user defined key/value pairs that are attached to each
    # feature view
    tags={"team": "driver_performance"},
    version="latest",
)

# Define a request data source which encodes features / information only
# available at request time (e.g. part of the user initiated HTTP request)
input_request = RequestSource(
    name="vals_to_add",
    schema=[
        Field(name="val_to_add", dtype=Int64),
        Field(name="val_to_add_2", dtype=Int64),
    ],
)

# Define an on demand feature view which can generate new features based on
# existing feature views and RequestSource features
@on_demand_feature_view(
    sources=[driver_stats_fv, input_request],
    schema=[
        Field(name="conv_rate_plus_val1", dtype=Float64),
        Field(name="conv_rate_plus_val2", dtype=Float64),
    ],
)
def transformed_conv_rate(inputs: pd.DataFrame) -> pd.DataFrame:
    df = pd.DataFrame()
    df["conv_rate_plus_val1"] = inputs["conv_rate"] + inputs["val_to_add"]
    df["conv_rate_plus_val2"] = inputs["conv_rate"] + inputs["val_to_add_2"]
    return df

# This groups features into a model version
driver_activity_v1 = FeatureService(
    name="driver_activity_v1",
    features=[
        driver_stats_fv[["conv_rate"]],  # Sub-selects a feature from a feature view
        transformed_conv_rate,  # Selects all features from the feature view
    ],
)
driver_activity_v2 = FeatureService(
    name="driver_activity_v2", features=[driver_stats_fv, transformed_conv_rate]
)

# Defines a way to push data (to be available offline, online or both) into Feast.
driver_stats_push_source = PushSource(
    name="driver_stats_push_source",
    batch_source=driver_stats_source,
)

# Defines a slightly modified version of the feature view from above, where the source
# has been changed to the push source. This allows fresh features to be directly pushed
# to the online store for this feature view.
driver_stats_fresh_fv = FeatureView(
    name="driver_hourly_stats_fresh",
    entities=[driver],
    ttl=timedelta(days=1),
    schema=[
        Field(name="conv_rate", dtype=Float32),
        Field(name="acc_rate", dtype=Float32),
        Field(name="avg_daily_trips", dtype=Int64),
    ],
    online=True,
    source=driver_stats_push_source,  # Changed from above
    tags={"team": "driver_performance"},
    version="latest",
)

# Define an on demand feature view which can generate new features based on
# existing feature views and RequestSource features
@on_demand_feature_view(
    sources=[driver_stats_fresh_fv, input_request],  # relies on fresh version of FV
    schema=[
        Field(name="conv_rate_plus_val1", dtype=Float64),
        Field(name="conv_rate_plus_val2", dtype=Float64),
    ],
)
def transformed_conv_rate_fresh(inputs: pd.DataFrame) -> pd.DataFrame:
    df = pd.DataFrame()
    df["conv_rate_plus_val1"] = inputs["conv_rate"] + inputs["val_to_add"]
    df["conv_rate_plus_val2"] = inputs["conv_rate"] + inputs["val_to_add_2"]
    return df

driver_activity_v3 = FeatureService(
    name="driver_activity_v3",
    features=[driver_stats_fresh_fv, transformed_conv_rate_fresh],
)


from datetime import timedelta

import pandas as pd

from feast import (
    BigQuerySource,
    Entity,
    FeatureService,
    FeatureView,
    Field,
    PushSource,
    RequestSource,
)
from feast.on_demand_feature_view import on_demand_feature_view
from feast.types import Float32, Float64, Int64

# Define an entity for the driver. You can think of an entity as a primary key used to
# fetch features.
driver = Entity(name="driver", join_keys=["driver_id"])

# Defines a data source from which feature values can be retrieved. Sources are queried when building training
# datasets or materializing features into an online store.
driver_stats_source = BigQuerySource(
    name="driver_hourly_stats_source",
    # The BigQuery table where features can be found
    table="feast-oss.demo_data.driver_hourly_stats_2",
    # The event timestamp is used for point-in-time joins and for ensuring only
    # features within the TTL are returned
    timestamp_field="event_timestamp",
    # The (optional) created timestamp is used to ensure there are no duplicate
    # feature rows in the offline store or when building training datasets
    created_timestamp_column="created",
)

# Feature views are a grouping based on how features are stored in either the
# online or offline store.
driver_stats_fv = FeatureView(
    # The unique name of this feature view. Two feature views in a single
    # project cannot have the same name
    name="driver_hourly_stats",
    # The list of entities specifies the keys required for joining or looking
    # up features from this feature view. The reference provided in this field
    # correspond to the name of a defined entity (or entities)
    entities=[driver],
    # The timedelta is the maximum age that each feature value may have
    # relative to its lookup time. For historical features (used in training),
    # TTL is relative to each timestamp provided in the entity dataframe.
    # TTL also allows for eviction of keys from online stores and limits the
    # amount of historical scanning required for historical feature values
    # during retrieval
    ttl=timedelta(weeks=52 * 10),  # Set to be very long for example purposes only
    # The list of features defined below act as a schema to both define features
    # for both materialization of features into a store, and are used as references
    # during retrieval for building a training dataset or serving features
    schema=[
        Field(name="conv_rate", dtype=Float32),
        Field(name="acc_rate", dtype=Float32),
        Field(name="avg_daily_trips", dtype=Int64),
    ],
    source=driver_stats_source,
    # Tags are user defined key/value pairs that are attached to each
    # feature view
    tags={"team": "driver_performance"},
    version="latest",
)

# Define a request data source which encodes features / information only
# available at request time (e.g. part of the user initiated HTTP request)
input_request = RequestSource(
    name="vals_to_add",
    schema=[
        Field(name="val_to_add", dtype=Int64),
        Field(name="val_to_add_2", dtype=Int64),
    ],
)

# Define an on demand feature view which can generate new features based on
# existing feature views and RequestSource features
@on_demand_feature_view(
    sources=[driver_stats_fv, input_request],
    schema=[
        Field(name="conv_rate_plus_val1", dtype=Float64),
        Field(name="conv_rate_plus_val2", dtype=Float64),
    ],
)
def transformed_conv_rate(inputs: pd.DataFrame) -> pd.DataFrame:
    df = pd.DataFrame()
    df["conv_rate_plus_val1"] = inputs["conv_rate"] + inputs["val_to_add"]
    df["conv_rate_plus_val2"] = inputs["conv_rate"] + inputs["val_to_add_2"]
    return df

# This groups features into a model version
driver_activity_v1 = FeatureService(
    name="driver_activity_v1",
    features=[
        driver_stats_fv[["conv_rate"]],  # Sub-selects a feature from a feature view
        transformed_conv_rate,  # Selects all features from the feature view
    ],
)
driver_activity_v2 = FeatureService(
    name="driver_activity_v2", features=[driver_stats_fv, transformed_conv_rate]
)

# Defines a way to push data (to be available offline, online or both) into Feast.
driver_stats_push_source = PushSource(
    name="driver_stats_push_source",
    batch_source=driver_stats_source,
)

# Defines a slightly modified version of the feature view from above, where the source
# has been changed to the push source. This allows fresh features to be directly pushed
# to the online store for this feature view.
driver_stats_fresh_fv = FeatureView(
    name="driver_hourly_stats_fresh",
    entities=[driver],
    ttl=timedelta(weeks=52 * 10),  # Set to be very long for example purposes only
    schema=[
        Field(name="conv_rate", dtype=Float32),
        Field(name="acc_rate", dtype=Float32),
        Field(name="avg_daily_trips", dtype=Int64),
    ],
    online=True,
    source=driver_stats_push_source,  # Changed from above
    tags={"team": "driver_performance"},
    version="latest",
)

# Define an on demand feature view which can generate new features based on
# existing feature views and RequestSource features
@on_demand_feature_view(
    sources=[driver_stats_fresh_fv, input_request],  # relies on fresh version of FV
    schema=[
        Field(name="conv_rate_plus_val1", dtype=Float64),
        Field(name="conv_rate_plus_val2", dtype=Float64),
    ],
)
def transformed_conv_rate_fresh(inputs: pd.DataFrame) -> pd.DataFrame:
    df = pd.DataFrame()
    df["conv_rate_plus_val1"] = inputs["conv_rate"] + inputs["val_to_add"]
    df["conv_rate_plus_val2"] = inputs["conv_rate"] + inputs["val_to_add_2"]
    return df

driver_activity_v3 = FeatureService(
    name="driver_activity_v3",
    features=[driver_stats_fresh_fv, transformed_conv_rate_fresh],
)


# This is an example feature definition file

from datetime import timedelta

import pandas as pd

from feast import (
    Entity,
    FeatureService,
    FeatureView,
    Field,
    FileSource,
    Project,
    PushSource,
    RequestSource,
)
from feast.feature_logging import LoggingConfig
from feast.infra.offline_stores.file_source import FileLoggingDestination
from feast.on_demand_feature_view import on_demand_feature_view
from feast.types import Float32, Float64, Int64, Json, Map, String, Struct

# Define a project for the feature repo
project = Project(name="%PROJECT_NAME%", description="A project for driver statistics")

# Define an entity for the driver. You can think of an entity as a primary key used to
# fetch features.
driver = Entity(name="driver", join_keys=["driver_id"])

# Read data from parquet files. Parquet is convenient for local development mode. For
# production, you can use your favorite DWH, such as BigQuery. See Feast documentation
# for more info.
driver_stats_source = FileSource(
    name="driver_hourly_stats_source",
    path="%PARQUET_PATH%",
    timestamp_field="event_timestamp",
    created_timestamp_column="created",
)

# Our parquet files contain sample data that includes a driver_id column, timestamps and
# three feature column. Here we define a Feature View that will allow us to serve this
# data to our model online.
driver_stats_fv = FeatureView(
    # The unique name of this feature view. Two feature views in a single
    # project cannot have the same name
    name="driver_hourly_stats",
    entities=[driver],
    ttl=timedelta(days=1),
    # The list of features defined below act as a schema to both define features
    # for both materialization of features into a store, and are used as references
    # during retrieval for building a training dataset or serving features
    schema=[
        Field(name="conv_rate", dtype=Float32),
        Field(name="acc_rate", dtype=Float32),
        Field(name="avg_daily_trips", dtype=Int64, description="Average daily trips"),
        Field(
            name="driver_metadata",
            dtype=Map,
            description="Driver metadata as key-value pairs",
        ),
        Field(
            name="driver_config", dtype=Json, description="Driver configuration as JSON"
        ),
        Field(
            name="driver_profile",
            dtype=Struct({"name": String, "age": String}),
            description="Driver profile as a typed struct",
        ),
    ],
    online=True,
    source=driver_stats_source,
    # Tags are user defined key/value pairs that are attached to each
    # feature view
    tags={"team": "driver_performance"},
    enable_validation=True,
    version="latest",
)

# Define a request data source which encodes features / information only
# available at request time (e.g. part of the user initiated HTTP request)
input_request = RequestSource(
    name="vals_to_add",
    schema=[
        Field(name="val_to_add", dtype=Int64),
        Field(name="val_to_add_2", dtype=Int64),
    ],
)

# Define an on demand feature view which can generate new features based on
# existing feature views and RequestSource features
@on_demand_feature_view(
    sources=[driver_stats_fv, input_request],
    schema=[
        Field(name="conv_rate_plus_val1", dtype=Float64),
        Field(name="conv_rate_plus_val2", dtype=Float64),
    ],
)
def transformed_conv_rate(inputs: pd.DataFrame) -> pd.DataFrame:
    df = pd.DataFrame()
    df["conv_rate_plus_val1"] = inputs["conv_rate"] + inputs["val_to_add"]
    df["conv_rate_plus_val2"] = inputs["conv_rate"] + inputs["val_to_add_2"]
    return df

# This groups features into a model version
driver_activity_v1 = FeatureService(
    name="driver_activity_v1",
    features=[
        driver_stats_fv[["conv_rate"]],  # Sub-selects a feature from a feature view
        transformed_conv_rate,  # Selects all features from the feature view
    ],
    logging_config=LoggingConfig(
        destination=FileLoggingDestination(path="%LOGGING_PATH%")
    ),
)
driver_activity_v2 = FeatureService(
    name="driver_activity_v2", features=[driver_stats_fv, transformed_conv_rate]
)

# Defines a way to push data (to be available offline, online or both) into Feast.
driver_stats_push_source = PushSource(
    name="driver_stats_push_source",
    batch_source=driver_stats_source,
)

# Defines a slightly modified version of the feature view from above, where the source
# has been changed to the push source. This allows fresh features to be directly pushed
# to the online store for this feature view.
driver_stats_fresh_fv = FeatureView(
    name="driver_hourly_stats_fresh",
    entities=[driver],
    ttl=timedelta(days=1),
    schema=[
        Field(name="conv_rate", dtype=Float32),
        Field(name="acc_rate", dtype=Float32),
        Field(name="avg_daily_trips", dtype=Int64),
        Field(name="driver_metadata", dtype=Map),
        Field(name="driver_config", dtype=Json),
        Field(name="driver_profile", dtype=Struct({"name": String, "age": String})),
    ],
    online=True,
    source=driver_stats_push_source,  # Changed from above
    tags={"team": "driver_performance"},
    version="latest",
)

# Define an on demand feature view which can generate new features based on
# existing feature views and RequestSource features
@on_demand_feature_view(
    sources=[driver_stats_fresh_fv, input_request],  # relies on fresh version of FV
    schema=[
        Field(name="conv_rate_plus_val1", dtype=Float64),
        Field(name="conv_rate_plus_val2", dtype=Float64),
    ],
)
def transformed_conv_rate_fresh(inputs: pd.DataFrame) -> pd.DataFrame:
    df = pd.DataFrame()
    df["conv_rate_plus_val1"] = inputs["conv_rate"] + inputs["val_to_add"]
    df["conv_rate_plus_val2"] = inputs["conv_rate"] + inputs["val_to_add_2"]
    return df

driver_activity_v3 = FeatureService(
    name="driver_activity_v3",
    features=[driver_stats_fresh_fv, transformed_conv_rate_fresh],
)


# This is an example feature definition file

from datetime import timedelta

import pandas as pd

from feast import (
    Entity,
    FeatureService,
    FeatureView,
    Field,
    FileSource,
    Project,
    PushSource,
    RequestSource,
)
from feast.feature_logging import LoggingConfig
from feast.infra.offline_stores.file_source import FileLoggingDestination
from feast.on_demand_feature_view import on_demand_feature_view
from feast.types import Float32, Float64, Int64

# Define a project for the feature repo
project = Project(name="%PROJECT_NAME%", description="A project for driver statistics")

# Define an entity for the driver. You can think of an entity as a primary key used to
# fetch features.
driver = Entity(name="driver", join_keys=["driver_id"])

# Read data from parquet files. Parquet is convenient for local development mode. For
# production, you can use your favorite DWH, such as BigQuery. See Feast documentation
# for more info.
driver_stats_source = FileSource(
    name="driver_hourly_stats_source",
    path="%PARQUET_PATH%",
    timestamp_field="event_timestamp",
    created_timestamp_column="created",
)

# Our parquet files contain sample data that includes a driver_id column, timestamps and
# three feature column. Here we define a Feature View that will allow us to serve this
# data to our model online.
driver_stats_fv = FeatureView(
    # The unique name of this feature view. Two feature views in a single
    # project cannot have the same name
    name="driver_hourly_stats",
    entities=[driver],
    ttl=timedelta(days=1),
    # The list of features defined below act as a schema to both define features
    # for both materialization of features into a store, and are used as references
    # during retrieval for building a training dataset or serving features
    schema=[
        Field(name="conv_rate", dtype=Float32),
        Field(name="acc_rate", dtype=Float32),
        Field(name="avg_daily_trips", dtype=Int64, description="Average daily trips"),
    ],
    online=True,
    source=driver_stats_source,
    # Tags are user defined key/value pairs that are attached to each
    # feature view
    tags={"team": "driver_performance"},
    version="latest",
)

# Define a request data source which encodes features / information only
# available at request time (e.g. part of the user initiated HTTP request)
input_request = RequestSource(
    name="vals_to_add",
    schema=[
        Field(name="val_to_add", dtype=Int64),
        Field(name="val_to_add_2", dtype=Int64),
    ],
)

# Define an on demand feature view which can generate new features based on
# existing feature views and RequestSource features
@on_demand_feature_view(
    sources=[driver_stats_fv, input_request],
    schema=[
        Field(name="conv_rate_plus_val1", dtype=Float64),
        Field(name="conv_rate_plus_val2", dtype=Float64),
    ],
)
def transformed_conv_rate(inputs: pd.DataFrame) -> pd.DataFrame:
    df = pd.DataFrame()
    df["conv_rate_plus_val1"] = inputs["conv_rate"] + inputs["val_to_add"]
    df["conv_rate_plus_val2"] = inputs["conv_rate"] + inputs["val_to_add_2"]
    return df

# This groups features into a model version
driver_activity_v1 = FeatureService(
    name="driver_activity_v1",
    features=[
        driver_stats_fv[["conv_rate"]],  # Sub-selects a feature from a feature view
        transformed_conv_rate,  # Selects all features from the feature view
    ],
    logging_config=LoggingConfig(
        destination=FileLoggingDestination(path="%LOGGING_PATH%")
    ),
)
driver_activity_v2 = FeatureService(
    name="driver_activity_v2", features=[driver_stats_fv, transformed_conv_rate]
)

# Defines a way to push data (to be available offline, online or both) into Feast.
driver_stats_push_source = PushSource(
    name="driver_stats_push_source",
    batch_source=driver_stats_source,
)

# Defines a slightly modified version of the feature view from above, where the source
# has been changed to the push source. This allows fresh features to be directly pushed
# to the online store for this feature view.
driver_stats_fresh_fv = FeatureView(
    name="driver_hourly_stats_fresh",
    entities=[driver],
    ttl=timedelta(days=1),
    schema=[
        Field(name="conv_rate", dtype=Float32),
        Field(name="acc_rate", dtype=Float32),
        Field(name="avg_daily_trips", dtype=Int64),
    ],
    online=True,
    source=driver_stats_push_source,  # Changed from above
    tags={"team": "driver_performance"},
    version="latest",
)

# Define an on demand feature view which can generate new features based on
# existing feature views and RequestSource features
@on_demand_feature_view(
    sources=[driver_stats_fresh_fv, input_request],  # relies on fresh version of FV
    schema=[
        Field(name="conv_rate_plus_val1", dtype=Float64),
        Field(name="conv_rate_plus_val2", dtype=Float64),
    ],
)
def transformed_conv_rate_fresh(inputs: pd.DataFrame) -> pd.DataFrame:
    df = pd.DataFrame()
    df["conv_rate_plus_val1"] = inputs["conv_rate"] + inputs["val_to_add"]
    df["conv_rate_plus_val2"] = inputs["conv_rate"] + inputs["val_to_add_2"]
    return df

driver_activity_v3 = FeatureService(
    name="driver_activity_v3",
    features=[driver_stats_fresh_fv, transformed_conv_rate_fresh],
)


# This is an example feature definition file

from datetime import timedelta

import pandas as pd

from feast import Entity, FeatureService, FeatureView, Field, PushSource, RequestSource
from feast.infra.offline_stores.contrib.postgres_offline_store.postgres_source import (
    PostgreSQLSource,
)
from feast.on_demand_feature_view import on_demand_feature_view
from feast.types import Float32, Float64, Int64

# Define an entity for the driver. You can think of an entity as a primary key used to
# fetch features.
driver = Entity(name="driver", join_keys=["driver_id"])

driver_stats_source = PostgreSQLSource(
    name="driver_hourly_stats_source",
    query="SELECT * FROM feast_driver_hourly_stats",
    timestamp_field="event_timestamp",
    created_timestamp_column="created",
)

# Our parquet files contain sample data that includes a driver_id column, timestamps and
# three feature column. Here we define a Feature View that will allow us to serve this
# data to our model online.
driver_stats_fv = FeatureView(
    # The unique name of this feature view. Two feature views in a single
    # project cannot have the same name
    name="driver_hourly_stats",
    entities=[driver],
    ttl=timedelta(days=1),
    # The list of features defined below act as a schema to both define features
    # for both materialization of features into a store, and are used as references
    # during retrieval for building a training dataset or serving features
    schema=[
        Field(name="conv_rate", dtype=Float32),
        Field(name="acc_rate", dtype=Float32),
        Field(name="avg_daily_trips", dtype=Int64),
    ],
    online=True,
    source=driver_stats_source,
    # Tags are user defined key/value pairs that are attached to each
    # feature view
    tags={"team": "driver_performance"},
    version="latest",
)

# Define a request data source which encodes features / information only
# available at request time (e.g. part of the user initiated HTTP request)
input_request = RequestSource(
    name="vals_to_add",
    schema=[
        Field(name="val_to_add", dtype=Int64),
        Field(name="val_to_add_2", dtype=Int64),
    ],
)

# Define an on demand feature view which can generate new features based on
# existing feature views and RequestSource features
@on_demand_feature_view(
    sources=[driver_stats_fv, input_request],
    schema=[
        Field(name="conv_rate_plus_val1", dtype=Float64),
        Field(name="conv_rate_plus_val2", dtype=Float64),
    ],
)
def transformed_conv_rate(inputs: pd.DataFrame) -> pd.DataFrame:
    df = pd.DataFrame()
    df["conv_rate_plus_val1"] = inputs["conv_rate"] + inputs["val_to_add"]
    df["conv_rate_plus_val2"] = inputs["conv_rate"] + inputs["val_to_add_2"]
    return df

# This groups features into a model version
driver_activity_v1 = FeatureService(
    name="driver_activity_v1",
    features=[
        driver_stats_fv[["conv_rate"]],  # Sub-selects a feature from a feature view
        transformed_conv_rate,  # Selects all features from the feature view
    ],
)
driver_activity_v2 = FeatureService(
    name="driver_activity_v2", features=[driver_stats_fv, transformed_conv_rate]
)

# Defines a way to push data (to be available offline, online or both) into Feast.
driver_stats_push_source = PushSource(
    name="driver_stats_push_source",
    batch_source=driver_stats_source,
)

# Defines a slightly modified version of the feature view from above, where the source
# has been changed to the push source. This allows fresh features to be directly pushed
# to the online store for this feature view.
driver_stats_fresh_fv = FeatureView(
    name="driver_hourly_stats_fresh",
    entities=[driver],
    ttl=timedelta(days=1),
    schema=[
        Field(name="conv_rate", dtype=Float32),
        Field(name="acc_rate", dtype=Float32),
        Field(name="avg_daily_trips", dtype=Int64),
    ],
    online=True,
    source=driver_stats_push_source,  # Changed from above
    tags={"team": "driver_performance"},
    version="latest",
)

# Define an on demand feature view which can generate new features based on
# existing feature views and RequestSource features
@on_demand_feature_view(
    sources=[driver_stats_fresh_fv, input_request],  # relies on fresh version of FV
    schema=[
        Field(name="conv_rate_plus_val1", dtype=Float64),
        Field(name="conv_rate_plus_val2", dtype=Float64),
    ],
)
def transformed_conv_rate_fresh(inputs: pd.DataFrame) -> pd.DataFrame:
    df = pd.DataFrame()
    df["conv_rate_plus_val1"] = inputs["conv_rate"] + inputs["val_to_add"]
    df["conv_rate_plus_val2"] = inputs["conv_rate"] + inputs["val_to_add_2"]
    return df

driver_activity_v3 = FeatureService(
    name="driver_activity_v3",
    features=[driver_stats_fresh_fv, transformed_conv_rate_fresh],
)


"""
PyTorch NLP Sentiment Analysis Feature Repository

This template demonstrates sentiment analysis using:
- Text feature engineering for NLP
- PyTorch + Hugging Face transformers integration
- On-demand sentiment prediction features
- Online and offline feature serving patterns
"""

from datetime import timedelta
from pathlib import Path

import pandas as pd

from feast import (
    Entity,
    FeatureService,
    FeatureView,
    Field,
    FileSource,
    RequestSource,
    ValueType,
)
from feast.on_demand_feature_view import on_demand_feature_view
from feast.types import Array, Float32, Int64, String

try:
    # Import static artifacts helpers (available when feature server loads artifacts)
    from static_artifacts import get_lookup_tables, get_sentiment_model
except ImportError:
    # Fallback for when static_artifacts.py is not available
    get_sentiment_model = None
    get_lookup_tables = None

# Global references for static artifacts (set by feature server)
_sentiment_model = None
_lookup_tables: dict = {}

# Configuration
repo_path = Path(__file__).parent
data_path = repo_path / "data"

# Define entities - primary keys for joining data
text_entity = Entity(
    name="text",
    join_keys=["text_id"],
    value_type=ValueType.STRING,
    description="Unique identifier for text samples",
)

user_entity = Entity(
    name="user",
    join_keys=["user_id"],
    value_type=ValueType.STRING,
    description="User who created the text content",
)

# Data source - points to the parquet file created by bootstrap
sentiment_source = FileSource(
    name="sentiment_data_source",
    path=str(data_path / "sentiment_data.parquet"),
    timestamp_field="event_timestamp",
    created_timestamp_column="created",
)

# Feature view for text metadata and engineered features
text_features_fv = FeatureView(
    name="text_features",
    entities=[text_entity],
    ttl=timedelta(days=7),  # Keep features for 7 days
    schema=[
        Field(name="text_content", dtype=String, description="Raw text content"),
        Field(
            name="sentiment_label",
            dtype=String,
            description="Ground truth sentiment label",
        ),
        Field(
            name="sentiment_score",
            dtype=Float32,
            description="Ground truth sentiment score",
        ),
        Field(name="text_length", dtype=Int64, description="Character count of text"),
        Field(name="word_count", dtype=Int64, description="Word count of text"),
        Field(
            name="exclamation_count",
            dtype=Int64,
            description="Number of exclamation marks",
        ),
        Field(name="caps_ratio", dtype=Float32, description="Ratio of capital letters"),
        Field(
            name="emoji_count", dtype=Int64, description="Number of emoji characters"
        ),
    ],
    online=True,
    source=sentiment_source,
    tags={"team": "nlp", "domain": "sentiment_analysis"},
    version="latest",
)

# Feature view for user-level aggregations
user_stats_fv = FeatureView(
    name="user_stats",
    entities=[user_entity],
    ttl=timedelta(days=30),  # User stats change less frequently
    schema=[
        Field(
            name="user_avg_sentiment",
            dtype=Float32,
            description="User's average sentiment score",
        ),
        Field(
            name="user_text_count",
            dtype=Int64,
            description="Total number of texts by user",
        ),
        Field(
            name="user_avg_text_length",
            dtype=Float32,
            description="User's average text length",
        ),
    ],
    online=True,
    source=sentiment_source,
    tags={"team": "nlp", "domain": "user_behavior"},
    version="latest",
)

# Request source for real-time inference
text_input_request = RequestSource(
    name="text_input",
    schema=[
        Field(
            name="input_text",
            dtype=String,
            description="Text to analyze at request time",
        ),
        Field(
            name="model_name", dtype=String, description="Model to use for prediction"
        ),
    ],
)

# On-demand feature view for real-time sentiment prediction
@on_demand_feature_view(
    sources=[text_input_request],
    schema=[
        Field(name="predicted_sentiment", dtype=String),
        Field(name="sentiment_confidence", dtype=Float32),
        Field(name="positive_prob", dtype=Float32),
        Field(name="negative_prob", dtype=Float32),
        Field(name="neutral_prob", dtype=Float32),
        Field(name="text_embedding", dtype=Array(Float32)),
    ],
)
def sentiment_prediction(inputs: pd.DataFrame) -> pd.DataFrame:
    """
    Real-time sentiment prediction using pre-loaded static artifacts.

    This function demonstrates how to use static artifacts (pre-loaded models,
    lookup tables) for efficient real-time inference. Models are loaded once
    at feature server startup rather than on each request.
    """
    try:
        import numpy as np
    except ImportError:
        # Fallback to dummy predictions if numpy isn't available

        df = pd.DataFrame()
        df["predicted_sentiment"] = ["neutral"] * len(inputs)
        df["sentiment_confidence"] = [0.5] * len(inputs)
        df["positive_prob"] = [0.33] * len(inputs)
        df["negative_prob"] = [0.33] * len(inputs)
        df["neutral_prob"] = [0.34] * len(inputs)
        df["text_embedding"] = [[0.0] * 384] * len(inputs)
        return df

    # Get pre-loaded static artifacts from global references
    # These are loaded once at startup via static_artifacts.py
    global _sentiment_model, _lookup_tables

    sentiment_model = _sentiment_model
    lookup_tables = _lookup_tables

    # Use lookup table for label mapping (from static artifacts)
    label_map = lookup_tables.get(
        "sentiment_labels",
        {"LABEL_0": "negative", "LABEL_1": "neutral", "LABEL_2": "positive"},
    )

    results = []

    for text in inputs["input_text"]:
        try:
            if sentiment_model is not None:
                # Use pre-loaded model for prediction
                predictions = sentiment_model(text)

                # Parse results using static lookup tables
                scores = {
                    label_map.get(pred["label"], pred["label"]): pred["score"]
                    for pred in predictions
                }

                # Get best prediction
                best_pred = max(predictions, key=lambda x: x["score"])
                predicted_sentiment = label_map.get(
                    best_pred["label"], best_pred["label"]
                )
                confidence = best_pred["score"]
            else:
                # Fallback when model is not available
                predicted_sentiment = "neutral"
                confidence = 0.5
                scores = {"positive": 0.33, "negative": 0.33, "neutral": 0.34}

            # Generate dummy embeddings (in production, use pre-loaded embeddings)
            embedding = np.random.rand(384).tolist()

            results.append(
                {
                    "predicted_sentiment": predicted_sentiment,
                    "sentiment_confidence": np.float32(confidence),
                    "positive_prob": np.float32(scores.get("positive", 0.0)),
                    "negative_prob": np.float32(scores.get("negative", 0.0)),
                    "neutral_prob": np.float32(scores.get("neutral", 0.0)),
                    "text_embedding": [np.float32(x) for x in embedding],
                }
            )

        except Exception:
            # Fallback for individual text processing errors
            results.append(
                {
                    "predicted_sentiment": "neutral",
                    "sentiment_confidence": np.float32(0.5),
                    "positive_prob": np.float32(0.33),
                    "negative_prob": np.float32(0.33),
                    "neutral_prob": np.float32(0.34),
                    "text_embedding": [np.float32(0.0)] * 384,
                }
            )

    return pd.DataFrame(results)

# Feature services group related features for model serving
sentiment_analysis_v1 = FeatureService(
    name="sentiment_analysis_v1",
    features=[
        text_features_fv[["text_content", "text_length", "word_count"]],
        sentiment_prediction,
    ],
    description="Basic sentiment analysis features for model v1",
)

sentiment_analysis_v2 = FeatureService(
    name="sentiment_analysis_v2",
    features=[
        text_features_fv,  # All text features
        user_stats_fv[["user_avg_sentiment", "user_text_count"]],  # User context
        sentiment_prediction,  # Real-time predictions
    ],
    description="Advanced sentiment analysis with user context for model v2",
)

# Feature service for training data (historical features only)
sentiment_training_features = FeatureService(
    name="sentiment_training_features",
    features=[
        text_features_fv,
        user_stats_fv,
    ],
    description="Historical features for model training and evaluation",
)


# # # # # # # # # # # # # # # # # # # # # # # #
# This is an example feature definition file  #
# showcasing Ray offline store and compute     #
# engine capabilities                          #
# # # # # # # # # # # # # # # # # # # # # # # #

from datetime import timedelta
from pathlib import Path

from feast import Entity, FeatureService, FeatureView, Field, ValueType
from feast.infra.offline_stores.contrib.ray_offline_store.ray_source import RaySource
from feast.on_demand_feature_view import on_demand_feature_view
from feast.types import Float32, Float64, Int64

# Constants related to the generated data sets
CURRENT_DIR = Path(__file__).parent

# Entity definitions
driver = Entity(
    name="driver",
    description="driver id",
    value_type=ValueType.INT64,
    join_keys=["driver_id"],
)

customer = Entity(
    name="customer",
    description="customer id",
    value_type=ValueType.INT64,
    join_keys=["customer_id"],
)

# Data sources — RaySource is the recommended data source for the Ray offline
# store. It tells Feast how to load a Ray Dataset from any Ray Data-supported
# format (Parquet, CSV, JSON, HuggingFace datasets, MongoDB, SQL, and more).
# Here we read local Parquet files that are generated by bootstrap.py.
driver_hourly_stats = RaySource(
    name="driver_hourly_stats",
    reader_type="parquet",
    path=f"{CURRENT_DIR}/%PARQUET_PATH%",
    timestamp_field="event_timestamp",
    created_timestamp_column="created",
)

customer_daily_profile = RaySource(
    name="customer_daily_profile",
    reader_type="parquet",
    path=f"{CURRENT_DIR}/data/customer_daily_profile.parquet",
    timestamp_field="event_timestamp",
    created_timestamp_column="created",
)

# Feature Views - These leverage Ray compute engine for distributed processing
driver_hourly_stats_view = FeatureView(
    name="driver_hourly_stats",
    entities=[driver],
    ttl=timedelta(days=7),
    schema=[
        Field(name="conv_rate", dtype=Float32),
        Field(name="acc_rate", dtype=Float32),
        Field(name="avg_daily_trips", dtype=Int64),
    ],
    online=True,
    source=driver_hourly_stats,
    tags={"team": "driver_performance", "processing": "ray"},
    version="latest",
)

customer_daily_profile_view = FeatureView(
    name="customer_daily_profile",
    entities=[customer],
    ttl=timedelta(days=7),
    schema=[
        Field(name="current_balance", dtype=Float32),
        Field(name="avg_passenger_count", dtype=Float32),
        Field(name="lifetime_trip_count", dtype=Int64),
    ],
    online=True,
    source=customer_daily_profile,
    tags={"team": "customer_analytics", "processing": "ray"},
    version="latest",
)

# On-demand feature view showcasing Ray compute engine capabilities
# This demonstrates real-time feature transformations using Ray
@on_demand_feature_view(
    sources=[driver_hourly_stats_view],
    schema=[
        Field(name="conv_rate_plus_acc_rate", dtype=Float64),
        Field(name="trips_per_day_normalized", dtype=Float64),
    ],
)
def driver_activity_v2(inputs: dict):
    """
    On-demand feature transformations processed by Ray compute engine.
    These calculations happen in real-time and can leverage Ray's
    distributed processing capabilities.
    """
    import pandas as pd

    conv_rate = inputs["conv_rate"]
    acc_rate = inputs["acc_rate"]
    avg_daily_trips = inputs["avg_daily_trips"]

    # Feature engineering using Ray's distributed processing
    conv_rate_plus_acc_rate = conv_rate + acc_rate

    # Normalize trips per day (example of more complex transformation)
    max_trips = avg_daily_trips.max() if len(avg_daily_trips) > 0 else 1
    trips_per_day_normalized = avg_daily_trips / max_trips

    return pd.DataFrame(
        {
            "conv_rate_plus_acc_rate": conv_rate_plus_acc_rate,
            "trips_per_day_normalized": trips_per_day_normalized,
        }
    )

# Feature Service - Groups related features for serving
# Ray compute engine optimizes the retrieval of these feature combinations
driver_activity_v1 = FeatureService(
    name="driver_activity_v1",
    features=[
        driver_hourly_stats_view,
        customer_daily_profile_view,
    ],
    tags={"version": "v1", "compute_engine": "ray"},
)

driver_activity_v2_service = FeatureService(
    name="driver_activity_v2",
    features=[
        driver_hourly_stats_view,
        customer_daily_profile_view,
        driver_activity_v2,  # Includes on-demand transformations
    ],
    tags={"version": "v2", "compute_engine": "ray", "transformations": "on_demand"},
)


"""
RAG Feature Repository - Movie Embeddings with Ray Native Processing

This template demonstrates distributed embedding generation using:
- Ray offline store for scalable data I/O
- Ray compute engine for parallel processing
- Milvus online store with vector search capabilities
"""

from datetime import timedelta
from pathlib import Path

import pandas as pd

from feast import BatchFeatureView, Entity, Field, FileSource, ValueType
from feast.types import Array, Float32, String

# Configuration
repo_path = Path(__file__).parent
data_path = repo_path / "data"
data_path.mkdir(exist_ok=True)

EMBED_MODEL_ID = "sentence-transformers/all-MiniLM-L6-v2"

# Entity definition
document = Entity(
    name="document_id",
    join_keys=["document_id_pk"],
    value_type=ValueType.STRING,
    description="Document identifier for RAG retrieval",
)

# Data source
movies_source = FileSource(
    path=str(data_path / "raw_movies.parquet"),
    timestamp_field="DatePublished",
)

# Embedding processor for distributed Ray processing
class EmbeddingProcessor:
    """
    Generate embeddings using SentenceTransformer model.
    Model is loaded once per worker and reused for all batches.
    """

    def __init__(self):
        """Initialize model once per worker."""
        import torch
        from sentence_transformers import SentenceTransformer

        device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model = SentenceTransformer(EMBED_MODEL_ID, device=device)

    def __call__(self, batch: pd.DataFrame) -> pd.DataFrame:
        """Process batch and generate embeddings."""
        if "id" in batch.columns:
            batch["document_id"] = "movie_" + batch["id"].astype(str)
            batch["document_id_pk"] = batch["document_id"]

        # Generate embeddings from descriptions
        if "Description" in batch.columns:
            descriptions = batch["Description"].fillna("").tolist()
            model_batch_size = min(128, max(32, len(descriptions)))
            embeddings = self.model.encode(
                descriptions,
                show_progress_bar=False,
                batch_size=model_batch_size,
                normalize_embeddings=True,
                convert_to_numpy=True,
            )
            batch["embedding"] = embeddings.tolist()
            batch["embedding_model"] = EMBED_MODEL_ID

        # Standardize movie-related metadata
        if "Name" in batch.columns:
            batch["movie_name"] = batch["Name"].fillna("")
        if "Director" in batch.columns:
            batch["movie_director"] = batch["Director"].fillna("")
        if "Genres" in batch.columns:
            batch["movie_genres"] = batch["Genres"].fillna("")
        if "RatingValue" in batch.columns:
            batch["movie_rating"] = batch["RatingValue"].fillna(0.0)

        return batch

# Ray native UDF - Fully adaptive to cluster resources
def generate_embeddings_ray_native(ds):
    """
    Distributed embedding generation using Ray Data.
    """
    # Ray transformation mode providing control to user over the resources
    # Optimize the resources for the transformation
    num_blocks = ds.num_blocks()
    max_workers = 9
    batch_size = 2500
    try:
        sample = ds.take(1)
        has_data = len(sample) > 0
    except Exception:
        has_data = False
    if has_data and num_blocks < max_workers:
        ds = ds.repartition(max_workers)

    result = ds.map_batches(
        EmbeddingProcessor,
        batch_format="pandas",
        concurrency=max_workers,
        batch_size=batch_size,
    )
    return result

# Batch feature view for embeddings with metadata
document_embeddings_view = BatchFeatureView(
    name="document_embeddings",
    entities=[document],
    mode="ray",  # Native Ray Dataset mode for distributed processing
    ttl=timedelta(days=365 * 100),
    schema=[
        Field(name="document_id", dtype=String),
        Field(name="document_id_pk", dtype=String),
        Field(name="embedding", dtype=Array(Float32), vector_index=True),
        Field(name="embedding_model", dtype=String),
        Field(name="movie_name", dtype=String),
        Field(name="movie_director", dtype=String),
        Field(name="movie_genres", dtype=String),
        Field(name="movie_rating", dtype=Float32),
    ],
    source=movies_source,
    udf=generate_embeddings_ray_native,
    online=True,
    tags={
        "team": "ml_platform",
        "use_case": "rag",
        "transformation_mode": "ray_native",
    },
)


# # # # # # # # # # # # # # # # # # # # # # # #
# This is an example feature definition file  #
# # # # # # # # # # # # # # # # # # # # # # # #

from datetime import timedelta
from pathlib import Path

from feast import Entity, FeatureService, FeatureView, Field
from feast.infra.offline_stores.contrib.spark_offline_store.spark_source import (
    SparkSource,
)
from feast.types import Float32, Int64

# Constants related to the generated data sets
CURRENT_DIR = Path(__file__).parent

# Entity definitions
driver = Entity(
    name="driver",
    description="driver id",
)
customer = Entity(
    name="customer",
    description="customer id",
)

# Sources
driver_hourly_stats = SparkSource(
    name="driver_hourly_stats",
    path=f"{CURRENT_DIR}/data/driver_hourly_stats.parquet",
    file_format="parquet",
    timestamp_field="event_timestamp",
    created_timestamp_column="created",
)
customer_daily_profile = SparkSource(
    name="customer_daily_profile",
    path=f"{CURRENT_DIR}/data/customer_daily_profile.parquet",
    file_format="parquet",
    timestamp_field="event_timestamp",
    created_timestamp_column="created",
)

# Feature Views
driver_hourly_stats_view = FeatureView(
    name="driver_hourly_stats",
    entities=[driver],
    ttl=timedelta(days=7),
    schema=[
        Field(name="conv_rate", dtype=Float32),
        Field(name="acc_rate", dtype=Float32),
        Field(name="avg_daily_trips", dtype=Int64),
    ],
    online=True,
    source=driver_hourly_stats,
    tags={},
    version="latest",
)
customer_daily_profile_view = FeatureView(
    name="customer_daily_profile",
    entities=[customer],
    ttl=timedelta(days=7),
    schema=[
        Field(name="current_balance", dtype=Float32),
        Field(name="avg_passenger_count", dtype=Float32),
        Field(name="lifetime_trip_count", dtype=Int64),
    ],
    online=True,
    source=customer_daily_profile,
    tags={},
    version="latest",
)

driver_stats_fs = FeatureService(
    name="driver_activity",
    features=[driver_hourly_stats_view, customer_daily_profile_view],
)