feat: Added RemoteDatasetProxy that executes Ray Data operations remotely

Signed-off-by: ntkathole <[email protected]>

Author: ntkathole

sdk/python/feast/infra/offline_stores/contrib/ray_offline_store/ray.py

@@ -43,6 +43,7 @@
     _build_required_columns,
     apply_field_mapping,
     ensure_timestamp_compatibility,
+    is_ray_data,
     normalize_timestamp_columns,
 )
 from feast.infra.registry.base_registry import BaseRegistry
@@ -61,7 +62,7 @@
 
 
 def _get_data_schema_info(
-    data: Union[pd.DataFrame, Dataset],
+    data: Union[pd.DataFrame, Dataset, Any],
 ) -> Tuple[Dict[str, Any], List[str]]:
     """
     Extract schema information from DataFrame or Dataset.
@@ -70,7 +71,7 @@ def _get_data_schema_info(
     Returns:
         Tuple of (dtypes_dict, column_names)
     """
-    if isinstance(data, Dataset):
+    if is_ray_data(data):
         schema = data.schema()
         dtypes = {}
         for i, col in enumerate(schema.names):
@@ -84,16 +85,17 @@ def _get_data_schema_info(
                 dtypes[col] = pd.api.types.pandas_dtype("object")
         columns = schema.names
     else:
+        assert isinstance(data, pd.DataFrame)
         dtypes = data.dtypes.to_dict()
         columns = list(data.columns)
     return dtypes, columns
 
 
 def _apply_to_data(
-    data: Union[pd.DataFrame, Dataset],
+    data: Union[pd.DataFrame, Dataset, Any],
     process_func: Callable[[pd.DataFrame], pd.DataFrame],
     inplace: bool = False,
-) -> Union[pd.DataFrame, Dataset]:
+) -> Union[pd.DataFrame, Dataset, Any]:
     """
     Apply a processing function to DataFrame or Dataset.
     Args:
@@ -103,9 +105,10 @@ def _apply_to_data(
     Returns:
         Processed DataFrame or Dataset
     """
-    if isinstance(data, Dataset):
+    if is_ray_data(data):
         return data.map_batches(process_func, batch_format="pandas")
     else:
+        assert isinstance(data, pd.DataFrame)
         if not inplace:
             data = data.copy()
         return process_func(data)
@@ -158,7 +161,7 @@ def _safe_infer_event_timestamp_column(
 
 
 def _safe_get_entity_timestamp_bounds(
-    data: Union[pd.DataFrame, Dataset], timestamp_column: str
+    data: Union[pd.DataFrame, Dataset, Any], timestamp_column: str
 ) -> Tuple[Optional[datetime], Optional[datetime]]:
     """
     Safely get entity timestamp bounds.
@@ -170,7 +173,7 @@ def _safe_get_entity_timestamp_bounds(
         Tuple of (min_timestamp, max_timestamp) or (None, None) if failed
     """
     try:
-        if isinstance(data, Dataset):
+        if is_ray_data(data):
             min_ts = data.min(timestamp_column)
             max_ts = data.max(timestamp_column)
         else:
@@ -192,7 +195,7 @@ def _safe_get_entity_timestamp_bounds(
             f"Timestamp bounds extraction failed: {e}, falling back to manual calculation"
         )
         try:
-            if isinstance(data, Dataset):
+            if is_ray_data(data):
 
                 def extract_bounds(batch: pd.DataFrame) -> pd.DataFrame:
                     if timestamp_column in batch.columns and not batch.empty:
@@ -212,6 +215,7 @@ def extract_bounds(batch: pd.DataFrame) -> pd.DataFrame:
                     if pd.notna(min_ts) and pd.notna(max_ts):
                         return min_ts.to_pydatetime(), max_ts.to_pydatetime()
             else:
+                assert isinstance(data, pd.DataFrame)
                 if timestamp_column in data.columns:
                     timestamps = pd.to_datetime(data[timestamp_column], utc=True)
                     return (

sdk/python/feast/infra/ray_shared_utils.py

@@ -1,16 +1,145 @@
-from typing import Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 import numpy as np
 import pandas as pd
+import pyarrow as pa
+import ray
 from ray.data import Dataset
 
 
+class RemoteDatasetProxy:
+    """Proxy class that executes Ray Data operations remotely on cluster workers."""
+
+    def __init__(self, dataset_ref: Any):
+        """Initialize with a reference to the remote dataset."""
+        self._dataset_ref = dataset_ref
+
+    def map_batches(self, func, **kwargs) -> "RemoteDatasetProxy":
+        """Execute map_batches remotely on cluster workers."""
+
+        @ray.remote
+        def _remote_map_batches(dataset, function, batch_kwargs):
+            return dataset.map_batches(function, **batch_kwargs)
+
+        new_ref = _remote_map_batches.remote(self._dataset_ref, func, kwargs)
+        return RemoteDatasetProxy(new_ref)
+
+    def filter(self, fn) -> "RemoteDatasetProxy":
+        """Execute filter remotely on cluster workers."""
+
+        @ray.remote
+        def _remote_filter(dataset, filter_fn):
+            return dataset.filter(filter_fn)
+
+        new_ref = _remote_filter.remote(self._dataset_ref, fn)
+        return RemoteDatasetProxy(new_ref)
+
+    def to_pandas(self) -> pd.DataFrame:
+        """Execute to_pandas remotely and transfer result to client."""
+
+        @ray.remote
+        def _remote_to_pandas(dataset):
+            return dataset.to_pandas()
+
+        result_ref = _remote_to_pandas.remote(self._dataset_ref)
+        return ray.get(result_ref)
+
+    def to_arrow(self) -> pa.Table:
+        """Execute to_arrow remotely and transfer result to client."""
+
+        @ray.remote
+        def _remote_to_arrow(dataset):
+            return dataset.to_arrow()
+
+        result_ref = _remote_to_arrow.remote(self._dataset_ref)
+        return ray.get(result_ref)
+
+    def schema(self) -> Any:
+        """Get dataset schema."""
+
+        @ray.remote
+        def _remote_schema(dataset):
+            return dataset.schema()
+
+        schema_ref = _remote_schema.remote(self._dataset_ref)
+        return ray.get(schema_ref)
+
+    def sort(self, key, descending=False) -> "RemoteDatasetProxy":
+        """Execute sort remotely on cluster workers."""
+
+        @ray.remote
+        def _remote_sort(dataset, sort_key, desc):
+            return dataset.sort(sort_key, descending=desc)
+
+        new_ref = _remote_sort.remote(self._dataset_ref, key, descending)
+        return RemoteDatasetProxy(new_ref)
+
+    def limit(self, count) -> "RemoteDatasetProxy":
+        """Execute limit remotely on cluster workers."""
+
+        @ray.remote
+        def _remote_limit(dataset, limit_count):
+            return dataset.limit(limit_count)
+
+        new_ref = _remote_limit.remote(self._dataset_ref, count)
+        return RemoteDatasetProxy(new_ref)
+
+    def union(self, other) -> "RemoteDatasetProxy":
+        """Execute union remotely on cluster workers."""
+
+        @ray.remote
+        def _remote_union(dataset1, dataset2):
+            return dataset1.union(dataset2)
+
+        new_ref = _remote_union.remote(self._dataset_ref, other._dataset_ref)
+        return RemoteDatasetProxy(new_ref)
+
+    def materialize(self) -> "RemoteDatasetProxy":
+        """Execute materialize remotely on cluster workers."""
+
+        @ray.remote
+        def _remote_materialize(dataset):
+            return dataset.materialize()
+
+        new_ref = _remote_materialize.remote(self._dataset_ref)
+        return RemoteDatasetProxy(new_ref)
+
+    def count(self) -> int:
+        """Execute count remotely and return result."""
+
+        @ray.remote
+        def _remote_count(dataset):
+            return dataset.count()
+
+        result_ref = _remote_count.remote(self._dataset_ref)
+        return ray.get(result_ref)
+
+    def take(self, n=20) -> list:
+        """Execute take remotely and return result."""
+
+        @ray.remote
+        def _remote_take(dataset, num):
+            return dataset.take(num)
+
+        result_ref = _remote_take.remote(self._dataset_ref, n)
+        return ray.get(result_ref)
+
+    def __getattr__(self, name):
+        """Catch any method calls that we haven't explicitly implemented."""
+        raise AttributeError(f"RemoteDatasetProxy has no attribute '{name}'")
+
+
+def is_ray_data(data: Any) -> bool:
+    """Check if data is a Ray Dataset or RemoteDatasetProxy."""
+    return isinstance(data, (Dataset, RemoteDatasetProxy))
+
+
 def normalize_timestamp_columns(
-    data: Union[pd.DataFrame, Dataset],
+    data: Union[pd.DataFrame, Dataset, Any],
     columns: Union[str, List[str]],
     inplace: bool = False,
     exclude_columns: Optional[List[str]] = None,
-) -> Union[pd.DataFrame, Dataset]:
+) -> Union[pd.DataFrame, Dataset, Any]:
     column_list = [columns] if isinstance(columns, str) else columns
     exclude_columns = exclude_columns or []
 
@@ -21,7 +150,7 @@ def apply_normalization(series: pd.Series) -> pd.Series:
             .astype("datetime64[ns, UTC]")
         )
 
-    if isinstance(data, Dataset):
+    if is_ray_data(data):
 
         def normalize_batch(batch: pd.DataFrame) -> pd.DataFrame:
             for column in column_list:
@@ -35,6 +164,7 @@ def normalize_batch(batch: pd.DataFrame) -> pd.DataFrame:
 
         return data.map_batches(normalize_batch, batch_format="pandas")
     else:
+        assert isinstance(data, pd.DataFrame)
         if not inplace:
             data = data.copy()
         for column in column_list:
@@ -44,13 +174,13 @@ def normalize_batch(batch: pd.DataFrame) -> pd.DataFrame:
 
 
 def ensure_timestamp_compatibility(
-    data: Union[pd.DataFrame, Dataset],
+    data: Union[pd.DataFrame, Dataset, Any],
     timestamp_fields: List[str],
     inplace: bool = False,
-) -> Union[pd.DataFrame, Dataset]:
+) -> Union[pd.DataFrame, Dataset, Any]:
     from feast.utils import make_df_tzaware
 
-    if isinstance(data, Dataset):
+    if is_ray_data(data):
 
         def ensure_compatibility(batch: pd.DataFrame) -> pd.DataFrame:
             batch = make_df_tzaware(batch)
@@ -65,6 +195,7 @@ def ensure_compatibility(batch: pd.DataFrame) -> pd.DataFrame:
 
         return data.map_batches(ensure_compatibility, batch_format="pandas")
     else:
+        assert isinstance(data, pd.DataFrame)
         if not inplace:
             data = data.copy()
         from feast.utils import make_df_tzaware
@@ -77,22 +208,24 @@ def ensure_compatibility(batch: pd.DataFrame) -> pd.DataFrame:
 
 
 def apply_field_mapping(
-    data: Union[pd.DataFrame, Dataset], field_mapping: Dict[str, str]
-) -> Union[pd.DataFrame, Dataset]:
+    data: Union[pd.DataFrame, Dataset, Any],
+    field_mapping: Dict[str, str],
+) -> Union[pd.DataFrame, Dataset, Any]:
     def rename_columns(df: pd.DataFrame) -> pd.DataFrame:
         return df.rename(columns=field_mapping)
 
-    if isinstance(data, Dataset):
+    if is_ray_data(data):
         return data.map_batches(rename_columns, batch_format="pandas")
     else:
+        assert isinstance(data, pd.DataFrame)
         return data.rename(columns=field_mapping)
 
 
 def deduplicate_by_keys_and_timestamp(
-    data: Union[pd.DataFrame, Dataset],
+    data: Union[pd.DataFrame, Dataset, Any],
     join_keys: List[str],
     timestamp_columns: List[str],
-) -> Union[pd.DataFrame, Dataset]:
+) -> Union[pd.DataFrame, Dataset, Any]:
     def deduplicate_batch(batch: pd.DataFrame) -> pd.DataFrame:
         if batch.empty:
             return batch
@@ -110,9 +243,10 @@ def deduplicate_batch(batch: pd.DataFrame) -> pd.DataFrame:
             return deduped_batch
         return batch
 
-    if isinstance(data, Dataset):
+    if is_ray_data(data):
         return data.map_batches(deduplicate_batch, batch_format="pandas")
     else:
+        assert isinstance(data, pd.DataFrame)
         return deduplicate_batch(data)