fugue 0.8.2.dev1__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."""
@@ -72,12 +75,15 @@ class DaskMapEngine(MapEngine):
         output_schema: Any,
         partition_spec: PartitionSpec,
         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)
         input_schema = df.schema
+        cursor = partition_spec.get_cursor(input_schema, 0)
         on_init_once: Any = (
             None
             if on_init is None
@@ -86,20 +92,21 @@ 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
             )
             if on_init_once is not None:
                 on_init_once(0, input_df)
-            cursor = partition_spec.get_cursor(input_schema, 0)
-            cursor.set(input_df.peek_array(), 0, 0)
+            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
@@ -112,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)
 
@@ -213,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:
@@ -252,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,
@@ -267,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,
@@ -285,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,
@@ -303,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,
@@ -324,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(
@@ -344,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,
@@ -362,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,
@@ -418,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_dask/registry.py

@@ -1,16 +1,15 @@
-import inspect
-from typing import Any, Optional
+from typing import Any
 
 import dask.dataframe as dd
 from dask.distributed import Client
 from triad import run_at_def
 
-from fugue import DataFrame, is_pandas_or, register_execution_engine
-from fugue._utils.interfaceless import (
+from fugue import DataFrame, register_execution_engine
+from fugue.dev import (
     DataFrameParam,
     ExecutionEngineParam,
-    SimpleAnnotationConverter,
-    register_annotation_converter,
+    fugue_annotated_param,
+    is_pandas_or,
 )
 from fugue.plugins import as_fugue_dataset, infer_execution_engine
 from fugue_dask._utils import DASK_UTILS
@@ -45,36 +44,13 @@ def _register_engines() -> None:
     )
 
 
-def _register_annotation_converters() -> None:
-    register_annotation_converter(
-        0.8,
-        SimpleAnnotationConverter(
-            DaskExecutionEngine,
-            lambda param: _DaskExecutionEngineParam(param),
-        ),
-    )
-    register_annotation_converter(
-        0.8,
-        SimpleAnnotationConverter(
-            dd.DataFrame, lambda param: _DaskDataFrameParam(param)
-        ),
-    )
-
-
+@fugue_annotated_param(DaskExecutionEngine)
 class _DaskExecutionEngineParam(ExecutionEngineParam):
-    def __init__(
-        self,
-        param: Optional[inspect.Parameter],
-    ):
-        super().__init__(
-            param, annotation="DaskExecutionEngine", engine_type=DaskExecutionEngine
-        )
+    pass
 
 
+@fugue_annotated_param(dd.DataFrame)
 class _DaskDataFrameParam(DataFrameParam):
-    def __init__(self, param: Optional[inspect.Parameter]):
-        super().__init__(param, annotation="dask.dataframe.DataFrame")
-
     def to_input_data(self, df: DataFrame, ctx: Any) -> Any:
         assert isinstance(ctx, DaskExecutionEngine)
         return ctx.to_df(df).native
@@ -99,4 +75,3 @@ def _register() -> None:
         >>> import fugue_dask
     """
     _register_engines()
-    _register_annotation_converters()

fugue_duckdb/_io.py

@@ -21,8 +21,14 @@ from fugue_duckdb.dataframe import DuckDataFrame
 def _get_single_files(
     fp: Iterable[FileParser], fs: FileSystem, fmt: str
 ) -> Iterable[FileParser]:
+    def _isdir(d: str) -> bool:
+        try:
+            return fs.isdir(d)
+        except Exception:  # pragma: no cover
+            return False
+
     for f in fp:
-        if f.glob_pattern == "" and fs.isdir(f.uri):
+        if f.glob_pattern == "" and _isdir(f.uri):
             yield f.with_glob("*." + fmt, fmt)
         else:
             yield f
@@ -211,7 +217,7 @@ class DuckDBIO:
         # for k, v in kw.items():
         #    params.append(f"{k}=" + encode_value_to_expr(v))
         pm = ", ".join(params)
-        query = f"SELECT {cols} FROM parquet_scan({pm})"
+        query = f"SELECT {cols} FROM parquet_scan([{pm}])"
         res = DuckDataFrame(self._con.from_query(query))
         return (
             res  # type: ignore

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,20 +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,
-    LocalDataFrame,
-)
+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,
 )
 
@@ -32,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:
@@ -104,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):
@@ -112,8 +101,11 @@ class DuckDataFrame(LocalBoundedDataFrame):
             return ArrowDataFrame(self.as_arrow()).as_pandas()
         return self._rel.to_df()
 
-    def as_local(self) -> LocalDataFrame:
-        return ArrowDataFrame(self.as_arrow())
+    def as_local_bounded(self) -> LocalBoundedDataFrame:
+        res = ArrowDataFrame(self.as_arrow())
+        if self.has_metadata:
+            res.reset_metadata(self.metadata)
+        return res
 
     def as_array(
         self, columns: Optional[List[str]] = None, type_safe: bool = False
@@ -172,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,9 +28,10 @@ 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"
 
 
 class DuckDBEngine(SQLEngine):
@@ -113,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()
 
@@ -161,6 +156,12 @@ class DuckExecutionEngine(ExecutionEngine):
         try:
             for pg in list(self._get_pragmas()):  # transactional
                 self._con.execute(pg)
+
+            for ext in self.conf.get(_FUGUE_DUCKDB_EXTENSIONS, "").split(","):
+                _ext = ext.strip()
+                if _ext != "":
+                    self._con.install_extension(_ext)
+                    self._con.load_extension(_ext)
         except Exception:
             self.stop()
             raise
@@ -228,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)
@@ -538,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_duckdb/registry.py

@@ -1,5 +1,4 @@
-import inspect
-from typing import Any, Optional
+from typing import Any
 
 from duckdb import DuckDBPyConnection, DuckDBPyRelation
 from triad import run_at_def
@@ -7,15 +6,14 @@ from triad import run_at_def
 from fugue import (
     DataFrame,
     ExecutionEngine,
-    is_pandas_or,
     register_execution_engine,
     register_sql_engine,
 )
-from fugue._utils.interfaceless import (
+from fugue.dev import (
     DataFrameParam,
     ExecutionEngineParam,
-    SimpleAnnotationConverter,
-    register_annotation_converter,
+    fugue_annotated_param,
+    is_pandas_or,
 )
 from fugue.plugins import infer_execution_engine
 from fugue_duckdb.dataframe import DuckDataFrame
@@ -69,40 +67,20 @@ def _register_engines() -> None:
     register_sql_engine("duckdb", lambda engine: DuckDBEngine(engine))
 
 
-def _register_annotation_converters() -> None:
-    register_annotation_converter(
-        0.8,
-        SimpleAnnotationConverter(
-            DuckDBPyConnection,
-            lambda param: _DuckDBPyConnectionParam(param),
-        ),
-    )
-    register_annotation_converter(
-        0.8,
-        SimpleAnnotationConverter(
-            DuckDBPyRelation,
-            lambda param: _DuckDBPyRelationParam(param),
-        ),
-    )
+@fugue_annotated_param(DuckExecutionEngine)
+class _DuckExecutionEngineParam(ExecutionEngineParam):
+    pass
 
 
+@fugue_annotated_param(DuckDBPyConnection)
 class _DuckDBPyConnectionParam(ExecutionEngineParam):
-    def __init__(
-        self,
-        param: Optional[inspect.Parameter],
-    ):
-        super().__init__(
-            param, annotation="DuckDBPyConnection", engine_type=DuckExecutionEngine
-        )
-
     def to_input(self, engine: ExecutionEngine) -> Any:
-        return super().to_input(engine).connection  # type:ignore
+        assert isinstance(engine, DuckExecutionEngine)
+        return engine.connection  # type:ignore
 
 
+@fugue_annotated_param(DuckDBPyRelation)
 class _DuckDBPyRelationParam(DataFrameParam):
-    def __init__(self, param: Optional[inspect.Parameter]):
-        super().__init__(param, annotation="DuckDBPyRelation")
-
     def to_input_data(self, df: DataFrame, ctx: Any) -> Any:
         assert isinstance(ctx, DuckExecutionEngine)
         return ctx.to_df(df).native  # type: ignore
@@ -127,4 +105,3 @@ def _register() -> None:
         >>> import fugue_duckdb
     """
     _register_engines()
-    _register_annotation_converters()

fugue_ibis/dataframe.py

@@ -5,13 +5,7 @@ import pandas as pd
 import pyarrow as pa
 from triad import Schema, assert_or_throw
 
-from fugue import (
-    DataFrame,
-    IterableDataFrame,
-    LocalBoundedDataFrame,
-    LocalDataFrame,
-    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
@@ -50,7 +44,9 @@ class IbisDataFrame(DataFrame):
     def _to_schema(self, schema: IbisSchema) -> Schema:
         return to_schema(schema)
 
-    def _to_local_df(self, table: IbisTable, schema: Any = None) -> LocalDataFrame:
+    def _to_local_df(
+        self, table: IbisTable, schema: Any = None
+    ) -> LocalBoundedDataFrame:
         raise NotImplementedError  # pragma: no cover
 
     def _to_iterable_df(
@@ -124,7 +120,7 @@ class IbisDataFrame(DataFrame):
     def as_pandas(self) -> pd.DataFrame:
         return self.as_local().as_pandas()
 
-    def as_local(self) -> LocalDataFrame:
+    def as_local_bounded(self) -> LocalBoundedDataFrame:
         res = self._to_local_df(self._table, schema=self.schema)
         if res is not self and self.has_metadata:
             res.reset_metadata(self.metadata)
@@ -152,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_ibis/execution_engine.py

@@ -324,10 +324,16 @@ class IbisMapEngine(MapEngine):
         output_schema: Any,
         partition_spec: PartitionSpec,
         on_init: Optional[Callable[[int, DataFrame], Any]] = None,
+        map_func_format_hint: Optional[str] = None,
     ) -> DataFrame:
         _df = self._ibis_engine._to_non_ibis_dataframe(df)
         return self._ibis_engine.non_ibis_engine.map_engine.map_dataframe(
-            _df, map_func, output_schema, partition_spec, on_init
+            _df,
+            map_func=map_func,
+            output_schema=output_schema,
+            partition_spec=partition_spec,
+            on_init=on_init,
+            map_func_format_hint=map_func_format_hint,
         )
 
     def map_bag(

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/__init__.py

@@ -0,0 +1,2 @@
+# flake8: noqa
+from .polars_dataframe import PolarsDataFrame

fugue_polars/_utils.py

@@ -0,0 +1,8 @@
+import polars as pl
+from triad import Schema
+
+from fugue.dataframe.arrow_dataframe import _build_empty_arrow
+
+
+def build_empty_pl(schema: Schema) -> pl.DataFrame:
+    return pl.from_arrow(_build_empty_arrow(schema))