feat: Compute Engine Initial Implementation

[HaoXuAI]

What this PR does / why we need it:

Introduce the ComputeEngine component that runs current materialization and get_historical_features on a compute engine such as Spark, Flink, etc.
This will hand over heavy computing work from online/offline stores to the engines. And let the online/offline deal with IO only.
The compute engine builds the execution plan in a DAG format named FeatureBuilder. It derives feature generation from Feature View definitiions including:

1. Transformation (via Transformation API)
2. Aggregation (via Aggregation API)
3. Join (join with an entity, adding Customized JOIN and join with another Feature View to do in the future)
4. Filter (Point in time filter, ttl filter, filter by custom expression)
5. ...

User facing API, e.g:

def transform_feature(df: DataFrame) -> DataFrame:
    df = df.withColumn("sum_conv_rate", df["sum_conv_rate"] * 2)
    df = df.withColumn("avg_acc_rate", df["avg_acc_rate"] * 2)
    return df

driver_stats_fv = BatchFeatureView(
    name="driver_hourly_stats",
    entities=[driver],
    mode="python",
    batch_engine=Field(...), # Not supported yet.
    aggregations=[
        Aggregation(column="conv_rate", function="sum"),
        Aggregation(column="acc_rate", function="avg"),
    ],
    join=[...], # Not supported yet
    udf=transform_feature,
    udf_string="transform_feature",
    ttl=timedelta(days=3), # Not supported for materialization yet
    filter="filter_expr", # Not fully support
    schema=[
        Field(name="conv_rate", dtype=Float32),
        Field(name="acc_rate", dtype=Float32),
        Field(name="avg_daily_trips", dtype=Int64),
        Field(name="driver_id", dtype=Int32),
    ],
    online=True,
    offline=False, # not supported yet
    source=data_source,
)

Feature build flow:

1. join with entity first if necessary
2. filter data based on entity event timestamp, ttl, and filter
3. feature aggregation based on provided aggregations
4. run transformation based on provided udf or feature_transformations
5. write data to online or offline

Current use case:

# Create retrieval task
task = MaterializationTask(
    project=spark_environment.project,
    feature_view=driver_stats_fv,
    start_time=now - timedelta(days=1),
    end_time=now,
    tqdm_builder=tqdm_builder,
)

# Create SparkComputeEngine
engine = SparkComputeEngine(
    repo_config=spark_environment.config,
    offline_store=SparkOfflineStore(),
    online_store=MagicMock(),
    registry=registry,
)

spark_materialize_job = engine.materialize(task)

TODOs:

• Update Batch/Stream feature view to support batch_engine and stream_engine configurations.
• Support offline write
• Change materialization read node to the engine, and merge the flow with get_historical_features
• Support filter by filter_expr in feature view
• Use engine APIs in feature_store API.
• Support compute engine in repo yaml config:
• Join with Feature View or Data Source.
• Update get_online_feature to be able to handle aggregated and transformed features
• Research the Transformation, Aggregation, and Filter to merge into single API.

Which issue(s) this PR fixes:

Misc

[franciscojavierarceo]

Nice! Left some small nits on naming.

[franciscojavierarceo]

I think we should remove the 'node' suffix from all of there, as I think it's implicit.

[HaoXuAI]

yeah sounds good

[HaoXuAI]

On second thought, I think it's better to keep the 'node' suffix, as it makes the output more descriptive and clearly indicates that it's a node, which can then be chained in the build method

[franciscojavierarceo]

not: build_dag(self)

[HaoXuAI]

I would suggest to keep build(), more consistent with the builder.build() pattern.

[franciscojavierarceo]

Native python?

[HaoXuAI]

Right, it's only a placeholder for now. The native python will be interesting as it works with on demand feature view. Will do it in the next PR

[franciscojavierarceo]

Serializing as pandas here could be expensive no?

[HaoXuAI]

Right, only to keep it consistent with the current get_historical_features API. But definitely needs some rework

[franciscojavierarceo]

Yeah just like above, don't think we need node suffix for these methods

[franciscojavierarceo]

@HaoXuAI is this ready to review? Or still draft?

[HaoXuAI]

Not yet. Still needs some work but should be ready soon

[franciscojavierarceo]

If I'm to be honest, I don't love putting the transformation in the FeatureView.

I think it'd be more intuitive to put Aggregation under Transformation and make the FeatureViews purely represent schemas available for online or offline.

[HaoXuAI]

I agree. This is only to be consistent with the current StreamFeatureView and Aggregation class. It makes sense to have Aggregation one type of Transformation (as well as for Filter), but that will need some work to refactor.
For Tecton, they put Aggregate in feature (what Feast called schema) as

features=[
        Aggregate(
            input_column=Field("transaction", Int64),
            function="count",
            time_window=TimeWindow(window_size=timedelta(days=1)),
        )
    ],

In some way that also make sense but for Feast it will needs much more work to refactor the schema API.

For Chrono they put it into a GroupBy API, which is similar to our FeatureView:

v1 = GroupBy(
    sources=[source],
    keys=["user_id"], # We are aggregating by user
    online=True,
    aggregations=[Aggregation(
            input_column="refund_amt",
            operation=Operation.SUM,
            windows=window_sizes
        ), # The sum of purchases prices in various windows
    ],
)

I can look into how to merge Aggregation with Transformation together in the next PR, and added it to a TODO for now.

[franciscojavierarceo]

I suppose my concern with calling it FeatureBuilder rather than DAGBuilder is because users will say that this isn't a fit for RAG.

[HaoXuAI]

I think RAG or embedding is still a type of feature, so that should still work?
Tecton defines Embedding into feature directly: https://docs.tecton.ai/docs/beta/defining-features/feature-views/embeddings which is also interesting.
For feast, RAG or embedding can be either custom udf in Transformation, and maybe like Aggregation a separate API (Or merge into a special type of Transformation).

[franciscojavierarceo]

Yeah, I don't think Tecton has captured the audience there candidly. But maybe that's fine for now and we can revisit later so no worries.

[franciscojavierarceo]

in this:

def transform_feature(df: DataFrame) -> DataFrame:
    df = df.withColumn("sum_conv_rate", df["sum_conv_rate"] * 2)
    df = df.withColumn("avg_acc_rate", df["avg_acc_rate"] * 2)
    return df

driver_stats_fv = BatchFeatureView(
    name="driver_hourly_stats",
    entities=[driver],
    mode="python",
    batch_engine=Field(...), # Not supported yet.
    aggregations=[
        Aggregation(column="conv_rate", function="sum"),
        Aggregation(column="acc_rate", function="avg"),
    ],
    join=[...], # Not supported yet
    udf=transform_feature,
    udf_string="transform_feature",
    ttl=timedelta(days=3), # Not supported for materialization yet
    filter="filter_expr", # Not fully support
    schema=[
        Field(name="conv_rate", dtype=Float32),
        Field(name="acc_rate", dtype=Float32),
        Field(name="avg_daily_trips", dtype=Int64),
        Field(name="driver_id", dtype=Int32),
    ],
    online=True,
    offline=False, # not supported yet
    source=data_source,
)

What is the point of the Aggregations in the Feature View?

    aggregations=[
        Aggregation(column="conv_rate", function="sum"),
        Aggregation(column="acc_rate", function="avg"),
    ],

[HaoXuAI]

in this:

def transform_feature(df: DataFrame) -> DataFrame:
    df = df.withColumn("sum_conv_rate", df["sum_conv_rate"] * 2)
    df = df.withColumn("avg_acc_rate", df["avg_acc_rate"] * 2)
    return df

driver_stats_fv = BatchFeatureView(
    name="driver_hourly_stats",
    entities=[driver],
    mode="python",
    batch_engine=Field(...), # Not supported yet.
    aggregations=[
        Aggregation(column="conv_rate", function="sum"),
        Aggregation(column="acc_rate", function="avg"),
    ],
    join=[...], # Not supported yet
    udf=transform_feature,
    udf_string="transform_feature",
    ttl=timedelta(days=3), # Not supported for materialization yet
    filter="filter_expr", # Not fully support
    schema=[
        Field(name="conv_rate", dtype=Float32),
        Field(name="acc_rate", dtype=Float32),
        Field(name="avg_daily_trips", dtype=Int64),
        Field(name="driver_id", dtype=Int32),
    ],
    online=True,
    offline=False, # not supported yet
    source=data_source,
)

What is the point of the Aggregations in the Feature View?

    aggregations=[
        Aggregation(column="conv_rate", function="sum"),
        Aggregation(column="acc_rate", function="avg"),
    ],

aggregation from the Design perspective is to simplify defining aggregation features with OOTB API. There are some benefits listed in Tecton: https://docs.tecton.ai/docs/beta/defining-features/feature-views/aggregation-engine. What I think important is the lifetime window support. For the Aggregation you can define
Aggregation(column="conv_rate", function="sum", time_window=timedelta(days=1)) where the compute will handle the window logic for you. For stream compute it is pretty helpful since the time window concept is bit complex for end users to understand such as the Flink window (https://nightlies.apache.org/flink/flink-docs-master/docs/dev/datastream/operators/windows/) we used before

[franciscojavierarceo]

For the Aggregation you can define
Aggregation(column="conv_rate", function="sum", time_window=timedelta(days=1)) where the compute will handle the window logic for you. For stream compute it is pretty helpful since the time window concept is bit complex for end users to understand such as the Flink window (https://nightlies.apache.org/flink/flink-docs-master/docs/dev/datastream/operators/windows/) we used before

Yeah that makes a lot of sense, I'm good with this.

[HaoXuAI]

For the Aggregation you can define
Aggregation(column="conv_rate", function="sum", time_window=timedelta(days=1)) where the compute will handle the window logic for you. For stream compute it is pretty helpful since the time window concept is bit complex for end users to understand such as the Flink window (https://nightlies.apache.org/flink/flink-docs-master/docs/dev/datastream/operators/windows/) we used before

Yeah that makes a lot of sense, I'm good with this.

This is helpful and a really good discussion. And I agree the interface is still immature and definitely subject to change.

[franciscojavierarceo]

Let's go! 🚀🚀🚀