|
| 1 | +from abc import ABC, abstractmethod |
| 2 | +from datetime import timedelta |
| 3 | + |
| 4 | + |
| 5 | +class DataFrameBackend(ABC): |
| 6 | + """ |
| 7 | + Abstract interface for DataFrame operations used by the LocalComputeEngine. |
| 8 | +
|
| 9 | + This interface defines the contract for implementing pluggable DataFrame backends |
| 10 | + such as Pandas, Polars, or DuckDB. Each backend must support core table operations |
| 11 | + such as joins, filtering, aggregation, conversion to/from Arrow, and deduplication. |
| 12 | +
|
| 13 | + The purpose of this abstraction is to allow seamless swapping of execution backends |
| 14 | + without changing DAGNode or ComputeEngine logic. All nodes operate on pyarrow.Table |
| 15 | + as the standard input/output format, while the backend defines how the computation |
| 16 | + is actually performed. |
| 17 | +
|
| 18 | + Expected implementations include: |
| 19 | + - PandasBackend |
| 20 | + - PolarsBackend |
| 21 | + - DuckDBBackend (future) |
| 22 | +
|
| 23 | + Methods |
| 24 | + ------- |
| 25 | + from_arrow(table: pa.Table) -> Any |
| 26 | + Convert a pyarrow.Table to the backend-native DataFrame format. |
| 27 | +
|
| 28 | + to_arrow(df: Any) -> pa.Table |
| 29 | + Convert a backend-native DataFrame to pyarrow.Table. |
| 30 | +
|
| 31 | + join(left: Any, right: Any, on: List[str], how: str) -> Any |
| 32 | + Join two dataframes on specified keys with given join type. |
| 33 | +
|
| 34 | + groupby_agg(df: Any, group_keys: List[str], agg_ops: Dict[str, Tuple[str, str]]) -> Any |
| 35 | + Group and aggregate the dataframe. `agg_ops` maps output column names |
| 36 | + to (aggregation function, source column name) pairs. |
| 37 | +
|
| 38 | + filter(df: Any, expr: str) -> Any |
| 39 | + Apply a filter expression (string-based) to the DataFrame. |
| 40 | +
|
| 41 | + to_timedelta_value(delta: timedelta) -> Any |
| 42 | + Convert a Python timedelta object to a backend-compatible value |
| 43 | + that can be subtracted from a timestamp column. |
| 44 | +
|
| 45 | + drop_duplicates(df: Any, keys: List[str], sort_by: List[str], ascending: bool = False) -> Any |
| 46 | + Deduplicate the DataFrame by key columns, keeping the first row |
| 47 | + by descending or ascending sort order. |
| 48 | +
|
| 49 | + rename_columns(df: Any, columns: Dict[str, str]) -> Any |
| 50 | + Rename columns in the DataFrame according to the provided mapping. |
| 51 | + """ |
| 52 | + |
| 53 | + @abstractmethod |
| 54 | + def columns(self, df): ... |
| 55 | + |
| 56 | + @abstractmethod |
| 57 | + def from_arrow(self, table): ... |
| 58 | + |
| 59 | + @abstractmethod |
| 60 | + def join(self, left, right, on, how): ... |
| 61 | + |
| 62 | + @abstractmethod |
| 63 | + def groupby_agg(self, df, group_keys, agg_ops): ... |
| 64 | + |
| 65 | + @abstractmethod |
| 66 | + def filter(self, df, expr): ... |
| 67 | + |
| 68 | + @abstractmethod |
| 69 | + def to_arrow(self, df): ... |
| 70 | + |
| 71 | + @abstractmethod |
| 72 | + def to_timedelta_value(self, delta: timedelta): ... |
| 73 | + |
| 74 | + @abstractmethod |
| 75 | + def drop_duplicates(self, df, keys, sort_by, ascending: bool = False): |
| 76 | + pass |
| 77 | + |
| 78 | + @abstractmethod |
| 79 | + def rename_columns(self, df, columns: dict[str, str]): ... |
0 commit comments