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 []

.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:

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:

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)

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

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

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)