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

fugue_ray/execution_engine.py

@@ -2,7 +2,7 @@ from typing import Any, Callable, Dict, List, Optional, Type, Union
 
 import pyarrow as pa
 import ray
-from duckdb import DuckDBPyConnection, DuckDBPyRelation
+from duckdb import DuckDBPyConnection
 from triad import Schema, assert_or_throw, to_uuid
 from triad.utils.threading import RunOnce
 
@@ -15,14 +15,14 @@ from fugue import (
     PartitionCursor,
     PartitionSpec,
 )
-from fugue.constants import KEYWORD_ROWCOUNT
+from fugue.constants import KEYWORD_PARALLELISM, KEYWORD_ROWCOUNT
 from fugue.dataframe.arrow_dataframe import _build_empty_arrow
 from fugue_duckdb.dataframe import DuckDataFrame
 from fugue_duckdb.execution_engine import DuckExecutionEngine
 
-from ._constants import FUGUE_RAY_DEFAULT_BATCH_SIZE
+from ._constants import FUGUE_RAY_DEFAULT_BATCH_SIZE, FUGUE_RAY_ZERO_COPY
 from ._utils.cluster import get_default_partitions, get_default_shuffle_partitions
-from ._utils.dataframe import add_partition_key
+from ._utils.dataframe import add_coarse_partition_key, add_partition_key
 from ._utils.io import RayIO
 from .dataframe import RayDataFrame
 
@@ -45,6 +45,7 @@ class RayMapEngine(MapEngine):
         output_schema: Any,
         partition_spec: PartitionSpec,
         on_init: Optional[Callable[[int, DataFrame], Any]] = None,
+        map_func_format_hint: Optional[str] = None,
     ) -> DataFrame:
         if len(partition_spec.partition_by) == 0:
             return self._map(
@@ -71,12 +72,15 @@ class RayMapEngine(MapEngine):
         partition_spec: PartitionSpec,
         on_init: Optional[Callable[[int, DataFrame], Any]] = None,
     ) -> DataFrame:
-        presort = partition_spec.presort
+        output_schema = Schema(output_schema)
+        input_schema = df.schema
+        presort = partition_spec.get_sorts(
+            input_schema, with_partition_keys=partition_spec.algo == "coarse"
+        )
         presort_tuples = [
             (k, "ascending" if v else "descending") for k, v in presort.items()
         ]
-        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
@@ -89,7 +93,7 @@ class RayMapEngine(MapEngine):
             if adf.shape[0] == 0:
                 return _build_empty_arrow(output_schema)
             adf = adf.remove_column(len(input_schema))  # remove partition key
-            if len(presort_tuples) > 0:
+            if len(partition_spec.presort) > 0:
                 if pa.__version__ < "7":  # pragma: no cover
                     idx = pa.compute.sort_indices(
                         adf, options=pa.compute.SortOptions(presort_tuples)
@@ -100,8 +104,7 @@ class RayMapEngine(MapEngine):
             input_df = ArrowDataFrame(adf)
             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_arrow()
 
@@ -117,12 +120,20 @@ class RayMapEngine(MapEngine):
                 _df = self.execution_engine.repartition(  # type: ignore
                     _df, PartitionSpec(num=n)
                 )
-        rdf, _ = add_partition_key(
-            _df.native,
-            keys=partition_spec.partition_by,
-            input_schema=input_schema,
-            output_key=_RAY_PARTITION_KEY,
-        )
+        if partition_spec.algo != "coarse":
+            rdf, _ = add_partition_key(
+                _df.native,
+                keys=partition_spec.partition_by,
+                input_schema=input_schema,
+                output_key=_RAY_PARTITION_KEY,
+            )
+        else:
+            rdf = add_coarse_partition_key(
+                _df.native,
+                keys=partition_spec.partition_by,
+                output_key=_RAY_PARTITION_KEY,
+                bucket=_df.num_partitions,
+            )
 
         gdf = rdf.groupby(_RAY_PARTITION_KEY)
         sdf = gdf.map_groups(
@@ -142,6 +153,7 @@ class RayMapEngine(MapEngine):
     ) -> DataFrame:
         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
@@ -156,8 +168,7 @@ class RayMapEngine(MapEngine):
             input_df = ArrowDataFrame(adf)
             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_arrow()
 
@@ -175,15 +186,17 @@ class RayMapEngine(MapEngine):
                 rdf = self.execution_engine.repartition(  # type: ignore
                     rdf, PartitionSpec(num=n)
                 )
-        batch_size = (
-            self.conf.get_or_throw(FUGUE_RAY_DEFAULT_BATCH_SIZE, object)
-            if FUGUE_RAY_DEFAULT_BATCH_SIZE in self.execution_engine.conf
-            else "default"
-        )
+        mb_args: Dict[str, Any] = {}
+        if FUGUE_RAY_DEFAULT_BATCH_SIZE in self.conf:
+            mb_args["batch_size"] = self.conf.get_or_throw(
+                FUGUE_RAY_DEFAULT_BATCH_SIZE, int
+            )
+        if ray.__version__ >= "2.3":
+            mb_args["zero_copy_batch"] = self.conf.get(FUGUE_RAY_ZERO_COPY, True)
         sdf = rdf.native.map_batches(
             _udf,
             batch_format="pyarrow",
-            batch_size=batch_size,
+            **mb_args,
             **self.execution_engine._get_remote_args(),  # type: ignore
         )
         return RayDataFrame(sdf, schema=output_schema, internal_schema=True)
@@ -202,6 +215,7 @@ class RayExecutionEngine(DuckExecutionEngine):
     ):
         if not ray.is_initialized():  # pragma: no cover
             ray.init()
+
         super().__init__(conf, connection)
         self._io = RayIO(self)
 
@@ -232,12 +246,15 @@ class RayExecutionEngine(DuckExecutionEngine):
 
         rdf = self._to_ray_df(df)
 
-        num_funcs = {KEYWORD_ROWCOUNT: lambda: _persist_and_count(rdf)}
+        num_funcs = {
+            KEYWORD_ROWCOUNT: lambda: _persist_and_count(rdf),
+            KEYWORD_PARALLELISM: lambda: self.get_current_parallelism(),
+        }
         num = partition_spec.get_num_partitions(**num_funcs)
         pdf = rdf.native
 
         if num > 0:
-            if partition_spec.algo in ["hash", "even"]:
+            if partition_spec.algo in ["hash", "even", "coarse"]:
                 pdf = pdf.repartition(num)
             elif partition_spec.algo == "rand":
                 pdf = pdf.repartition(num, shuffle=True)
@@ -264,6 +281,16 @@ class RayExecutionEngine(DuckExecutionEngine):
             return df if not as_local else df.as_local()
         return super().convert_yield_dataframe(df, as_local)
 
+    def union(self, df1: DataFrame, df2: DataFrame, distinct: bool = True) -> DataFrame:
+        if distinct:
+            return super().union(df1, df2, distinct)
+        assert_or_throw(
+            df1.schema == df2.schema, ValueError(f"{df1.schema} != {df2.schema}")
+        )
+        tdf1 = self._to_ray_df(df1)
+        tdf2 = self._to_ray_df(df2)
+        return RayDataFrame(tdf1.native.union(tdf2.native), df1.schema)
+
     def load_df(  # type:ignore
         self,
         path: Union[str, List[str]],
@@ -312,12 +339,6 @@ class RayExecutionEngine(DuckExecutionEngine):
                 ValueError("schema must be None when df is a DataFrame"),
             )
             return df
-        if isinstance(df, DuckDBPyRelation):
-            assert_or_throw(
-                schema is None,
-                ValueError("schema must be None when df is a DuckDBPyRelation"),
-            )
-            return DuckDataFrame(df)
         return RayDataFrame(df, schema)
 
     def _get_remote_args(self) -> Dict[str, Any]:

fugue_ray/registry.py

@@ -1,19 +1,17 @@
-import inspect
-from typing import Any, Optional
+from typing import Any
 
 import ray.data as rd
 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 .dataframe import RayDataFrame
 from .execution_engine import RayExecutionEngine
 
@@ -36,34 +34,13 @@ def _register_engines() -> None:
     )
 
 
-def _register_annotation_converters() -> None:
-    register_annotation_converter(
-        0.8,
-        SimpleAnnotationConverter(
-            RayExecutionEngine,
-            lambda param: _RayExecutionEngineParam(param),
-        ),
-    )
-    register_annotation_converter(
-        0.8,
-        SimpleAnnotationConverter(rd.Dataset, lambda param: _RayDatasetParam(param)),
-    )
-
-
+@fugue_annotated_param(RayExecutionEngine)
 class _RayExecutionEngineParam(ExecutionEngineParam):
-    def __init__(
-        self,
-        param: Optional[inspect.Parameter],
-    ):
-        super().__init__(
-            param, annotation="RayExecutionEngine", engine_type=RayExecutionEngine
-        )
+    pass
 
 
+@fugue_annotated_param(rd.Dataset)
 class _RayDatasetParam(DataFrameParam):
-    def __init__(self, param: Optional[inspect.Parameter]):
-        super().__init__(param, annotation="ray.data.Dataset")
-
     def to_input_data(self, df: DataFrame, ctx: Any) -> Any:
         assert isinstance(ctx, RayExecutionEngine)
         return ctx._to_ray_df(df).native
@@ -81,4 +58,3 @@ class _RayDatasetParam(DataFrameParam):
 def _register() -> None:
     """Register Ray Execution Engine"""
     _register_engines()
-    _register_annotation_converters()

fugue_spark/_utils/convert.py

@@ -1,29 +1,24 @@
 from typing import Any, Iterable, List, Tuple
 
+import cloudpickle
+import pandas as pd
 import pyarrow as pa
 import pyspark.sql as ps
 import pyspark.sql.types as pt
-
-try:  # pyspark < 3
-    from pyspark.sql.types import from_arrow_type, to_arrow_type  # type: ignore
-
-    # https://issues.apache.org/jira/browse/SPARK-29041
-    pt._acceptable_types[pt.BinaryType] = (bytearray, bytes)  # type: ignore  # pragma: no cover  # noqa: E501  # pylint: disable=line-too-long
-except ImportError:  # pyspark >=3
-    from pyspark.sql.pandas.types import from_arrow_type, to_arrow_type
-
 from pyarrow.types import is_list, is_struct, is_timestamp
+from pyspark.sql.pandas.types import from_arrow_type, to_arrow_type
 from triad.collections import Schema
 from triad.utils.assertion import assert_arg_not_none, assert_or_throw
 from triad.utils.pyarrow import TRIAD_DEFAULT_TIMESTAMP
 from triad.utils.schema import quote_name
+from .misc import is_spark_dataframe
 
 
 def to_spark_schema(obj: Any) -> pt.StructType:
     assert_arg_not_none(obj, "schema")
     if isinstance(obj, pt.StructType):
         return obj
-    if isinstance(obj, ps.DataFrame):
+    if is_spark_dataframe(obj):
         return obj.schema
     return _from_arrow_schema(Schema(obj).pa_schema)
 
@@ -32,7 +27,7 @@ def to_schema(obj: Any) -> Schema:
     assert_arg_not_none(obj, "obj")
     if isinstance(obj, pt.StructType):
         return Schema(_to_arrow_schema(obj))
-    if isinstance(obj, ps.DataFrame):
+    if is_spark_dataframe(obj):
         return to_schema(obj.schema)
     return Schema(obj)
 
@@ -113,6 +108,22 @@ def to_type_safe_input(rows: Iterable[ps.Row], schema: Schema) -> Iterable[List[
             yield r
 
 
+def to_pandas(df: ps.DataFrame) -> pd.DataFrame:
+    if pd.__version__ < "2" or not any(
+        isinstance(x.dataType, (pt.TimestampType, pt.TimestampNTZType))
+        for x in df.schema.fields
+    ):
+        return df.toPandas()
+
+    def serialize(dfs):  # pragma: no cover
+        for df in dfs:
+            data = cloudpickle.dumps(df)
+            yield pd.DataFrame([[data]], columns=["data"])
+
+    sdf = df.mapInPandas(serialize, schema="data binary")
+    return pd.concat(cloudpickle.loads(x.data) for x in sdf.collect())
+
+
 # TODO: the following function always set nullable to true,
 # but should we use field.nullable?
 def _to_arrow_type(dt: pt.DataType) -> pa.DataType:

fugue_spark/_utils/io.py

@@ -21,7 +21,6 @@ class SparkIO(object):
             "csv": self._load_csv,
             "parquet": self._load_parquet,
             "json": self._load_json,
-            "avro": self._load_avro,
         }
 
     def load_df(
@@ -136,15 +135,3 @@ class SparkIO(object):
             return SparkDataFrame(reader.load(p))[columns]
         schema = Schema(columns)
         return SparkDataFrame(reader.load(p)[schema.names], schema)
-
-    def _load_avro(self, p: List[str], columns: Any = None, **kwargs: Any) -> DataFrame:
-        reader = self._session.read.format(
-            "avro"
-        )  # avro is an external data source that has built-in support since spark 2.4
-        reader.options(**kwargs)
-        if columns is None:
-            return SparkDataFrame(reader.load(p))
-        if isinstance(columns, list):  # column names
-            return SparkDataFrame(reader.load(p))[columns]
-        schema = Schema(columns)
-        return SparkDataFrame(reader.load(p)[schema.names], schema)

fugue_spark/_utils/misc.py

@@ -0,0 +1,27 @@
+from typing import Any
+
+try:
+    from pyspark.sql.connect.session import SparkSession as SparkConnectSession
+    from pyspark.sql.connect.dataframe import DataFrame as SparkConnectDataFrame
+except ImportError:  # pragma: no cover
+    SparkConnectSession = None
+    SparkConnectDataFrame = None
+import pyspark.sql as ps
+
+
+def is_spark_connect(session: Any) -> bool:
+    return SparkConnectSession is not None and isinstance(
+        session, (SparkConnectSession, SparkConnectDataFrame)
+    )
+
+
+def is_spark_dataframe(df: Any) -> bool:
+    return isinstance(df, ps.DataFrame) or (
+        SparkConnectDataFrame is not None and isinstance(df, SparkConnectDataFrame)
+    )
+
+
+def is_spark_session(session: Any) -> bool:
+    return isinstance(session, ps.SparkSession) or (
+        SparkConnectSession is not None and isinstance(session, SparkConnectSession)
+    )

fugue_spark/_utils/partition.py

@@ -1,11 +1,12 @@
-import random
 from typing import Any, Iterable, List
 
 import pyspark.sql as ps
-from fugue_spark._utils.convert import to_schema, to_spark_schema
+import pyspark.sql.functions as psf
 from pyspark import RDD
 from pyspark.sql import SparkSession
-from pyspark.sql.functions import lit
+import warnings
+from .convert import to_schema, to_spark_schema
+from .misc import is_spark_connect
 
 _PARTITION_DUMMY_KEY = "__partition_dummy_key__"
 
@@ -28,16 +29,10 @@ def rand_repartition(
     if len(cols) > 0 or num <= 1:
         return hash_repartition(session, df, num, cols)
 
-    def _rand(rows: Iterable[Any], n: int) -> Iterable[Any]:  # pragma: no cover
-        for row in rows:
-            yield (random.randint(0, n - 1), row)
-
-    rdd = (
-        df.rdd.mapPartitions(lambda r: _rand(r, num))
-        .partitionBy(num, lambda k: k)
-        .mapPartitions(_to_rows)
+    tdf = df.withColumn(
+        _PARTITION_DUMMY_KEY, (psf.rand(0) * psf.lit(2**15 - 1)).cast("long")
     )
-    return session.createDataFrame(rdd, df.schema)
+    return tdf.repartition(num, _PARTITION_DUMMY_KEY)[df.schema.names]
 
 
 def even_repartition(
@@ -45,6 +40,9 @@ def even_repartition(
 ) -> ps.DataFrame:
     if num == 1:
         return _single_repartition(df)
+    if is_spark_connect(session):  # pragma: no cover
+        warnings.warn("Even repartitioning is not supported by Spark Connect")
+        return hash_repartition(session, df, num, cols)
     if len(cols) == 0:
         if num == 0:
             return df
@@ -82,7 +80,7 @@ def even_repartition(
 
 def _single_repartition(df: ps.DataFrame) -> ps.DataFrame:
     return (
-        df.withColumn(_PARTITION_DUMMY_KEY, lit(0))
+        df.withColumn(_PARTITION_DUMMY_KEY, psf.lit(0))
         .repartition(_PARTITION_DUMMY_KEY)
         .drop(_PARTITION_DUMMY_KEY)
     )
@@ -93,11 +91,6 @@ def _to_rows(rdd: Iterable[Any]) -> Iterable[Any]:  # pragma: no cover
         yield item[1]
 
 
-def _to_rows_with_key(rdd: Iterable[Any]) -> Iterable[Any]:  # pragma: no cover
-    for item in rdd:
-        yield list(item[1]) + [item[0]]
-
-
 def _zipWithIndex(rdd: RDD, to_rows: bool = False) -> RDD:
     """
     Modified from

fugue_spark/dataframe.py

@@ -13,7 +13,6 @@ from fugue.dataframe import (
     DataFrame,
     IterableDataFrame,
     LocalBoundedDataFrame,
-    LocalDataFrame,
     PandasDataFrame,
 )
 from fugue.exceptions import FugueDataFrameOperationError
@@ -31,7 +30,9 @@ from fugue.plugins import (
     rename,
     select_columns,
 )
-from fugue_spark._utils.convert import to_cast_expression, to_schema, to_type_safe_input
+
+from ._utils.convert import to_cast_expression, to_pandas, to_schema, to_type_safe_input
+from ._utils.misc import is_spark_connect, is_spark_dataframe
 
 
 class SparkDataFrame(DataFrame):
@@ -52,12 +53,12 @@ class SparkDataFrame(DataFrame):
 
     def __init__(self, df: Any = None, schema: Any = None):  # noqa: C901
         self._lock = SerializableRLock()
-        if isinstance(df, ps.DataFrame):
+        if is_spark_dataframe(df):
             if schema is not None:
                 schema = to_schema(schema).assert_not_empty()
                 has_cast, expr = to_cast_expression(df, schema, True)
                 if has_cast:
-                    df = df.selectExpr(*expr)
+                    df = df.selectExpr(*expr)  # type: ignore
             else:
                 schema = to_schema(df).assert_not_empty()
             self._native = df
@@ -90,12 +91,12 @@ class SparkDataFrame(DataFrame):
     def is_bounded(self) -> bool:
         return True
 
-    def as_local(self) -> LocalDataFrame:
+    def as_local_bounded(self) -> LocalBoundedDataFrame:
         if any(pa.types.is_nested(t) for t in self.schema.types):
             data = list(to_type_safe_input(self.native.collect(), self.schema))
-            res: LocalDataFrame = ArrayDataFrame(data, self.schema)
+            res: LocalBoundedDataFrame = ArrayDataFrame(data, self.schema)
         else:
-            res = PandasDataFrame(self.native.toPandas(), self.schema)
+            res = PandasDataFrame(self.as_pandas(), self.schema)
         if self.has_metadata:
             res.reset_metadata(self.metadata)
         return res
@@ -127,7 +128,7 @@ class SparkDataFrame(DataFrame):
         return SparkDataFrame(self.native[schema.names])
 
     def as_pandas(self) -> pd.DataFrame:
-        return self.native.toPandas()
+        return to_pandas(self.native)
 
     def rename(self, columns: Dict[str, str]) -> DataFrame:
         try:
@@ -151,6 +152,9 @@ class SparkDataFrame(DataFrame):
     def as_array_iterable(
         self, columns: Optional[List[str]] = None, type_safe: bool = False
     ) -> Iterable[Any]:
+        if is_spark_connect(self.native):  # pragma: no cover
+            yield from self.as_array(columns, type_safe=type_safe)
+            return
         sdf = self._select_columns(columns)
         if not type_safe:
             for row in to_type_safe_input(sdf.native.rdd.toLocalIterator(), sdf.schema):
@@ -183,47 +187,47 @@ class SparkDataFrame(DataFrame):
         return SparkDataFrame(self.native.select(*columns))
 
 
-@is_df.candidate(lambda df: isinstance(df, ps.DataFrame))
+@is_df.candidate(lambda df: is_spark_dataframe(df))
 def _spark_is_df(df: ps.DataFrame) -> bool:
     return True
 
 
-@get_num_partitions.candidate(lambda df: isinstance(df, ps.DataFrame))
+@get_num_partitions.candidate(lambda df: is_spark_dataframe(df))
 def _spark_num_partitions(df: ps.DataFrame) -> int:
     return df.rdd.getNumPartitions()
 
 
-@count.candidate(lambda df: isinstance(df, ps.DataFrame))
+@count.candidate(lambda df: is_spark_dataframe(df))
 def _spark_df_count(df: ps.DataFrame) -> int:
     return df.count()
 
 
-@is_bounded.candidate(lambda df: isinstance(df, ps.DataFrame))
+@is_bounded.candidate(lambda df: is_spark_dataframe(df))
 def _spark_df_is_bounded(df: ps.DataFrame) -> bool:
     return True
 
 
-@is_empty.candidate(lambda df: isinstance(df, ps.DataFrame))
+@is_empty.candidate(lambda df: is_spark_dataframe(df))
 def _spark_df_is_empty(df: ps.DataFrame) -> bool:
     return df.first() is None
 
 
-@is_local.candidate(lambda df: isinstance(df, ps.DataFrame))
+@is_local.candidate(lambda df: is_spark_dataframe(df))
 def _spark_df_is_local(df: ps.DataFrame) -> bool:
     return False
 
 
-@as_local_bounded.candidate(lambda df: isinstance(df, ps.DataFrame))
+@as_local_bounded.candidate(lambda df: is_spark_dataframe(df))
 def _spark_df_as_local(df: ps.DataFrame) -> pd.DataFrame:
-    return df.toPandas()
+    return to_pandas(df)
 
 
-@get_column_names.candidate(lambda df: isinstance(df, ps.DataFrame))
+@get_column_names.candidate(lambda df: is_spark_dataframe(df))
 def _get_spark_df_columns(df: ps.DataFrame) -> List[Any]:
     return df.columns
 
 
-@rename.candidate(lambda df, *args, **kwargs: isinstance(df, ps.DataFrame))
+@rename.candidate(lambda df, *args, **kwargs: is_spark_dataframe(df))
 def _rename_spark_df(
     df: ps.DataFrame, columns: Dict[str, Any], as_fugue: bool = False
 ) -> ps.DataFrame:
@@ -233,7 +237,7 @@ def _rename_spark_df(
     return _adjust_df(_rename_spark_dataframe(df, columns), as_fugue=as_fugue)
 
 
-@drop_columns.candidate(lambda df, *args, **kwargs: isinstance(df, ps.DataFrame))
+@drop_columns.candidate(lambda df, *args, **kwargs: is_spark_dataframe(df))
 def _drop_spark_df_columns(
     df: ps.DataFrame, columns: List[str], as_fugue: bool = False
 ) -> Any:
@@ -245,7 +249,7 @@ def _drop_spark_df_columns(
     return _adjust_df(df[cols], as_fugue=as_fugue)
 
 
-@select_columns.candidate(lambda df, *args, **kwargs: isinstance(df, ps.DataFrame))
+@select_columns.candidate(lambda df, *args, **kwargs: is_spark_dataframe(df))
 def _select_spark_df_columns(
     df: ps.DataFrame, columns: List[Any], as_fugue: bool = False
 ) -> Any:
@@ -255,7 +259,7 @@ def _select_spark_df_columns(
     return _adjust_df(df[columns], as_fugue=as_fugue)
 
 
-@head.candidate(lambda df, *args, **kwargs: isinstance(df, ps.DataFrame))
+@head.candidate(lambda df, *args, **kwargs: is_spark_dataframe(df))
 def _spark_df_head(
     df: ps.DataFrame,
     n: int,
@@ -265,7 +269,7 @@ def _spark_df_head(
     if columns is not None:
         df = df[columns]
     res = df.limit(n)
-    return SparkDataFrame(res).as_local() if as_fugue else res.toPandas()
+    return SparkDataFrame(res).as_local() if as_fugue else to_pandas(res)
 
 
 def _rename_spark_dataframe(df: ps.DataFrame, names: Dict[str, Any]) -> ps.DataFrame: