Add FeastDataFrame

HaoXuAI · HaoXuAI · commit 37923de7cf0d · 2025-09-12T21:46:39.000-07:00
Signed-off-by: HaoXuAI &lt;sduxuhao@gmail.com&gt;
diff --git a/sdk/python/feast/__init__.py b/sdk/python/feast/__init__.py
@@ -11,6 +11,7 @@
 
 from .batch_feature_view import BatchFeatureView
 from .data_source import KafkaSource, KinesisSource, PushSource, RequestSource
+from .dataframe import DataFrameEngine, FeastDataFrame
 from .entity import Entity
 from .feature import Feature
 from .feature_service import FeatureService
@@ -32,9 +33,11 @@
 
 __all__ = [
     "BatchFeatureView",
+    "DataFrameEngine",
     "Entity",
     "KafkaSource",
     "KinesisSource",
+    "FeastDataFrame",
     "Feature",
     "Field",
     "FeatureService",
diff --git a/sdk/python/feast/dataframe.py b/sdk/python/feast/dataframe.py
@@ -0,0 +1,83 @@
+"""FeastDataFrame: A lightweight container for DataFrame-like objects in Feast."""
+
+from enum import Enum
+from typing import Any, Dict, Optional, Union
+
+import pandas as pd
+import pyarrow as pa
+
+
+class DataFrameEngine(str, Enum):
+    """Supported DataFrame engines."""
+
+    PANDAS = "pandas"
+    SPARK = "spark"
+    DASK = "dask"
+    RAY = "ray"
+    ARROW = "arrow"
+    POLARS = "polars"
+    UNKNOWN = "unknown"
+
+
+class FeastDataFrame:
+    """
+    A lightweight container for DataFrame-like objects in Feast.
+
+    This class wraps any DataFrame implementation and provides metadata
+    about the engine type for proper routing in Feast's processing pipeline.
+    """
+
+    def __init__(
+        self,
+        data: Any,
+        engine: Optional[DataFrameEngine] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ):
+        """
+        Initialize a FeastDataFrame.
+
+        Args:
+            data: The wrapped DataFrame object (pandas, Spark, Dask, etc.)
+            engine: Explicitly specify the engine type (auto-detected if None)
+            metadata: Additional metadata (schema hints, etc.)
+        """
+        self.data = data
+        self.metadata = metadata or {}
+        self._engine = engine or self._detect_engine()
+
+    def _detect_engine(self) -> DataFrameEngine:
+        """Auto-detect the DataFrame engine based on type."""
+        if isinstance(self.data, pd.DataFrame):
+            return DataFrameEngine.PANDAS
+        elif isinstance(self.data, pa.Table):
+            return DataFrameEngine.ARROW
+        
+        # For optional dependencies, check module name to avoid import errors
+        module = type(self.data).__module__
+        if "pyspark" in module:
+            return DataFrameEngine.SPARK
+        elif "dask" in module:
+            return DataFrameEngine.DASK
+        elif "ray" in module:
+            return DataFrameEngine.RAY
+        elif "polars" in module:
+            return DataFrameEngine.POLARS
+        else:
+            return DataFrameEngine.UNKNOWN
+
+    @property
+    def engine(self) -> DataFrameEngine:
+        """Get the detected or specified engine type."""
+        return self._engine
+
+    def __repr__(self):
+        return f"FeastDataFrame(engine={self.engine}, type={type(self.data).__name__})"
+
+    @property
+    def is_lazy(self) -> bool:
+        """Check if the underlying DataFrame is lazy (Spark, Dask, Ray)."""
+        return self.engine in [
+            DataFrameEngine.SPARK,
+            DataFrameEngine.DASK,
+            DataFrameEngine.RAY,
+        ]
diff --git a/sdk/python/tests/unit/test_dataframe.py b/sdk/python/tests/unit/test_dataframe.py
@@ -0,0 +1,94 @@
+"""Unit tests for FeastDataFrame."""
+
+import pandas as pd
+import pyarrow as pa
+import pytest
+
+from feast.dataframe import DataFrameEngine, FeastDataFrame
+
+
+class TestFeastDataFrame:
+    """Test suite for FeastDataFrame functionality."""
+
+    def test_pandas_detection(self):
+        """Test auto-detection of pandas DataFrame."""
+        df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]})
+        feast_df = FeastDataFrame(df)
+        
+        assert feast_df.engine == DataFrameEngine.PANDAS
+        assert not feast_df.is_lazy
+        assert isinstance(feast_df.data, pd.DataFrame)
+
+    def test_arrow_detection(self):
+        """Test auto-detection of Arrow Table."""
+        table = pa.table({"a": [1, 2, 3], "b": [4, 5, 6]})
+        feast_df = FeastDataFrame(table)
+        
+        assert feast_df.engine == DataFrameEngine.ARROW
+        assert not feast_df.is_lazy
+        assert isinstance(feast_df.data, pa.Table)
+
+    def test_explicit_engine(self):
+        """Test explicit engine specification."""
+        data = {"mock": "data"}
+        feast_df = FeastDataFrame(data, engine=DataFrameEngine.SPARK)
+        
+        assert feast_df.engine == DataFrameEngine.SPARK
+        assert feast_df.is_lazy
+
+    def test_unknown_engine(self):
+        """Test handling of unknown DataFrame types."""
+        data = {"some": "dict"}
+        feast_df = FeastDataFrame(data)
+        
+        assert feast_df.engine == DataFrameEngine.UNKNOWN
+
+    def test_metadata(self):
+        """Test metadata handling."""
+        df = pd.DataFrame({"a": [1, 2, 3]})
+        metadata = {"features": ["a"], "source": "test"}
+        feast_df = FeastDataFrame(df, metadata=metadata)
+        
+        assert feast_df.metadata == metadata
+        assert feast_df.metadata["features"] == ["a"]
+
+    def test_repr(self):
+        """Test string representation."""
+        df = pd.DataFrame({"a": [1, 2, 3]})
+        feast_df = FeastDataFrame(df)
+        
+        repr_str = repr(feast_df)
+        assert "FeastDataFrame" in repr_str
+        assert "engine=pandas" in repr_str
+        assert "DataFrame" in repr_str
+
+    @pytest.mark.parametrize(
+        "engine,expected_lazy",
+        [
+            (DataFrameEngine.PANDAS, False),
+            (DataFrameEngine.ARROW, False),
+            (DataFrameEngine.POLARS, False),
+            (DataFrameEngine.SPARK, True),
+            (DataFrameEngine.DASK, True),
+            (DataFrameEngine.RAY, True),
+            (DataFrameEngine.UNKNOWN, False),
+        ],
+    )
+    def test_is_lazy_property(self, engine, expected_lazy):
+        """Test is_lazy property for different engines."""
+        feast_df = FeastDataFrame({"mock": "data"}, engine=engine)
+        assert feast_df.is_lazy == expected_lazy
+
+    def test_polars_detection(self):
+        """Test detection of polars DataFrame (using mock)."""
+        # Mock polars DataFrame
+        class MockPolarsDF:
+            def __init__(self):
+                self.__module__ = "polars.dataframe.frame"
+                self.__class__.__name__ = "DataFrame"
+        
+        polars_df = MockPolarsDF()
+        feast_df = FeastDataFrame(polars_df)
+        
+        assert feast_df.engine == DataFrameEngine.POLARS
+        assert not feast_df.is_lazy
