Repository: feast-dev/feast

Issue:
Substrait-based on demand feature views
**Is your feature request related to a problem? Please describe.**
On demand feature views as implemented right now are very limited. The only way to specify odfvs is through a python function that takes in pandas Dataframe as input and outputs another pandas Dataframe. This leads to problems for both offline and online interfaces:

- Even the most scalable offline stores are forced to collect the whole dataset as a single pandas Dataframe to apply odfv function. There's no way for offline stores to incorporate computation in their engines.
- udfs in odfvs are inherently bound to pandas and python runtime. Non-python feature servers are stuck with the problem of figuring out how to run this functions if necessary. Java feature server for example has a separate python transformation service only for this reason, but that's obviously a subpar solution as the whole point of a java feature server was to avoid python runtime in feature serving in the first place.

**Describe the solution you'd like**
Allow constructing odfvs as substrait plans. [Substrait](https://substrait.io/) is a protobuf-based serialization format for relational algebra operations. It is meant to be used as a cross-language and cross-engine format for sharing logical or physical execution plans. It has a number of producers (tools that can generate substrait) and consumers (engines that can run substrait) in different languages.

- Different offline stores will be able to inspect and incorporate substrait plans in their transformations. Even if that's impossible the default implementation inside feast to apply these functions will avoid pandas.
- Most importantly, non-python feature servers like a java feature server will be able to apply the functions without a separate python component. Apache Arrow java implementation comes with java bindings to Acero query engine that can consume substrait plans. (https://arrow.apache.org/docs/java/substrait.html#executing-queries-using-substrait-plans)

The example code in my PoC implementation looks something like this:

```
def generate_substrait():
    import ibis
    from ibis_substrait.compiler.core import SubstraitCompiler

    compiler = SubstraitCompiler()

    t = ibis.table([("conv_rate", "float"), ("acc_rate", "float")], "t")

    expr = t.select((t['conv_rate'] + t['acc_rate']).name('conv_rate_plus_acc_substrait'))

    return compiler.compile(expr).SerializeToString()

substrait_odfv = OnDemandFeatureView(
    name='substrait_view',
    sources=[driver_stats_fv],
    schema=[
        Field(name="conv_rate_plus_acc_substrait", dtype=Float64)
    ],
    substrait_plan=generate_substrait()
)
```
Substait plan object that feast accepts is `bytes` and introduces no external dependency. I'm using `ibis` and `ibis-substrait` to generate the plan. Right now that's the most practical way to generate substrait plan in python with DataFrame-like API, but this could have been any other substrait producer.

**Describe alternatives you've considered**
An obvious alternative to substrait is sql-based odfvs, but using SQL has a number of important downsides:

1. The presence of different sql dialects means that, it will be especially hard to ensure that sql-based feature functions will behave the same way across different offline store and online store implementations.
2. The user is implicitly bound to their offline store and online store of choice, because the dialect used in sql strings has to match offline store engine.

Having said that, it probably makes sense to support both substrait-based and sql-based odfvs, because at the moment it might be easier for sql-based logic to be incorporated inside offline store engines.


Discussion:

Resolving patch:
diff --git a/sdk/python/feast/on_demand_feature_view.py b/sdk/python/feast/on_demand_feature_view.py
--- a/sdk/python/feast/on_demand_feature_view.py
+++ b/sdk/python/feast/on_demand_feature_view.py
@@ -1,5 +1,6 @@
 import copy
 import functools
+import inspect
 import warnings
 from datetime import datetime
 from types import FunctionType
@@ -17,6 +18,7 @@
 from feast.feature_view_projection import FeatureViewProjection
 from feast.field import Field, from_value_type
 from feast.on_demand_pandas_transformation import OnDemandPandasTransformation
+from feast.on_demand_substrait_transformation import OnDemandSubstraitTransformation
 from feast.protos.feast.core.OnDemandFeatureView_pb2 import (
     OnDemandFeatureView as OnDemandFeatureViewProto,
 )
@@ -210,6 +212,9 @@ def to_proto(self) -> OnDemandFeatureViewProto:
             user_defined_function=self.transformation.to_proto()
             if type(self.transformation) == OnDemandPandasTransformation
             else None,
+            on_demand_substrait_transformation=self.transformation.to_proto()  # type: ignore
+            if type(self.transformation) == OnDemandSubstraitTransformation
+            else None,
             description=self.description,
             tags=self.tags,
             owner=self.owner,
@@ -255,6 +260,13 @@ def from_proto(cls, on_demand_feature_view_proto: OnDemandFeatureViewProto):
             transformation = OnDemandPandasTransformation.from_proto(
                 on_demand_feature_view_proto.spec.user_defined_function
             )
+        elif (
+            on_demand_feature_view_proto.spec.WhichOneof("transformation")
+            == "on_demand_substrait_transformation"
+        ):
+            transformation = OnDemandSubstraitTransformation.from_proto(
+                on_demand_feature_view_proto.spec.on_demand_substrait_transformation
+            )
         else:
             raise Exception("At least one transformation type needs to be provided")
 
@@ -460,10 +472,47 @@ def mainify(obj) -> None:
             obj.__module__ = "__main__"
 
     def decorator(user_function):
-        udf_string = dill.source.getsource(user_function)
-        mainify(user_function)
+        return_annotation = inspect.signature(user_function).return_annotation
+        if (
+            return_annotation
+            and return_annotation.__module__ == "ibis.expr.types.relations"
+            and return_annotation.__name__ == "Table"
+        ):
+            import ibis
+            import ibis.expr.datatypes as dt
+            from ibis_substrait.compiler.core import SubstraitCompiler
+
+            compiler = SubstraitCompiler()
+
+            input_fields: Field = []
+
+            for s in sources:
+                if type(s) == FeatureView:
+                    fields = s.projection.features
+                else:
+                    fields = s.features
+
+                input_fields.extend(
+                    [
+                        (
+                            f.name,
+                            dt.dtype(
+                                feast_value_type_to_pandas_type(f.dtype.to_value_type())
+                            ),
+                        )
+                        for f in fields
+                    ]
+                )
+
+            expr = user_function(ibis.table(input_fields, "t"))
 
-        transformation = OnDemandPandasTransformation(user_function, udf_string)
+            transformation = OnDemandSubstraitTransformation(
+                substrait_plan=compiler.compile(expr).SerializeToString()
+            )
+        else:
+            udf_string = dill.source.getsource(user_function)
+            mainify(user_function)
+            transformation = OnDemandPandasTransformation(user_function, udf_string)
 
         on_demand_feature_view_obj = OnDemandFeatureView(
             name=user_function.__name__,
diff --git a/sdk/python/feast/on_demand_substrait_transformation.py b/sdk/python/feast/on_demand_substrait_transformation.py
new file mode 100644
--- /dev/null
+++ b/sdk/python/feast/on_demand_substrait_transformation.py
@@ -0,0 +1,50 @@
+import pandas as pd
+import pyarrow
+import pyarrow.substrait as substrait  # type: ignore # noqa
+
+from feast.protos.feast.core.OnDemandFeatureView_pb2 import (
+    OnDemandSubstraitTransformation as OnDemandSubstraitTransformationProto,
+)
+
+
+class OnDemandSubstraitTransformation:
+    def __init__(self, substrait_plan: bytes):
+        """
+        Creates an OnDemandSubstraitTransformation object.
+
+        Args:
+            substrait_plan: The user-provided substrait plan.
+        """
+        self.substrait_plan = substrait_plan
+
+    def transform(self, df: pd.DataFrame) -> pd.DataFrame:
+        def table_provider(names, schema: pyarrow.Schema):
+            return pyarrow.Table.from_pandas(df[schema.names])
+
+        table: pyarrow.Table = pyarrow.substrait.run_query(
+            self.substrait_plan, table_provider=table_provider
+        ).read_all()
+        return table.to_pandas()
+
+    def __eq__(self, other):
+        if not isinstance(other, OnDemandSubstraitTransformation):
+            raise TypeError(
+                "Comparisons should only involve OnDemandSubstraitTransformation class objects."
+            )
+
+        if not super().__eq__(other):
+            return False
+
+        return self.substrait_plan == other.substrait_plan
+
+    def to_proto(self) -> OnDemandSubstraitTransformationProto:
+        return OnDemandSubstraitTransformationProto(substrait_plan=self.substrait_plan)
+
+    @classmethod
+    def from_proto(
+        cls,
+        on_demand_substrait_transformation_proto: OnDemandSubstraitTransformationProto,
+    ):
+        return OnDemandSubstraitTransformation(
+            substrait_plan=on_demand_substrait_transformation_proto.substrait_plan
+        )
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -100,7 +100,7 @@
 
 AWS_REQUIRED = ["boto3>=1.17.0,<2", "docker>=5.0.2", "fsspec<=2024.1.0"]
 
-BYTEWAX_REQUIRED = ["bytewax==0.18.2", "docker>=5.0.2", "kubernetes<=20.13.0"]
+BYTEWAX_REQUIRED = ["bytewax==0.15.1", "docker>=5.0.2", "kubernetes<=20.13.0"]
 
 SNOWFLAKE_REQUIRED = [
     "snowflake-connector-python[pandas]>=3,<4",
@@ -144,6 +144,11 @@
     "hazelcast-python-client>=5.1",
 ]
 
+IBIS_REQUIRED = [
+    "ibis-framework",
+    "ibis-substrait"
+]
+
 CI_REQUIRED = (
     [
         "build",
@@ -201,6 +206,7 @@
     + AZURE_REQUIRED
     + ROCKSET_REQUIRED
     + HAZELCAST_REQUIRED
+    + IBIS_REQUIRED
 )
 
 
@@ -368,6 +374,7 @@ def run(self):
         "cassandra": CASSANDRA_REQUIRED,
         "hazelcast": HAZELCAST_REQUIRED,
         "rockset": ROCKSET_REQUIRED,
+        "ibis": IBIS_REQUIRED
     },
     include_package_data=True,
     license="Apache",


Test patch:
diff --git a/sdk/python/tests/unit/test_on_demand_substrait_transformation.py b/sdk/python/tests/unit/test_on_demand_substrait_transformation.py
new file mode 100644
--- /dev/null
+++ b/sdk/python/tests/unit/test_on_demand_substrait_transformation.py
@@ -0,0 +1,112 @@
+import os
+import tempfile
+from datetime import datetime, timedelta
+
+import pandas as pd
+
+from feast import Entity, FeatureStore, FeatureView, FileSource, RepoConfig
+from feast.driver_test_data import create_driver_hourly_stats_df
+from feast.field import Field
+from feast.infra.online_stores.sqlite import SqliteOnlineStoreConfig
+from feast.on_demand_feature_view import on_demand_feature_view
+from feast.types import Float32, Float64, Int64
+
+
+def test_ibis_pandas_parity():
+    with tempfile.TemporaryDirectory() as data_dir:
+        store = FeatureStore(
+            config=RepoConfig(
+                project="test_on_demand_substrait_transformation",
+                registry=os.path.join(data_dir, "registry.db"),
+                provider="local",
+                entity_key_serialization_version=2,
+                online_store=SqliteOnlineStoreConfig(
+                    path=os.path.join(data_dir, "online.db")
+                ),
+            )
+        )
+
+        # Generate test data.
+        end_date = datetime.now().replace(microsecond=0, second=0, minute=0)
+        start_date = end_date - timedelta(days=15)
+
+        driver_entities = [1001, 1002, 1003, 1004, 1005]
+        driver_df = create_driver_hourly_stats_df(driver_entities, start_date, end_date)
+        driver_stats_path = os.path.join(data_dir, "driver_stats.parquet")
+        driver_df.to_parquet(path=driver_stats_path, allow_truncated_timestamps=True)
+
+        driver = Entity(name="driver", join_keys=["driver_id"])
+
+        driver_stats_source = FileSource(
+            name="driver_hourly_stats_source",
+            path=driver_stats_path,
+            timestamp_field="event_timestamp",
+            created_timestamp_column="created",
+        )
+
+        driver_stats_fv = FeatureView(
+            name="driver_hourly_stats",
+            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_source,
+        )
+
+        @on_demand_feature_view(
+            sources=[driver_stats_fv],
+            schema=[Field(name="conv_rate_plus_acc", dtype=Float64)],
+        )
+        def pandas_view(inputs: pd.DataFrame) -> pd.DataFrame:
+            df = pd.DataFrame()
+            df["conv_rate_plus_acc"] = inputs["conv_rate"] + inputs["acc_rate"]
+            return df
+
+        from ibis.expr.types import Table
+
+        @on_demand_feature_view(
+            sources=[driver_stats_fv[["conv_rate", "acc_rate"]]],
+            schema=[Field(name="conv_rate_plus_acc_substrait", dtype=Float64)],
+        )
+        def substrait_view(inputs: Table) -> Table:
+            return inputs.select(
+                (inputs["conv_rate"] + inputs["acc_rate"]).name(
+                    "conv_rate_plus_acc_substrait"
+                )
+            )
+
+        store.apply(
+            [driver, driver_stats_source, driver_stats_fv, substrait_view, pandas_view]
+        )
+
+        entity_df = pd.DataFrame.from_dict(
+            {
+                # entity's join key -> entity values
+                "driver_id": [1001, 1002, 1003],
+                # "event_timestamp" (reserved key) -> timestamps
+                "event_timestamp": [
+                    datetime(2021, 4, 12, 10, 59, 42),
+                    datetime(2021, 4, 12, 8, 12, 10),
+                    datetime(2021, 4, 12, 16, 40, 26),
+                ],
+            }
+        )
+
+        training_df = store.get_historical_features(
+            entity_df=entity_df,
+            features=[
+                "driver_hourly_stats:conv_rate",
+                "driver_hourly_stats:acc_rate",
+                "driver_hourly_stats:avg_daily_trips",
+                "substrait_view:conv_rate_plus_acc_substrait",
+                "pandas_view:conv_rate_plus_acc",
+            ],
+        ).to_df()
+
+        assert training_df["conv_rate_plus_acc"].equals(
+            training_df["conv_rate_plus_acc_substrait"]
+        )