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

fugue_polars/polars_dataframe.py

@@ -0,0 +1,234 @@
+from typing import Any, Dict, Iterable, List, Optional
+
+import pandas as pd
+import polars as pl
+import pyarrow as pa
+from triad.collections.schema import Schema
+from triad.exceptions import InvalidOperationError
+from triad.utils.assertion import assert_or_throw
+from triad.utils.pyarrow import (
+    LARGE_TYPES_REPLACEMENT,
+    replace_types_in_schema,
+    replace_types_in_table,
+)
+
+from fugue import ArrowDataFrame
+from fugue.api import (
+    as_arrow,
+    drop_columns,
+    get_column_names,
+    get_schema,
+    is_df,
+    rename,
+    select_columns,
+)
+from fugue.dataframe.dataframe import DataFrame, LocalBoundedDataFrame, _input_schema
+from fugue.dataset.api import (
+    as_local,
+    as_local_bounded,
+    count,
+    get_num_partitions,
+    is_bounded,
+    is_empty,
+    is_local,
+)
+from fugue.exceptions import FugueDataFrameOperationError
+
+from ._utils import build_empty_pl
+
+
+class PolarsDataFrame(LocalBoundedDataFrame):
+    """DataFrame that wraps :func:`pyarrow.Table <pa:pyarrow.table>`. Please also read
+    |DataFrameTutorial| to understand this Fugue concept
+
+    :param df: polars DataFrame or None, defaults to None
+    :param schema: |SchemaLikeObject|
+    """
+
+    def __init__(
+        self,
+        df: Optional[pl.DataFrame] = None,
+        schema: Any = None,
+    ):
+        if df is None:
+            schema = _input_schema(schema).assert_not_empty()
+            self._native: pa.Table = build_empty_pl(schema)
+            super().__init__(schema)
+            return
+        else:
+            assert_or_throw(
+                schema is None,
+                InvalidOperationError("can't reset schema for pl.DataFrame"),
+            )
+            self._native = df
+            super().__init__(_get_pl_schema(df))
+
+    @property
+    def native(self) -> pl.DataFrame:
+        """:func:`pyarrow.Table <pa:pyarrow.table>`"""
+        return self._native
+
+    def native_as_df(self) -> pl.DataFrame:
+        return self._native
+
+    @property
+    def empty(self) -> bool:
+        return self._native.shape[0] == 0
+
+    def peek_array(self) -> List[Any]:
+        self.assert_not_empty()
+        return list(self._native.row(0))
+
+    def peek_dict(self) -> Dict[str, Any]:
+        self.assert_not_empty()
+        return self._native.row(0, named=True)
+
+    def count(self) -> int:
+        return self.native.shape[0]
+
+    def as_pandas(self) -> pd.DataFrame:
+        return self.native.to_pandas()
+
+    def head(
+        self, n: int, columns: Optional[List[str]] = None
+    ) -> LocalBoundedDataFrame:
+        adf = self.native if columns is None else self.native.select(columns)
+        n = min(n, self.count())
+        if n == 0:
+            schema = self.schema if columns is None else self.schema.extract(columns)
+            return PolarsDataFrame(None, schema=schema)
+        return PolarsDataFrame(adf.head(n))
+
+    def _drop_cols(self, cols: List[str]) -> DataFrame:
+        return PolarsDataFrame(self.native.drop(cols))
+
+    def _select_cols(self, keys: List[Any]) -> DataFrame:
+        return PolarsDataFrame(self.native.select(keys))
+
+    def rename(self, columns: Dict[str, str]) -> DataFrame:
+        return PolarsDataFrame(_rename_pl_dataframe(self.native, columns))
+
+    def alter_columns(self, columns: Any) -> DataFrame:
+        adf = ArrowDataFrame(self.as_arrow()).alter_columns(columns)
+        return PolarsDataFrame(pl.from_arrow(adf.native))
+
+    def as_arrow(self, type_safe: bool = False) -> pa.Table:
+        return _pl_as_arrow(self.native)
+
+    def as_array(
+        self, columns: Optional[List[str]] = None, type_safe: bool = False
+    ) -> List[Any]:
+        tdf = self.native
+        if columns is not None:
+            tdf = tdf.select(columns)
+        return [list(row) for row in tdf.rows()]
+
+    def as_array_iterable(
+        self, columns: Optional[List[str]] = None, type_safe: bool = False
+    ) -> Iterable[Any]:
+        if not self.empty:
+            yield from ArrowDataFrame(_pl_as_arrow(self.native)).as_array_iterable(
+                columns=columns
+            )
+
+    def as_dict_iterable(
+        self, columns: Optional[List[str]] = None
+    ) -> Iterable[Dict[str, Any]]:
+        if not self.empty:
+            yield from ArrowDataFrame(_pl_as_arrow(self.native)).as_dict_iterable(
+                columns=columns
+            )
+
+
+@as_local.candidate(lambda df: isinstance(df, pl.DataFrame))
+def _pl_as_local(df: pl.DataFrame) -> pl.DataFrame:
+    return df
+
+
+@as_local_bounded.candidate(lambda df: isinstance(df, pl.DataFrame))
+def _pl_as_local_bounded(df: pl.DataFrame) -> pl.DataFrame:
+    return df
+
+
+@as_arrow.candidate(lambda df: isinstance(df, pl.DataFrame))
+def _pl_as_arrow(df: pl.DataFrame) -> pa.Table:
+    adf = df.to_arrow()
+    adf = replace_types_in_table(adf, LARGE_TYPES_REPLACEMENT)
+    return adf
+
+
+@is_df.candidate(lambda df: isinstance(df, pl.DataFrame))
+def _pl_is_df(df: pl.DataFrame) -> bool:
+    return True
+
+
+@count.candidate(lambda df: isinstance(df, pl.DataFrame))
+def _pl_count(df: pl.DataFrame) -> int:
+    return df.shape[0]
+
+
+@is_bounded.candidate(lambda df: isinstance(df, pl.DataFrame))
+def _pl_is_bounded(df: pl.DataFrame) -> bool:
+    return True
+
+
+@is_empty.candidate(lambda df: isinstance(df, pl.DataFrame))
+def _pl_is_empty(df: pl.DataFrame) -> bool:
+    return df.shape[0] == 0
+
+
+@is_local.candidate(lambda df: isinstance(df, pl.DataFrame))
+def _pl_is_local(df: pl.DataFrame) -> bool:
+    return True
+
+
+@get_num_partitions.candidate(lambda df: isinstance(df, pl.DataFrame))
+def _pl_get_num_partitions(df: pl.DataFrame) -> int:
+    return 1
+
+
+@get_column_names.candidate(lambda df: isinstance(df, pl.DataFrame))
+def _get_pl_columns(df: pl.DataFrame) -> List[Any]:
+    return list(df.schema.keys())
+
+
+@get_schema.candidate(lambda df: isinstance(df, pl.DataFrame))
+def _get_pl_schema(df: pl.DataFrame) -> Schema:
+    adf = df.to_arrow()
+    schema = replace_types_in_schema(adf.schema, LARGE_TYPES_REPLACEMENT)
+    return Schema(schema)
+
+
+@rename.candidate(lambda df, *args, **kwargs: isinstance(df, pl.DataFrame))
+def _rename_pl_dataframe(df: pl.DataFrame, columns: Dict[str, Any]) -> pl.DataFrame:
+    if len(columns) == 0:
+        return df
+    assert_or_throw(
+        set(columns.keys()).issubset(set(df.columns)),
+        FugueDataFrameOperationError(f"invalid {columns}"),
+    )
+    return df.rename(columns)
+
+
+@drop_columns.candidate(lambda df, *args, **kwargs: isinstance(df, pl.DataFrame))
+def _drop_pa_columns(df: pl.DataFrame, columns: List[str]) -> pl.DataFrame:
+    cols = [x for x in df.schema.keys() if x not in columns]
+    if len(cols) == 0:
+        raise FugueDataFrameOperationError("cannot drop all columns")
+    if len(cols) + len(columns) != len(df.columns):
+        _assert_no_missing(df, columns)
+    return df.select(cols)
+
+
+@select_columns.candidate(lambda df, *args, **kwargs: isinstance(df, pl.DataFrame))
+def _select_pa_columns(df: pl.DataFrame, columns: List[Any]) -> pl.DataFrame:
+    if len(columns) == 0:
+        raise FugueDataFrameOperationError("must select at least one column")
+    _assert_no_missing(df, columns=columns)
+    return df.select(columns)
+
+
+def _assert_no_missing(df: pl.DataFrame, columns: Iterable[Any]) -> None:
+    missing = [x for x in columns if x not in df.schema.keys()]
+    if len(missing) > 0:
+        raise FugueDataFrameOperationError("found nonexistent columns: {missing}")

fugue_polars/registry.py

@@ -0,0 +1,86 @@
+from typing import Any, Iterable, Iterator, Optional, no_type_check
+
+import polars as pl
+import pyarrow as pa
+from triad import Schema, make_empty_aware
+from triad.utils.pyarrow import get_alter_func
+
+from fugue import (
+    ArrowDataFrame,
+    DataFrame,
+    IterableArrowDataFrame,
+    LocalDataFrameIterableDataFrame,
+)
+from fugue.dev import LocalDataFrameParam, fugue_annotated_param
+from .polars_dataframe import PolarsDataFrame
+from fugue.plugins import as_fugue_dataset
+
+
+@as_fugue_dataset.candidate(lambda df, **kwargs: isinstance(df, pl.DataFrame))
+def _pl_as_fugue_df(df: pl.DataFrame, **kwargs: Any) -> PolarsDataFrame:
+    return PolarsDataFrame(df, **kwargs)
+
+
+@fugue_annotated_param(pl.DataFrame)
+class _PolarsParam(LocalDataFrameParam):
+    def to_input_data(self, df: DataFrame, ctx: Any) -> Any:
+        return pl.from_arrow(df.as_arrow())
+
+    def to_output_df(self, output: Any, schema: Any, ctx: Any) -> DataFrame:
+        assert isinstance(output, pl.DataFrame)
+        return _to_adf(output, schema=schema)
+
+    def count(self, df: Any) -> int:  # pragma: no cover
+        return df.shape[0]
+
+    def format_hint(self) -> Optional[str]:
+        return "pyarrow"
+
+
+@fugue_annotated_param(
+    Iterable[pl.DataFrame],
+    matcher=lambda x: x == Iterable[pl.DataFrame] or x == Iterator[pl.DataFrame],
+)
+class _IterablePolarsParam(LocalDataFrameParam):
+    @no_type_check
+    def to_input_data(self, df: DataFrame, ctx: Any) -> Iterable[pa.Table]:
+        if not isinstance(df, LocalDataFrameIterableDataFrame):
+            yield pl.from_arrow(df.as_arrow())
+        else:  # pragma: no cover # spark code coverage can't be included
+            for sub in df.native:
+                yield pl.from_arrow(sub.as_arrow())
+
+    @no_type_check
+    def to_output_df(
+        self, output: Iterable[pl.DataFrame], schema: Any, ctx: Any
+    ) -> DataFrame:
+        def dfs(_schema: Schema) -> Iterable[ArrowDataFrame]:
+            if output is not None:
+                for df in output:
+                    yield _to_adf(df, _schema)
+
+        _schema: Optional[Schema] = (
+            None
+            if schema is None
+            else (schema if isinstance(schema, Schema) else Schema(schema))
+        )
+        _dfs = make_empty_aware(dfs(_schema))
+        if not _dfs.empty:
+            return IterableArrowDataFrame(_dfs)
+        return IterableArrowDataFrame([], schema=_schema)
+
+    @no_type_check
+    def count(self, df: Iterable[pl.DataFrame]) -> int:  # pragma: no cover
+        return sum(_.shape[0] for _ in df)
+
+    def format_hint(self) -> Optional[str]:
+        return "pyarrow"
+
+
+def _to_adf(output: pl.DataFrame, schema: Any) -> ArrowDataFrame:
+    adf = output.to_arrow()
+    if schema is None:  # pragma: no cover
+        return ArrowDataFrame(adf)
+    _schema = schema if isinstance(schema, Schema) else Schema(schema)
+    f = get_alter_func(adf.schema, _schema.pa_schema, safe=False)
+    return ArrowDataFrame(f(adf))

fugue_ray/_constants.py

@@ -1,10 +1,19 @@
-from typing import Dict, Any
+from typing import Any, Dict
+
+import ray
 
 FUGUE_RAY_CONF_SHUFFLE_PARTITIONS = "fugue.ray.shuffle.partitions"
 FUGUE_RAY_DEFAULT_PARTITIONS = "fugue.ray.default.partitions"
 FUGUE_RAY_DEFAULT_BATCH_SIZE = "fugue.ray.default.batch_size"
+FUGUE_RAY_ZERO_COPY = "fugue.ray.zero_copy"
 
 FUGUE_RAY_DEFAULT_CONF: Dict[str, Any] = {
     FUGUE_RAY_CONF_SHUFFLE_PARTITIONS: -1,
     FUGUE_RAY_DEFAULT_PARTITIONS: 0,
+    FUGUE_RAY_ZERO_COPY: True,
 }
+
+if ray.__version__ >= "2.3":
+    _ZERO_COPY: Dict[str, Any] = {"zero_copy_batch": True}
+else:  # pragma: no cover
+    _ZERO_COPY = {}

fugue_ray/_utils/dataframe.py

@@ -1,21 +1,27 @@
 import pickle
-from typing import List, Optional, Tuple, Dict, Any
+from typing import Any, Dict, List, Optional, Tuple
 
+import pandas as pd
 import pyarrow as pa
 import ray.data as rd
-from fugue.dataframe.arrow_dataframe import _build_empty_arrow
 from triad import Schema
 
+from fugue.dataframe.arrow_dataframe import _build_empty_arrow
+
+from .._constants import _ZERO_COPY
+
 _RAY_NULL_REPR = "__RAY_NULL__"
 
 
 def get_dataset_format(df: rd.Dataset) -> Optional[str]:
-    try:  # pragma: no cover
-        if hasattr(df, "_dataset_format"):  # ray<2.2
-            return df._dataset_format()
-        return df.dataset_format()  # ray>=2.2
-    except Exception:
+    df.fully_executed()
+    if df.count() == 0:
         return None
+    if hasattr(df, "_dataset_format"):  # pragma: no cover
+        return df._dataset_format()  # ray<2.2
+    ctx = rd.context.DatasetContext.get_current()
+    ctx.use_streaming_executor = False
+    return df.dataset_format()  # ray>=2.2
 
 
 def build_empty(schema: Schema) -> rd.Dataset:
@@ -50,7 +56,7 @@ def add_partition_key(
             )
 
         return df.map_batches(
-            add_simple_key, batch_format="pyarrow", **ray_remote_args
+            add_simple_key, batch_format="pyarrow", **_ZERO_COPY, **ray_remote_args
         ), input_schema + (
             output_key,
             str,
@@ -67,8 +73,29 @@ def add_partition_key(
             return fdf.append_column(output_key, sarr)
 
         return df.map_batches(
-            add_key, batch_format="pyarrow", **ray_remote_args
+            add_key, batch_format="pyarrow", **_ZERO_COPY, **ray_remote_args
         ), input_schema + (
             output_key,
             pa.binary(),
         )
+
+
+def add_coarse_partition_key(
+    df: rd.Dataset,
+    keys: List[str],
+    output_key: str,
+    bucket: int,
+) -> rd.Dataset:
+    ray_remote_args: Dict[str, Any] = {"num_cpus": 1}
+
+    def add_coarse_key(arrow_df: pa.Table) -> pa.Table:  # pragma: no cover
+        hdf = arrow_df.select(keys).to_pandas()
+        _hash = pd.util.hash_pandas_object(hdf, index=False).mod(bucket)
+        return arrow_df.append_column(output_key, pa.Array.from_pandas(_hash))
+
+    return df.map_batches(
+        add_coarse_key,
+        batch_format="pyarrow",
+        **_ZERO_COPY,
+        **ray_remote_args,
+    )

fugue_ray/_utils/io.py

@@ -5,7 +5,7 @@ from typing import Any, Callable, Dict, Iterable, List, Optional, Union
 import pyarrow as pa
 import ray.data as rd
 from fugue import ExecutionEngine
-from fugue._utils.io import FileParser, load_df, save_df
+from fugue._utils.io import FileParser, save_df
 from fugue.collections.partition import PartitionSpec
 from fugue.dataframe import DataFrame
 from fugue_ray.dataframe import RayDataFrame
@@ -49,8 +49,6 @@ class RayIO(object):
             len(fmts) == 1, NotImplementedError("can't support multiple formats")
         )
         fmt = fmts[0]
-        if fmt == "avro":  # TODO: remove avro support
-            return load_df(uri, format_hint=format_hint, columns=columns, **kwargs)
         files = [f.uri for f in fp]
         return self._loads[fmt](files, columns, **kwargs)
 
@@ -75,7 +73,7 @@ class RayIO(object):
                 except Exception:  # pragma: no cover
                     pass
         p = FileParser(uri, format_hint)
-        if not force_single and p.file_format != "avro":
+        if not force_single:
             df = self._prepartition(df, partition_spec=partition_spec)
 
             self._saves[p.file_format](df=df, uri=p.uri, **kwargs)

fugue_ray/dataframe.py

@@ -6,12 +6,7 @@ import ray
 import ray.data as rd
 from triad.collections.schema import Schema
 
-from fugue.dataframe import (
-    ArrowDataFrame,
-    DataFrame,
-    LocalBoundedDataFrame,
-    LocalDataFrame,
-)
+from fugue.dataframe import ArrowDataFrame, DataFrame, LocalBoundedDataFrame
 from fugue.dataframe.dataframe import _input_schema
 from fugue.exceptions import FugueDataFrameOperationError, FugueDatasetEmptyError
 from fugue.plugins import (
@@ -22,6 +17,7 @@ from fugue.plugins import (
     rename,
 )
 
+from ._constants import _ZERO_COPY
 from ._utils.dataframe import build_empty, get_dataset_format
 
 
@@ -115,7 +111,7 @@ class RayDataFrame(DataFrame):
     def is_local(self) -> bool:
         return False
 
-    def as_local(self) -> LocalDataFrame:
+    def as_local_bounded(self) -> LocalBoundedDataFrame:
         adf = self.as_arrow()
         if adf.shape[0] == 0:
             res = ArrowDataFrame([], self.schema)
@@ -145,7 +141,10 @@ class RayDataFrame(DataFrame):
         if cols == self.columns:
             return self
         rdf = self.native.map_batches(
-            lambda b: b.select(cols), batch_format="pyarrow", **self._remote_args()
+            lambda b: b.select(cols),
+            batch_format="pyarrow",
+            **_ZERO_COPY,
+            **self._remote_args(),
         )
         return RayDataFrame(rdf, self.schema.extract(cols), internal_schema=True)
 
@@ -158,7 +157,7 @@ class RayDataFrame(DataFrame):
     def persist(self, **kwargs: Any) -> "RayDataFrame":
         # TODO: it mutates the dataframe, is this a good bahavior
         if not self.native.is_fully_executed():  # pragma: no cover
-            self._native = self.native.fully_executed()
+            self.native.fully_executed()
         return self
 
     def count(self) -> int:
@@ -179,6 +178,7 @@ class RayDataFrame(DataFrame):
         rdf = self.native.map_batches(
             lambda b: b.rename_columns(new_cols),
             batch_format="pyarrow",
+            **_ZERO_COPY,
             **self._remote_args(),
         )
         return RayDataFrame(rdf, schema=new_schema, internal_schema=True)
@@ -193,7 +193,7 @@ class RayDataFrame(DataFrame):
         if self.schema == new_schema:
             return self
         rdf = self.native.map_batches(
-            _alter, batch_format="pyarrow", **self._remote_args()
+            _alter, batch_format="pyarrow", **_ZERO_COPY, **self._remote_args()
         )
         return RayDataFrame(rdf, schema=new_schema, internal_schema=True)
 
@@ -236,7 +236,9 @@ class RayDataFrame(DataFrame):
             return ArrowDataFrame(table).alter_columns(schema).native  # type: ignore
 
         return (
-            rdf.map_batches(_alter, batch_format="pyarrow", **self._remote_args()),
+            rdf.map_batches(
+                _alter, batch_format="pyarrow", **_ZERO_COPY, **self._remote_args()
+            ),
             schema,
         )
 
@@ -278,7 +280,9 @@ def _rename_ray_dataframe(df: rd.Dataset, columns: Dict[str, Any]) -> rd.Dataset
     if len(missing) > 0:
         raise FugueDataFrameOperationError("found nonexistent columns: {missing}")
     new_cols = [columns.get(name, name) for name in cols]
-    return df.map_batches(lambda b: b.rename_columns(new_cols), batch_format="pyarrow")
+    return df.map_batches(
+        lambda b: b.rename_columns(new_cols), batch_format="pyarrow", **_ZERO_COPY
+    )
 
 
 def _get_arrow_tables(df: rd.Dataset) -> Iterable[pa.Table]: