# feast_data_stream.pyi
from typing import List, Optional, Type, Union

from denormalized._internal import PyDataStream
from denormalized.datafusion import Expr
from feast import FeatureStore, Field

from .data_stream import DataStream

class FeastDataStreamMeta(type):
    def __new__(cls, name: str, bases: tuple, attrs: dict) -> Type[FeastDataStream]: ...

class FeastDataStream(DataStream, metaclass=FeastDataStreamMeta):
    def __init__(self, stream: Union[PyDataStream, DataStream]) -> None: ...
    def select(self, expr_list: List[Expr]) -> FeastDataStream: ...
    def filter(self, predicate: Expr) -> FeastDataStream: ...
    def with_column(self, name: str, predicate: Expr) -> FeastDataStream: ...
    def join_on(
        self, right: DataStream, join_type: str, on_exprs: List[Expr]
    ) -> FeastDataStream: ...
    def join(
        self,
        right: DataStream,
        join_type: str,
        left_cols: List[str],
        right_cols: List[str],
        filter: Optional[Expr] = None,
    ) -> FeastDataStream: ...
    def window(
        self,
        group_exprs: List[Expr],
        aggr_exprs: List[Expr],
        window_length_millis: int,
        slide_millis: Optional[int] = None,
    ) -> FeastDataStream: ...
    def print_schema(self) -> FeastDataStream: ...
    def print_plan(self) -> FeastDataStream: ...
    def print_physical_plan(self) -> FeastDataStream: ...
    def get_feast_schema(self) -> List[Field]: ...
    def write_feast_feature(
        self, feature_store: FeatureStore, source_name: str
    ) -> None: ...
