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

fugue_dask/execution_engine.py

@@ -3,6 +3,7 @@ import os
 from typing import Any, Callable, Dict, List, Optional, Type, Union
 
 import dask.dataframe as dd
+import pandas as pd
 from distributed import Client
 from qpd_dask import run_sql_on_dask
 from triad.collections import Schema
@@ -18,7 +19,7 @@ from fugue.collections.partition import (
     PartitionSpec,
     parse_presort_exp,
 )
-from fugue.constants import KEYWORD_CORECOUNT, KEYWORD_ROWCOUNT
+from fugue.constants import KEYWORD_PARALLELISM, KEYWORD_ROWCOUNT
 from fugue.dataframe import (
     AnyDataFrame,
     DataFrame,
@@ -34,6 +35,8 @@ from fugue_dask._io import load_df, save_df
 from fugue_dask._utils import DASK_UTILS, DaskUtils
 from fugue_dask.dataframe import DaskDataFrame
 
+_DASK_PARTITION_KEY = "__dask_partition_key__"
+
 
 class QPDDaskEngine(SQLEngine):
     """QPD execution implementation."""
@@ -74,7 +77,8 @@ class DaskMapEngine(MapEngine):
         on_init: Optional[Callable[[int, DataFrame], Any]] = None,
         map_func_format_hint: Optional[str] = None,
     ) -> DataFrame:
-        presort = partition_spec.presort
+        is_coarse = partition_spec.algo == "coarse"
+        presort = partition_spec.get_sorts(df.schema, with_partition_keys=is_coarse)
         presort_keys = list(presort.keys())
         presort_asc = list(presort.values())
         output_schema = Schema(output_schema)
@@ -88,10 +92,12 @@ class DaskMapEngine(MapEngine):
             )
         )
 
-        def _map(pdf: Any) -> dd.DataFrame:
+        def _map(pdf: Any) -> pd.DataFrame:
             if pdf.shape[0] == 0:
                 return PandasDataFrame([], output_schema).as_pandas()
-            if len(presort_keys) > 0:
+            if is_coarse:
+                pdf = pdf.drop(columns=[_DASK_PARTITION_KEY])
+            if len(partition_spec.presort) > 0:
                 pdf = pdf.sort_values(presort_keys, ascending=presort_asc)
             input_df = PandasDataFrame(
                 pdf.reset_index(drop=True), input_schema, pandas_df_wrapper=True
@@ -100,7 +106,7 @@ class DaskMapEngine(MapEngine):
                 on_init_once(0, input_df)
             cursor.set(lambda: input_df.peek_array(), 0, 0)
             output_df = map_func(cursor, input_df)
-            return output_df.as_pandas()
+            return output_df.as_pandas()[output_schema.names]
 
         df = self.to_df(df)
         meta = self.execution_engine.pl_utils.safe_to_pandas_dtype(  # type: ignore
@@ -113,8 +119,28 @@ class DaskMapEngine(MapEngine):
             df = self.execution_engine.repartition(
                 df, PartitionSpec(num=partition_spec.num_partitions)
             )
+            if is_coarse:
+                input_num_partitions = df.num_partitions
+                _utils = self.execution_engine.pl_utils  # type: ignore
+                input_meta = _utils.safe_to_pandas_dtype(
+                    (input_schema + (_DASK_PARTITION_KEY, "uint64")).pa_schema
+                )
+                tddf = df.native.map_partitions(
+                    lambda pdf: pdf.assign(
+                        **{
+                            _DASK_PARTITION_KEY: pd.util.hash_pandas_object(
+                                pdf[partition_spec.partition_by], index=False
+                            ).mod(input_num_partitions)
+                        }
+                    ),
+                    meta=input_meta,
+                )
+                keys = [_DASK_PARTITION_KEY]
+            else:
+                tddf = df.native
+                keys = partition_spec.partition_by
             result = self.execution_engine.pl_utils.safe_groupby_apply(  # type: ignore
-                df.native, partition_spec.partition_by, _map, meta=meta  # type: ignore
+                tddf, keys, _map, meta=meta  # type: ignore
             )
         return DaskDataFrame(result, output_schema)
 
@@ -214,7 +240,7 @@ class DaskExecutionEngine(ExecutionEngine):
         p = partition_spec.get_num_partitions(
             **{
                 KEYWORD_ROWCOUNT: lambda: df.persist().count(),  # type: ignore
-                KEYWORD_CORECOUNT: lambda: self.get_current_parallelism(),
+                KEYWORD_PARALLELISM: lambda: self.get_current_parallelism(),
             }
         )
         if p > 0:
@@ -253,7 +279,7 @@ class DaskExecutionEngine(ExecutionEngine):
             join_type=how,
             on=key_schema.names,
         )
-        return DaskDataFrame(d, output_schema)
+        return DaskDataFrame(d, output_schema, type_safe=False)
 
     def union(
         self,
@@ -268,7 +294,7 @@ class DaskExecutionEngine(ExecutionEngine):
         d = self.pl_utils.union(
             self.to_df(df1).native, self.to_df(df2).native, unique=distinct
         )
-        return DaskDataFrame(d, df1.schema)
+        return DaskDataFrame(d, df1.schema, type_safe=False)
 
     def subtract(
         self,
@@ -286,7 +312,7 @@ class DaskExecutionEngine(ExecutionEngine):
         d = self.pl_utils.except_df(
             self.to_df(df1).native, self.to_df(df2).native, unique=distinct
         )
-        return DaskDataFrame(d, df1.schema)
+        return DaskDataFrame(d, df1.schema, type_safe=False)
 
     def intersect(
         self,
@@ -304,11 +330,11 @@ class DaskExecutionEngine(ExecutionEngine):
         d = self.pl_utils.intersect(
             self.to_df(df1).native, self.to_df(df2).native, unique=distinct
         )
-        return DaskDataFrame(d, df1.schema)
+        return DaskDataFrame(d, df1.schema, type_safe=False)
 
     def distinct(self, df: DataFrame) -> DataFrame:
         d = self.pl_utils.drop_duplicates(self.to_df(df).native)
-        return DaskDataFrame(d, df.schema)
+        return DaskDataFrame(d, df.schema, type_safe=False)
 
     def dropna(
         self,
@@ -325,7 +351,7 @@ class DaskExecutionEngine(ExecutionEngine):
         if how == "any" and thresh is not None:
             del kw["how"]  # to deal with a dask logic flaw
         d = self.to_df(df).native.dropna(**kw)
-        return DaskDataFrame(d, df.schema)
+        return DaskDataFrame(d, df.schema, type_safe=False)
 
     def fillna(self, df: DataFrame, value: Any, subset: List[str] = None) -> DataFrame:
         assert_or_throw(
@@ -345,7 +371,7 @@ class DaskExecutionEngine(ExecutionEngine):
             subset = subset or df.columns
             mapping = {col: value for col in subset}
         d = self.to_df(df).native.fillna(mapping)
-        return DaskDataFrame(d, df.schema)
+        return DaskDataFrame(d, df.schema, type_safe=False)
 
     def sample(
         self,
@@ -363,7 +389,7 @@ class DaskExecutionEngine(ExecutionEngine):
         d = self.to_df(df).native.sample(
             n=n, frac=frac, replace=replace, random_state=seed
         )
-        return DaskDataFrame(d, df.schema)
+        return DaskDataFrame(d, df.schema, type_safe=False)
 
     def take(
         self,
@@ -419,7 +445,7 @@ class DaskExecutionEngine(ExecutionEngine):
                 .reset_index(drop=True)
             )
 
-        return DaskDataFrame(d, df.schema)
+        return DaskDataFrame(d, df.schema, type_safe=False)
 
     def load_df(
         self,

fugue_duckdb/_utils.py

@@ -27,7 +27,11 @@ _DUCK_TYPES_TO_PA: Dict[str, pa.DataType] = {
     "TIME": pa.time32("ms"),
 }
 
-_PA_TYPES_TO_DUCK: Dict[pa.DataType, str] = {v: k for k, v in _DUCK_TYPES_TO_PA.items()}
+_PA_TYPES_TO_DUCK: Dict[pa.DataType, str] = {
+    v: k
+    for k, v in list(_DUCK_TYPES_TO_PA.items())
+    + [("VARCHAR", pa.large_string()), ("BLOB", pa.large_binary())]
+}
 
 
 def encode_column_name(name: str) -> str:
@@ -94,8 +98,9 @@ def to_duck_type(tp: pa.DataType) -> str:
         raise ValueError(f"can't convert {tp} to DuckDB data type")
 
 
-def to_pa_type(duck_type: str) -> pa.DataType:
+def to_pa_type(duck_type_raw: Any) -> pa.DataType:
     try:
+        duck_type = str(duck_type_raw)  # for duckdb >= 0.8.0
         if duck_type.endswith("[]"):
             return pa.list_(to_pa_type(duck_type[:-2]))
         p = duck_type.find("(")

fugue_duckdb/dask.py

@@ -50,7 +50,7 @@ class DuckDaskExecutionEngine(DuckExecutionEngine):
                 res = DuckDataFrame(self.connection.from_df(ddf.as_pandas()))
             else:
                 res = DuckDataFrame(
-                    duckdb.arrow(ddf.as_arrow(), connection=self.connection)
+                    duckdb.from_arrow(ddf.as_arrow(), connection=self.connection)
                 )
             if ddf.has_metadata:  # pragma: no cover
                 res.reset_metadata(ddf.metadata)

fugue_duckdb/dataframe.py

@@ -4,14 +4,17 @@ import pandas as pd
 import pyarrow as pa
 from duckdb import DuckDBPyRelation
 from triad import Schema
+from triad.utils.pyarrow import LARGE_TYPES_REPLACEMENT, replace_types_in_table
 
 from fugue import ArrayDataFrame, ArrowDataFrame, DataFrame, LocalBoundedDataFrame
 from fugue.exceptions import FugueDataFrameOperationError, FugueDatasetEmptyError
 from fugue.plugins import (
+    as_arrow,
     as_fugue_dataset,
     as_local_bounded,
     get_column_names,
     get_num_partitions,
+    get_schema,
     is_df,
 )
 
@@ -26,15 +29,7 @@ class DuckDataFrame(LocalBoundedDataFrame):
 
     def __init__(self, rel: DuckDBPyRelation):
         self._rel = rel
-        super().__init__(schema=self._get_schema)
-
-    def _get_schema(self) -> Schema:
-        return Schema(
-            [
-                pa.field(x, to_pa_type(y))
-                for x, y in zip(self._rel.columns, self._rel.types)
-            ]
-        )
+        super().__init__(schema=lambda: _duck_get_schema(self._rel))
 
     @property
     def alias(self) -> str:
@@ -98,7 +93,7 @@ class DuckDataFrame(LocalBoundedDataFrame):
         return DuckDataFrame(self._rel.project(", ".join(fields)))
 
     def as_arrow(self, type_safe: bool = False) -> pa.Table:
-        return self._rel.arrow()
+        return _duck_as_arrow(self._rel)
 
     def as_pandas(self) -> pd.DataFrame:
         if any(pa.types.is_nested(f.type) for f in self.schema.fields):
@@ -169,6 +164,18 @@ def _duck_as_local(df: DuckDBPyRelation) -> DuckDBPyRelation:
     return df
 
 
+@as_arrow.candidate(lambda df: isinstance(df, DuckDBPyRelation))
+def _duck_as_arrow(df: DuckDBPyRelation) -> pa.Table:
+    _df = df.arrow()
+    _df = replace_types_in_table(_df, LARGE_TYPES_REPLACEMENT, recursive=True)
+    return _df
+
+
+@get_schema.candidate(lambda df: isinstance(df, DuckDBPyRelation))
+def _duck_get_schema(df: DuckDBPyRelation) -> Schema:
+    return Schema([pa.field(x, to_pa_type(y)) for x, y in zip(df.columns, df.types)])
+
+
 @get_column_names.candidate(lambda df: isinstance(df, DuckDBPyRelation))
 def _get_duckdb_columns(df: DuckDBPyRelation) -> List[Any]:
     return list(df.columns)

fugue_duckdb/execution_engine.py

@@ -2,12 +2,11 @@ import logging
 from typing import Any, Dict, Iterable, List, Optional, Union
 
 import duckdb
-import pyarrow as pa
 from duckdb import DuckDBPyConnection, DuckDBPyRelation
 from triad import SerializableRLock
 from triad.collections.fs import FileSystem
-from triad.utils.schema import quote_name
 from triad.utils.assertion import assert_or_throw
+from triad.utils.schema import quote_name
 
 from fugue import (
     ArrowDataFrame,
@@ -19,12 +18,7 @@ from fugue import (
 )
 from fugue.collections.partition import PartitionSpec, parse_presort_exp
 from fugue.collections.sql import StructuredRawSQL, TempTableName
-from fugue.dataframe import (
-    DataFrame,
-    DataFrames,
-    LocalBoundedDataFrame,
-    PandasDataFrame,
-)
+from fugue.dataframe import DataFrame, DataFrames, LocalBoundedDataFrame
 from fugue.dataframe.utils import get_join_schemas
 
 from ._io import DuckDBIO
@@ -34,7 +28,7 @@ from ._utils import (
     encode_schema_names,
     encode_value_to_expr,
 )
-from .dataframe import DuckDataFrame
+from .dataframe import DuckDataFrame, _duck_as_arrow
 
 _FUGUE_DUCKDB_PRAGMA_CONFIG_PREFIX = "fugue.duckdb.pragma."
 _FUGUE_DUCKDB_EXTENSIONS = "fugue.duckdb.extensions"
@@ -114,8 +108,8 @@ class DuckDBEngine(SQLEngine):
         conn = duckdb.connect()
         try:
             for k, v in dfs.items():
-                duckdb.arrow(v.as_arrow(), connection=conn).create_view(k)
-            return ArrowDataFrame(conn.execute(statement).arrow())
+                duckdb.from_arrow(v.as_arrow(), connection=conn).create_view(k)
+            return ArrowDataFrame(_duck_as_arrow(conn.execute(statement)))
         finally:
             conn.close()
 
@@ -235,7 +229,7 @@ class DuckExecutionEngine(ExecutionEngine):
         # TODO: we should create DuckDB table, but it has bugs, so can't use by 0.3.1
         if isinstance(df, DuckDataFrame):
             # materialize
-            res: DataFrame = ArrowDataFrame(df.native.arrow())
+            res: DataFrame = ArrowDataFrame(df.as_arrow())
         else:
             res = self.to_df(df)
         res.reset_metadata(df.metadata)
@@ -545,19 +539,15 @@ def _to_duck_df(
             )
             if isinstance(df, DuckDataFrame):
                 return df
-
-            if isinstance(df, PandasDataFrame) and all(
-                not pa.types.is_nested(f.type) for f in df.schema.fields
-            ):
-                rdf = DuckDataFrame(engine.connection.from_df(df.as_pandas()))
-            else:
-                rdf = DuckDataFrame(
-                    duckdb.arrow(df.as_arrow(), connection=engine.connection)
-                )
+            rdf = DuckDataFrame(
+                duckdb.from_arrow(df.as_arrow(), connection=engine.connection)
+            )
             rdf.reset_metadata(df.metadata if df.has_metadata else None)
             return rdf
         tdf = ArrowDataFrame(df, schema)
-        return DuckDataFrame(duckdb.arrow(tdf.native, connection=engine.connection))
+        return DuckDataFrame(
+            duckdb.from_arrow(tdf.native, connection=engine.connection)
+        )
 
     res = _gen_duck()
     if create_view:

fugue_ibis/dataframe.py

@@ -5,12 +5,7 @@ import pandas as pd
 import pyarrow as pa
 from triad import Schema, assert_or_throw
 
-from fugue import (
-    DataFrame,
-    IterableDataFrame,
-    LocalBoundedDataFrame,
-    to_local_bounded_df,
-)
+from fugue import DataFrame, IterableDataFrame, LocalBoundedDataFrame
 from fugue.dataframe.dataframe import _input_schema
 from fugue.exceptions import FugueDataFrameOperationError, FugueDatasetEmptyError
 from fugue.plugins import drop_columns, get_column_names, is_df, rename
@@ -153,7 +148,7 @@ class IbisDataFrame(DataFrame):
     ) -> LocalBoundedDataFrame:
         if columns is not None:
             return self[columns].head(n)
-        return to_local_bounded_df(self._to_local_df(self._table.head(n)))
+        return self._to_local_df(self._table.head(n)).as_local_bounded()
 
     def _alter_table_columns(self, table: IbisTable, new_schema: Schema) -> IbisTable:
         fields: Dict[str, Any] = {}

fugue_notebook/env.py

@@ -3,21 +3,16 @@ import html
 import json
 from typing import Any, Dict, List, Optional
 
-from IPython.core.magic import Magics, cell_magic, magics_class, needs_local_scope
 from IPython import get_ipython
+from IPython.core.magic import Magics, cell_magic, magics_class, needs_local_scope
 from IPython.display import HTML, display
 from triad import ParamDict
 from triad.utils.convert import to_instance
 from triad.utils.pyarrow import _field_to_expression
 
-import fugue_sql
-from fugue import (
-    DataFrame,
-    DataFrameDisplay,
-    ExecutionEngine,
-    get_dataset_display,
-    make_execution_engine,
-)
+from fugue import DataFrame, DataFrameDisplay, ExecutionEngine
+from fugue import fsql as fugue_sql
+from fugue import get_dataset_display, make_execution_engine
 from fugue.dataframe import YieldedDataFrame
 from fugue.exceptions import FugueSQLSyntaxError
 
@@ -58,7 +53,7 @@ class _FugueSQLMagics(Magics):
     @cell_magic("fsql")
     def fsql(self, line: str, cell: str, local_ns: Any = None) -> None:
         try:
-            dag = fugue_sql.fsql(
+            dag = fugue_sql(
                 "\n" + cell, local_ns, fsql_ignore_case=self._fsql_ignore_case
             )
         except FugueSQLSyntaxError as ex:

fugue_polars/_utils.py

@@ -1,48 +1,8 @@
 import polars as pl
-import pyarrow as pa
 from triad import Schema
-from triad.utils.pyarrow import get_alter_func
 
 from fugue.dataframe.arrow_dataframe import _build_empty_arrow
 
 
-def pl_as_arrow(df: pl.DataFrame) -> pa.Table:
-    adf = df.to_arrow()
-    schema = convert_schema(adf.schema)
-    func = get_alter_func(adf.schema, schema, safe=False)
-    return func(adf)
-
-
-def to_schema(df: pl.DataFrame) -> Schema:
-    return Schema(convert_schema(pl.DataFrame(schema=df.schema).to_arrow().schema))
-
-
 def build_empty_pl(schema: Schema) -> pl.DataFrame:
     return pl.from_arrow(_build_empty_arrow(schema))
-
-
-def convert_schema(schema: pa.Schema) -> pa.Schema:
-    fields = [convert_field(f) for f in schema]
-    return pa.schema(fields)
-
-
-def convert_field(field: pa.Field) -> pa.Field:
-    tp = convert_type(field.type)
-    if tp == field.type:
-        return field
-    print(field.type, tp)
-    return pa.field(field.name, tp)
-
-
-def convert_type(tp: pa.DataType) -> pa.DataType:
-    if pa.types.is_struct(tp):
-        return pa.struct([convert_field(f) for f in tp])
-    if pa.types.is_list(tp) or pa.types.is_large_list(tp):
-        return pa.list_(convert_type(tp.value_type))
-    if pa.types.is_map(tp):  # pragma: no cover
-        return pa.map_(convert_type(tp.key_type), convert_type(tp.value_type))
-    if pa.types.is_large_string(tp):
-        return pa.string()
-    if pa.types.is_large_binary(tp):
-        return pa.binary()
-    return tp

fugue_polars/polars_dataframe.py

@@ -6,9 +6,15 @@ 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,
@@ -28,7 +34,7 @@ from fugue.dataset.api import (
 )
 from fugue.exceptions import FugueDataFrameOperationError
 
-from ._utils import build_empty_pl, pl_as_arrow, to_schema
+from ._utils import build_empty_pl
 
 
 class PolarsDataFrame(LocalBoundedDataFrame):
@@ -55,7 +61,7 @@ class PolarsDataFrame(LocalBoundedDataFrame):
                 InvalidOperationError("can't reset schema for pl.DataFrame"),
             )
             self._native = df
-            super().__init__(to_schema(df))
+            super().__init__(_get_pl_schema(df))
 
     @property
     def native(self) -> pl.DataFrame:
@@ -75,7 +81,7 @@ class PolarsDataFrame(LocalBoundedDataFrame):
 
     def peek_dict(self) -> Dict[str, Any]:
         self.assert_not_empty()
-        return dict(zip(self._native.columns, self._native.row(0)))
+        return self._native.row(0, named=True)
 
     def count(self) -> int:
         return self.native.shape[0]
@@ -107,7 +113,7 @@ class PolarsDataFrame(LocalBoundedDataFrame):
         return PolarsDataFrame(pl.from_arrow(adf.native))
 
     def as_arrow(self, type_safe: bool = False) -> pa.Table:
-        return pl_as_arrow(self.native)
+        return _pl_as_arrow(self.native)
 
     def as_array(
         self, columns: Optional[List[str]] = None, type_safe: bool = False
@@ -121,7 +127,7 @@ class PolarsDataFrame(LocalBoundedDataFrame):
         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(
+            yield from ArrowDataFrame(_pl_as_arrow(self.native)).as_array_iterable(
                 columns=columns
             )
 
@@ -129,7 +135,7 @@ class PolarsDataFrame(LocalBoundedDataFrame):
         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(
+            yield from ArrowDataFrame(_pl_as_arrow(self.native)).as_dict_iterable(
                 columns=columns
             )
 
@@ -144,6 +150,13 @@ 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
@@ -181,7 +194,9 @@ def _get_pl_columns(df: pl.DataFrame) -> List[Any]:
 
 @get_schema.candidate(lambda df: isinstance(df, pl.DataFrame))
 def _get_pl_schema(df: pl.DataFrame) -> Schema:
-    return to_schema(df)
+    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))

fugue_ray/_constants.py

@@ -1,4 +1,6 @@
-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"
@@ -10,3 +12,8 @@ FUGUE_RAY_DEFAULT_CONF: Dict[str, Any] = {
     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,11 +1,15 @@
 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__"
 
 
@@ -15,6 +19,8 @@ def get_dataset_format(df: rd.Dataset) -> Optional[str]:
         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
 
 
@@ -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)