fugue 0.8.2.dev1__py3-none-any.whl → 0.8.4__py3-none-any.whl

fugue/dataframe/function_wrapper.py

@@ -0,0 +1,432 @@
+import inspect
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    Iterable,
+    Iterator,
+    List,
+    Optional,
+    no_type_check,
+)
+
+import pandas as pd
+import pyarrow as pa
+from triad import Schema, assert_or_throw
+from triad.collections.function_wrapper import (
+    AnnotatedParam,
+    FunctionWrapper,
+    KeywordParam,
+    PositionalParam,
+    function_wrapper,
+)
+from triad.utils.iter import EmptyAwareIterable, make_empty_aware
+
+from ..constants import FUGUE_ENTRYPOINT
+from .array_dataframe import ArrayDataFrame
+from .arrow_dataframe import ArrowDataFrame
+from .dataframe import DataFrame, LocalDataFrame
+from .dataframe_iterable_dataframe import (
+    IterableArrowDataFrame,
+    IterablePandasDataFrame,
+    LocalDataFrameIterableDataFrame,
+)
+from .dataframes import DataFrames
+from .iterable_dataframe import IterableDataFrame
+from .pandas_dataframe import PandasDataFrame
+
+
+@function_wrapper(FUGUE_ENTRYPOINT)
+class DataFrameFunctionWrapper(FunctionWrapper):
+    @property
+    def need_output_schema(self) -> Optional[bool]:
+        return (
+            self._rt.need_schema()
+            if isinstance(self._rt, _DataFrameParamBase)
+            else False
+        )
+
+    def get_format_hint(self) -> Optional[str]:
+        for v in self._params.values():
+            if isinstance(v, _DataFrameParamBase):
+                if v.format_hint() is not None:
+                    return v.format_hint()
+        if isinstance(self._rt, _DataFrameParamBase):
+            return self._rt.format_hint()
+        return None
+
+    def run(  # noqa: C901
+        self,
+        args: List[Any],
+        kwargs: Dict[str, Any],
+        ignore_unknown: bool = False,
+        output_schema: Any = None,
+        output: bool = True,
+        ctx: Any = None,
+    ) -> Any:
+        p: Dict[str, Any] = {}
+        for i in range(len(args)):
+            p[self._params.get_key_by_index(i)] = args[i]
+        p.update(kwargs)
+        has_kw = False
+        rargs: Dict[str, Any] = {}
+        for k, v in self._params.items():
+            if isinstance(v, (PositionalParam, KeywordParam)):
+                if isinstance(v, KeywordParam):
+                    has_kw = True
+            elif k in p:
+                if isinstance(v, _DataFrameParamBase):
+                    assert_or_throw(
+                        isinstance(p[k], DataFrame),
+                        lambda: TypeError(f"{p[k]} is not a DataFrame"),
+                    )
+                    rargs[k] = v.to_input_data(p[k], ctx=ctx)
+                else:
+                    rargs[k] = p[k]  # TODO: should we do auto type conversion?
+                del p[k]
+            elif v.required:
+                raise ValueError(f"{k} is required by not given")
+        if has_kw:
+            rargs.update(p)
+        elif not ignore_unknown and len(p) > 0:
+            raise ValueError(f"{p} are not acceptable parameters")
+        rt = self._func(**rargs)
+        if not output:
+            if isinstance(self._rt, _DataFrameParamBase):
+                self._rt.count(rt)
+            return
+        if isinstance(self._rt, _DataFrameParamBase):
+            return self._rt.to_output_df(rt, output_schema, ctx=ctx)
+        return rt
+
+
+fugue_annotated_param = DataFrameFunctionWrapper.annotated_param
+
+
+@fugue_annotated_param(
+    "Callable",
+    "F",
+    lambda annotation: (
+        annotation == Callable
+        or annotation == callable  # pylint: disable=comparison-with-callable
+        or str(annotation).startswith("typing.Callable")
+    ),
+)
+class _CallableParam(AnnotatedParam):
+    pass
+
+
+@fugue_annotated_param(
+    "Callable",
+    "f",
+    lambda annotation: (
+        annotation == Optional[Callable]
+        or annotation == Optional[callable]
+        or str(annotation).startswith("typing.Union[typing.Callable")  # 3.8-
+        or str(annotation).startswith("typing.Optional[typing.Callable")  # 3.9+
+    ),
+)
+class _OptionalCallableParam(AnnotatedParam):
+    pass
+
+
+class _DataFrameParamBase(AnnotatedParam):
+    def __init__(self, param: Optional[inspect.Parameter]):
+        super().__init__(param)
+        assert_or_throw(self.required, lambda: TypeError(f"{self} must be required"))
+
+    def to_input_data(self, df: DataFrame, ctx: Any) -> Any:  # pragma: no cover
+        raise NotImplementedError
+
+    def to_output_df(
+        self, df: Any, schema: Any, ctx: Any
+    ) -> DataFrame:  # pragma: no cover
+        raise NotImplementedError
+
+    def count(self, df: Any) -> int:  # pragma: no cover
+        raise NotImplementedError
+
+    def need_schema(self) -> Optional[bool]:
+        return False
+
+    def format_hint(self) -> Optional[str]:
+        return None
+
+
+@fugue_annotated_param(DataFrame, "d", child_can_reuse_code=True)
+class DataFrameParam(_DataFrameParamBase):
+    def to_input_data(self, df: DataFrame, ctx: Any) -> Any:
+        return df
+
+    def to_output_df(self, output: Any, schema: Any, ctx: Any) -> DataFrame:
+        assert_or_throw(
+            schema is None or output.schema == schema,
+            lambda: f"Output schema mismatch {output.schema} vs {schema}",
+        )
+        return output
+
+    def count(self, df: Any) -> int:
+        if df.is_bounded:
+            return df.count()
+        else:
+            return sum(1 for _ in df.as_array_iterable())
+
+
+@fugue_annotated_param(LocalDataFrame, "l", child_can_reuse_code=True)
+class LocalDataFrameParam(DataFrameParam):
+    def to_input_data(self, df: DataFrame, ctx: Any) -> LocalDataFrame:
+        return df.as_local()
+
+    def to_output_df(self, output: LocalDataFrame, schema: Any, ctx: Any) -> DataFrame:
+        assert_or_throw(
+            schema is None or output.schema == schema,
+            lambda: f"Output schema mismatch {output.schema} vs {schema}",
+        )
+        return output
+
+    def count(self, df: LocalDataFrame) -> int:
+        if df.is_bounded:
+            return df.count()
+        else:
+            return sum(1 for _ in df.as_array_iterable())
+
+
+@fugue_annotated_param(
+    "[NoSchema]", "s", matcher=lambda x: False, child_can_reuse_code=True
+)
+class _LocalNoSchemaDataFrameParam(LocalDataFrameParam):
+    def need_schema(self) -> Optional[bool]:
+        return True
+
+
+@fugue_annotated_param(List[List[Any]])
+class _ListListParam(_LocalNoSchemaDataFrameParam):
+    @no_type_check
+    def to_input_data(self, df: DataFrame, ctx: Any) -> List[List[Any]]:
+        return df.as_array(type_safe=True)
+
+    @no_type_check
+    def to_output_df(self, output: List[List[Any]], schema: Any, ctx: Any) -> DataFrame:
+        return ArrayDataFrame(output, schema)
+
+    @no_type_check
+    def count(self, df: List[List[Any]]) -> int:
+        return len(df)
+
+
+@fugue_annotated_param(
+    Iterable[List[Any]],
+    matcher=lambda x: x == Iterable[List[Any]] or x == Iterator[List[Any]],
+)
+class _IterableListParam(_LocalNoSchemaDataFrameParam):
+    @no_type_check
+    def to_input_data(self, df: DataFrame, ctx: Any) -> Iterable[List[Any]]:
+        return df.as_array_iterable(type_safe=True)
+
+    @no_type_check
+    def to_output_df(
+        self, output: Iterable[List[Any]], schema: Any, ctx: Any
+    ) -> DataFrame:
+        return IterableDataFrame(output, schema)
+
+    @no_type_check
+    def count(self, df: Iterable[List[Any]]) -> int:
+        return sum(1 for _ in df)
+
+
+@fugue_annotated_param(EmptyAwareIterable[List[Any]])
+class _EmptyAwareIterableListParam(_LocalNoSchemaDataFrameParam):
+    @no_type_check
+    def to_input_data(self, df: DataFrame, ctx: Any) -> EmptyAwareIterable[List[Any]]:
+        return make_empty_aware(df.as_array_iterable(type_safe=True))
+
+    @no_type_check
+    def to_output_df(
+        self, output: EmptyAwareIterable[List[Any]], schema: Any, ctx: Any
+    ) -> DataFrame:
+        return IterableDataFrame(output, schema)
+
+    @no_type_check
+    def count(self, df: EmptyAwareIterable[List[Any]]) -> int:
+        return sum(1 for _ in df)
+
+
+@fugue_annotated_param(List[Dict[str, Any]])
+class _ListDictParam(_LocalNoSchemaDataFrameParam):
+    @no_type_check
+    def to_input_data(self, df: DataFrame, ctx: Any) -> List[Dict[str, Any]]:
+        return list(df.as_local().as_dict_iterable())
+
+    @no_type_check
+    def to_output_df(
+        self, output: List[Dict[str, Any]], schema: Any, ctx: Any
+    ) -> DataFrame:
+        schema = schema if isinstance(schema, Schema) else Schema(schema)
+
+        def get_all() -> Iterable[List[Any]]:
+            for row in output:
+                yield [row[x] for x in schema.names]
+
+        return IterableDataFrame(get_all(), schema)
+
+    @no_type_check
+    def count(self, df: List[Dict[str, Any]]) -> int:
+        return len(df)
+
+
+@fugue_annotated_param(
+    Iterable[Dict[str, Any]],
+    matcher=lambda x: x == Iterable[Dict[str, Any]] or x == Iterator[Dict[str, Any]],
+)
+class _IterableDictParam(_LocalNoSchemaDataFrameParam):
+    @no_type_check
+    def to_input_data(self, df: DataFrame, ctx: Any) -> Iterable[Dict[str, Any]]:
+        return df.as_dict_iterable()
+
+    @no_type_check
+    def to_output_df(
+        self, output: Iterable[Dict[str, Any]], schema: Any, ctx: Any
+    ) -> DataFrame:
+        schema = schema if isinstance(schema, Schema) else Schema(schema)
+
+        def get_all() -> Iterable[List[Any]]:
+            for row in output:
+                yield [row[x] for x in schema.names]
+
+        return IterableDataFrame(get_all(), schema)
+
+    @no_type_check
+    def count(self, df: Iterable[Dict[str, Any]]) -> int:
+        return sum(1 for _ in df)
+
+
+@fugue_annotated_param(EmptyAwareIterable[Dict[str, Any]])
+class _EmptyAwareIterableDictParam(_LocalNoSchemaDataFrameParam):
+    @no_type_check
+    def to_input_data(
+        self, df: DataFrame, ctx: Any
+    ) -> EmptyAwareIterable[Dict[str, Any]]:
+        return make_empty_aware(df.as_dict_iterable())
+
+    @no_type_check
+    def to_output_df(
+        self, output: EmptyAwareIterable[Dict[str, Any]], schema: Any, ctx: Any
+    ) -> DataFrame:
+        schema = schema if isinstance(schema, Schema) else Schema(schema)
+
+        def get_all() -> Iterable[List[Any]]:
+            for row in output:
+                yield [row[x] for x in schema.names]
+
+        return IterableDataFrame(get_all(), schema)
+
+    @no_type_check
+    def count(self, df: EmptyAwareIterable[Dict[str, Any]]) -> int:
+        return sum(1 for _ in df)
+
+
+@fugue_annotated_param(pd.DataFrame, "p")
+class _PandasParam(LocalDataFrameParam):
+    @no_type_check
+    def to_input_data(self, df: DataFrame, ctx: Any) -> pd.DataFrame:
+        return df.as_pandas()
+
+    @no_type_check
+    def to_output_df(self, output: pd.DataFrame, schema: Any, ctx: Any) -> DataFrame:
+        return PandasDataFrame(output, schema)
+
+    @no_type_check
+    def count(self, df: pd.DataFrame) -> int:
+        return df.shape[0]
+
+    def format_hint(self) -> Optional[str]:
+        return "pandas"
+
+
+@fugue_annotated_param(
+    Iterable[pd.DataFrame],
+    matcher=lambda x: x == Iterable[pd.DataFrame] or x == Iterator[pd.DataFrame],
+)
+class _IterablePandasParam(LocalDataFrameParam):
+    @no_type_check
+    def to_input_data(self, df: DataFrame, ctx: Any) -> Iterable[pd.DataFrame]:
+        if not isinstance(df, LocalDataFrameIterableDataFrame):
+            yield df.as_pandas()
+        else:
+            for sub in df.native:
+                yield sub.as_pandas()
+
+    @no_type_check
+    def to_output_df(
+        self, output: Iterable[pd.DataFrame], schema: Any, ctx: Any
+    ) -> DataFrame:
+        def dfs():
+            for df in output:
+                yield PandasDataFrame(df, schema)
+
+        return IterablePandasDataFrame(dfs())
+
+    @no_type_check
+    def count(self, df: Iterable[pd.DataFrame]) -> int:
+        return sum(_.shape[0] for _ in df)
+
+    def format_hint(self) -> Optional[str]:
+        return "pandas"
+
+
+@fugue_annotated_param(pa.Table)
+class _PyArrowTableParam(LocalDataFrameParam):
+    def to_input_data(self, df: DataFrame, ctx: Any) -> Any:
+        return df.as_arrow()
+
+    def to_output_df(self, output: Any, schema: Any, ctx: Any) -> DataFrame:
+        assert isinstance(output, pa.Table)
+        return ArrowDataFrame(output, schema=schema)
+
+    def count(self, df: Any) -> int:  # pragma: no cover
+        return df.count()
+
+    def format_hint(self) -> Optional[str]:
+        return "pyarrow"
+
+
+@fugue_annotated_param(
+    Iterable[pa.Table],
+    matcher=lambda x: x == Iterable[pa.Table] or x == Iterator[pa.Table],
+)
+class _IterableArrowParam(LocalDataFrameParam):
+    @no_type_check
+    def to_input_data(self, df: DataFrame, ctx: Any) -> Iterable[pa.Table]:
+        if not isinstance(df, LocalDataFrameIterableDataFrame):
+            yield df.as_arrow()
+        else:
+            for sub in df.native:
+                yield sub.as_arrow()
+
+    @no_type_check
+    def to_output_df(
+        self, output: Iterable[pa.Table], schema: Any, ctx: Any
+    ) -> DataFrame:
+        def dfs():
+            _schema: Optional[Schema] = None if schema is None else Schema(schema)
+            for df in output:
+                adf = ArrowDataFrame(df)
+                if _schema is not None and not (  # pylint: disable-all
+                    adf.schema == schema
+                ):
+                    adf = adf[_schema.names].alter_columns(_schema)
+                yield adf
+
+        return IterableArrowDataFrame(dfs())
+
+    @no_type_check
+    def count(self, df: Iterable[pa.Table]) -> int:
+        return sum(_.shape[0] for _ in df)
+
+    def format_hint(self) -> Optional[str]:
+        return "pyarrow"
+
+
+@fugue_annotated_param(DataFrames, "c")
+class _DataFramesParam(AnnotatedParam):
+    pass

fugue/dataframe/iterable_dataframe.py

@@ -97,6 +97,9 @@ class IterableDataFrame(LocalUnboundedDataFrame):
             return self
         return IterableDataFrame(self.native, new_schema)
 
+    def as_local_bounded(self) -> LocalBoundedDataFrame:
+        return ArrayDataFrame(self.as_array(), schema=self.schema)
+
     def as_array(
         self, columns: Optional[List[str]] = None, type_safe: bool = False
     ) -> List[Any]:

fugue/dataframe/utils.py

@@ -13,11 +13,9 @@ from triad.exceptions import InvalidOperationError
 from triad.utils.assertion import assert_arg_not_none
 from triad.utils.assertion import assert_or_throw as aot
 
-from .api import get_column_names, normalize_column_names, rename
+from .api import get_column_names, normalize_column_names, rename, as_fugue_df
 from .array_dataframe import ArrayDataFrame
-from .arrow_dataframe import ArrowDataFrame
-from .dataframe import DataFrame, LocalBoundedDataFrame, LocalDataFrame
-from .iterable_dataframe import IterableDataFrame
+from .dataframe import DataFrame, LocalBoundedDataFrame
 from .pandas_dataframe import PandasDataFrame
 
 # For backward compatibility, TODO: remove!
@@ -29,7 +27,7 @@ rename_dataframe_column_names = rename
 def _pa_type_eq(t1: pa.DataType, t2: pa.DataType) -> bool:
     # should ignore the name difference of list
     # e.g. list<item: string> == list<l: string>
-    if pa.types.is_list(t1) and pa.types.is_list(t2):
+    if pa.types.is_list(t1) and pa.types.is_list(t2):  # pragma: no cover
         return _pa_type_eq(t1.value_type, t2.value_type)
     return t1 == t2
 
@@ -74,8 +72,11 @@ def _df_eq(
     :param throw: if to throw error if not equal, defaults to False
     :return: if they equal
     """
-    df1 = to_local_bounded_df(df)
-    df2 = to_local_bounded_df(data, schema)
+    df1 = as_fugue_df(df).as_local_bounded()
+    if schema is not None:
+        df2 = as_fugue_df(data, schema=schema).as_local_bounded()
+    else:
+        df2 = as_fugue_df(data).as_local_bounded()
     try:
         assert (
             df1.count() == df2.count()
@@ -99,7 +100,7 @@ def _df_eq(
         d1 = d1.reset_index(drop=True)
         d2 = d2.reset_index(drop=True)
         pd.testing.assert_frame_equal(
-            d1, d2, check_less_precise=digits, check_dtype=False
+            d1, d2, rtol=0, atol=10 ** (-digits), check_dtype=False, check_exact=False
         )
         return True
     except AssertionError:
@@ -108,78 +109,9 @@ def _df_eq(
         return False
 
 
-def to_local_df(df: Any, schema: Any = None) -> LocalDataFrame:
-    """Convert a data structure to :class:`~fugue.dataframe.dataframe.LocalDataFrame`
-
-    :param df: :class:`~fugue.dataframe.dataframe.DataFrame`, pandas DataFramme and
-      list or iterable of arrays
-    :param schema: |SchemaLikeObject|, defaults to None, it should not be set for
-      :class:`~fugue.dataframe.dataframe.DataFrame` type
-    :raises ValueError: if ``df`` is :class:`~fugue.dataframe.dataframe.DataFrame`
-      but you set ``schema``
-    :raises TypeError: if ``df`` is not compatible
-    :return: the dataframe itself if it's
-      :class:`~fugue.dataframe.dataframe.LocalDataFrame` else a converted one
-
-    .. admonition:: Examples
-
-        >>> a = to_local_df([[0,'a'],[1,'b']],"a:int,b:str")
-        >>> assert to_local_df(a) is a
-        >>> to_local_df(SparkDataFrame([[0,'a'],[1,'b']],"a:int,b:str"))
-    """
-    assert_arg_not_none(df, "df")
-    if isinstance(df, DataFrame):
-        aot(
-            schema is None,
-            ValueError("schema and metadata must be None when df is a DataFrame"),
-        )
-        return df.as_local()
-    if isinstance(df, pd.DataFrame):
-        return PandasDataFrame(df, schema)
-    if isinstance(df, pa.Table):
-        return ArrowDataFrame(df, schema)
-    if isinstance(df, List):
-        return ArrayDataFrame(df, schema)
-    if isinstance(df, Iterable):
-        return IterableDataFrame(df, schema)
-    raise TypeError(f"{df} cannot convert to a LocalDataFrame")
-
-
-def to_local_bounded_df(df: Any, schema: Any = None) -> LocalBoundedDataFrame:
-    """Convert a data structure to
-    :class:`~fugue.dataframe.dataframe.LocalBoundedDataFrame`
-
-    :param df: :class:`~fugue.dataframe.dataframe.DataFrame`, pandas DataFramme and
-      list or iterable of arrays
-    :param schema: |SchemaLikeObject|, defaults to None, it should not be set for
-      :class:`~fugue.dataframe.dataframe.DataFrame` type
-    :raises ValueError: if ``df`` is :class:`~fugue.dataframe.dataframe.DataFrame`
-      but you set ``schema``
-    :raises TypeError: if ``df`` is not compatible
-    :return: the dataframe itself if it's
-      :class:`~fugue.dataframe.dataframe.LocalBoundedDataFrame` else a converted one
-
-    .. admonition:: Examples
-
-        >>> a = IterableDataFrame([[0,'a'],[1,'b']],"a:int,b:str")
-        >>> assert isinstance(to_local_bounded_df(a), LocalBoundedDataFrame)
-        >>> to_local_bounded_df(SparkDataFrame([[0,'a'],[1,'b']],"a:int,b:str"))
-
-    .. note::
-
-        Compared to :func:`.to_local_df`, this function makes sure the dataframe is also
-        bounded, so :class:`~fugue.dataframe.iterable_dataframe.IterableDataFrame` will
-        be converted although it's local.
-    """
-    df = to_local_df(df, schema)
-    if isinstance(df, LocalBoundedDataFrame):
-        return df
-    return ArrayDataFrame(df.as_array(), df.schema)
-
-
 def pickle_df(df: DataFrame) -> bytes:
     """Pickles a dataframe to bytes array. It firstly converts the dataframe
-    using :func:`.to_local_bounded_df`, and then serialize the underlying data.
+    local bounded, and then serialize the underlying data.
 
     :param df: input DataFrame
     :return: pickled binary data
@@ -189,7 +121,7 @@ def pickle_df(df: DataFrame) -> bytes:
         Be careful to use on large dataframes or non-local, un-materialized dataframes,
         it can be slow. You should always use :func:`.unpickle_df` to deserialize.
     """
-    df = to_local_bounded_df(df)
+    df = df.as_local_bounded()
     o: List[Any] = [df.schema]
     if isinstance(df, PandasDataFrame):
         o.append("p")

fugue/dataset/api.py

@@ -41,8 +41,6 @@ def as_local(data: AnyDataset) -> AnyDataset:
 
     :param data: the dataset that can be recognized by Fugue
     """
-    if isinstance(data, Dataset) and data.is_local:
-        return data
     return as_local_bounded(data)
 
 
@@ -52,8 +50,6 @@ def as_local_bounded(data: AnyDataset) -> AnyDataset:
 
     :param data: the dataset that can be recognized by Fugue
     """
-    if isinstance(data, Dataset) and data.is_local and data.is_bounded:
-        return data
     raise NotImplementedError(
         f"no registered function to convert {type(data)} to a local bounded dataset"
     )

fugue/dev.py

@@ -0,0 +1,47 @@
+"""
+All modeuls for developing and extending Fugue
+"""
+# flake8: noqa
+# pylint: disable-all
+
+from triad.collections.function_wrapper import AnnotatedParam
+
+from fugue.bag.bag import BagDisplay
+from fugue.collections.partition import PartitionCursor, PartitionSpec
+from fugue.collections.sql import StructuredRawSQL, TempTableName
+from fugue.collections.yielded import PhysicalYielded, Yielded
+from fugue.dataframe.function_wrapper import (
+    DataFrameFunctionWrapper,
+    DataFrameParam,
+    LocalDataFrameParam,
+    fugue_annotated_param,
+)
+from fugue.dataset import DatasetDisplay
+from fugue.execution.execution_engine import (
+    EngineFacet,
+    ExecutionEngineParam,
+    MapEngine,
+    SQLEngine,
+)
+from fugue.execution.factory import (
+    is_pandas_or,
+    make_execution_engine,
+    make_sql_engine,
+    register_default_execution_engine,
+    register_default_sql_engine,
+    register_execution_engine,
+    register_sql_engine,
+)
+from fugue.execution.native_execution_engine import PandasMapEngine, QPDPandasEngine
+from fugue.rpc import (
+    EmptyRPCHandler,
+    RPCClient,
+    RPCFunc,
+    RPCHandler,
+    RPCServer,
+    make_rpc_server,
+    to_rpc_handler,
+)
+from fugue.workflow._workflow_context import FugueWorkflowContext
+from fugue.workflow.module import module
+from fugue.workflow.workflow import FugueWorkflow, WorkflowDataFrame, WorkflowDataFrames

fugue/execution/__init__.py

@@ -10,8 +10,4 @@ from .factory import (
     register_execution_engine,
     register_sql_engine,
 )
-from .native_execution_engine import (
-    NativeExecutionEngine,
-    QPDPandasEngine,
-    SqliteEngine,
-)
+from .native_execution_engine import NativeExecutionEngine, QPDPandasEngine