fugue 0.8.7.dev5__py3-none-any.whl → 0.8.7.dev7__py3-none-any.whl

fugue_duckdb/dataframe.py

@@ -3,21 +3,33 @@ from typing import Any, Dict, Iterable, List, Optional
 import pandas as pd
 import pyarrow as pa
 from duckdb import DuckDBPyRelation
-from triad import Schema
+from triad import Schema, assert_or_throw
 from triad.utils.pyarrow import LARGE_TYPES_REPLACEMENT, replace_types_in_table
 
-from fugue import ArrayDataFrame, ArrowDataFrame, DataFrame, LocalBoundedDataFrame
+from fugue import ArrowDataFrame, DataFrame, LocalBoundedDataFrame
 from fugue.dataframe.arrow_dataframe import _pa_table_as_pandas
+from fugue.dataframe.utils import (
+    pa_table_as_array,
+    pa_table_as_array_iterable,
+    pa_table_as_dict_iterable,
+    pa_table_as_dicts,
+)
 from fugue.exceptions import FugueDataFrameOperationError, FugueDatasetEmptyError
 from fugue.plugins import (
+    as_array,
+    as_array_iterable,
     as_arrow,
+    as_dict_iterable,
+    as_dicts,
     as_fugue_dataset,
     as_local_bounded,
     as_pandas,
+    drop_columns,
     get_column_names,
     get_num_partitions,
     get_schema,
     is_df,
+    select_columns,
 )
 
 from ._utils import encode_column_name, to_duck_type, to_pa_type
@@ -59,13 +71,10 @@ class DuckDataFrame(LocalBoundedDataFrame):
         return len(self._rel)
 
     def _drop_cols(self, cols: List[str]) -> DataFrame:
-        cols = [col for col in self._rel.columns if col not in cols]
-        rel = self._rel.project(",".join(encode_column_name(n) for n in cols))
-        return DuckDataFrame(rel)
+        return DuckDataFrame(_drop_duckdb_columns(self._rel, cols))
 
     def _select_cols(self, keys: List[Any]) -> DataFrame:
-        rel = self._rel.project(",".join(encode_column_name(n) for n in keys))
-        return DuckDataFrame(rel)
+        return DuckDataFrame(_select_duckdb_columns(self._rel, keys))
 
     def rename(self, columns: Dict[str, str]) -> DataFrame:
         _assert_no_missing(self._rel, columns.keys())
@@ -109,38 +118,29 @@ class DuckDataFrame(LocalBoundedDataFrame):
     def as_array(
         self, columns: Optional[List[str]] = None, type_safe: bool = False
     ) -> List[Any]:
-        if columns is not None:
-            return self[columns].as_array(type_safe=type_safe)
-        return self._fetchall(self._rel)
+        return _duck_as_array(self._rel, columns=columns, type_safe=type_safe)
 
     def as_array_iterable(
         self, columns: Optional[List[str]] = None, type_safe: bool = False
     ) -> Iterable[Any]:
-        if columns is not None:
-            yield from self[columns].as_array_iterable(type_safe=type_safe)
-        else:
-            yield from self._fetchall(self._rel)
+        yield from _duck_as_array_iterable(
+            self._rel, columns=columns, type_safe=type_safe
+        )
+
+    def as_dicts(self, columns: Optional[List[str]] = None) -> List[Dict[str, Any]]:
+        return _duck_as_dicts(self._rel, columns=columns)
+
+    def as_dict_iterable(
+        self, columns: Optional[List[str]] = None
+    ) -> Iterable[Dict[str, Any]]:
+        yield from _duck_as_dict_iterable(self._rel, columns=columns)
 
     def head(
         self, n: int, columns: Optional[List[str]] = None
     ) -> LocalBoundedDataFrame:
         if columns is not None:
             return self[columns].head(n)
-        return ArrayDataFrame(self._fetchall(self._rel.limit(n)), schema=self.schema)
-
-    def _fetchall(self, rel: DuckDBPyRelation) -> List[List[Any]]:
-        map_pos = [i for i, t in enumerate(self.schema.types) if pa.types.is_map(t)]
-        if len(map_pos) == 0:
-            return [list(x) for x in rel.fetchall()]
-        else:
-
-            def to_list(row: Any) -> List[Any]:
-                res = list(row)
-                for p in map_pos:
-                    res[p] = list(zip(row[p]["key"], row[p]["value"]))
-                return res
-
-            return [to_list(x) for x in rel.fetchall()]
+        return ArrowDataFrame(_duck_as_arrow(self._rel.limit(n)))
 
 
 @as_fugue_dataset.candidate(lambda df, **kwargs: isinstance(df, DuckDBPyRelation))
@@ -186,6 +186,64 @@ def _get_duckdb_columns(df: DuckDBPyRelation) -> List[Any]:
     return list(df.columns)
 
 
+@select_columns.candidate(lambda df, *args, **kwargs: isinstance(df, DuckDBPyRelation))
+def _select_duckdb_columns(
+    df: DuckDBPyRelation, columns: List[Any]
+) -> DuckDBPyRelation:
+    if len(columns) == 0:
+        raise FugueDataFrameOperationError("must select at least one column")
+    _assert_no_missing(df, columns)
+    return df.project(",".join(encode_column_name(n) for n in columns))
+
+
+@drop_columns.candidate(lambda df, *args, **kwargs: isinstance(df, DuckDBPyRelation))
+def _drop_duckdb_columns(df: DuckDBPyRelation, columns: List[str]) -> DuckDBPyRelation:
+    # if len(columns) == 0:
+    #   return df
+    _columns = {c: 1 for c in columns}
+    cols = [col for col in df.columns if _columns.pop(col, None) is None]
+    assert_or_throw(
+        len(cols) > 0, FugueDataFrameOperationError("must keep at least one column")
+    )
+    assert_or_throw(
+        len(_columns) == 0,
+        FugueDataFrameOperationError("found nonexistent columns {_columns}"),
+    )
+    return df.project(",".join(encode_column_name(n) for n in cols))
+
+
+@as_array.candidate(lambda df, *args, **kwargs: isinstance(df, DuckDBPyRelation))
+def _duck_as_array(
+    df: DuckDBPyRelation, columns: Optional[List[str]] = None, type_safe: bool = False
+) -> List[Any]:
+    return pa_table_as_array(df.arrow(), columns=columns)
+
+
+@as_array_iterable.candidate(
+    lambda df, *args, **kwargs: isinstance(df, DuckDBPyRelation)
+)
+def _duck_as_array_iterable(
+    df: DuckDBPyRelation, columns: Optional[List[str]] = None, type_safe: bool = False
+) -> Iterable[Any]:
+    yield from pa_table_as_array_iterable(df.arrow(), columns=columns)
+
+
+@as_dicts.candidate(lambda df, *args, **kwargs: isinstance(df, DuckDBPyRelation))
+def _duck_as_dicts(
+    df: DuckDBPyRelation, columns: Optional[List[str]] = None
+) -> List[Dict[str, Any]]:
+    return pa_table_as_dicts(df.arrow(), columns=columns)
+
+
+@as_dict_iterable.candidate(
+    lambda df, *args, **kwargs: isinstance(df, DuckDBPyRelation)
+)
+def _duck_as_dict_iterable(
+    df: DuckDBPyRelation, columns: Optional[List[str]] = None
+) -> Iterable[Dict[str, Any]]:
+    yield from pa_table_as_dict_iterable(df.arrow(), columns=columns)
+
+
 def _assert_no_missing(df: DuckDBPyRelation, columns: Iterable[Any]) -> None:
     missing = set(columns) - set(df.columns)
     if len(missing) > 0:

fugue_duckdb/execution_engine.py

@@ -4,7 +4,6 @@ from typing import Any, Dict, Iterable, List, Optional, Union
 import duckdb
 from duckdb import DuckDBPyConnection, DuckDBPyRelation
 from triad import SerializableRLock
-from triad.collections.fs import FileSystem
 from triad.utils.assertion import assert_or_throw
 from triad.utils.schema import quote_name
 
@@ -195,10 +194,6 @@ class DuckExecutionEngine(ExecutionEngine):
     def log(self) -> logging.Logger:
         return self._native_engine.log
 
-    @property
-    def fs(self) -> FileSystem:
-        return self._native_engine.fs
-
     def create_default_sql_engine(self) -> SQLEngine:
         return DuckDBEngine(self)
 
@@ -488,7 +483,7 @@ class DuckExecutionEngine(ExecutionEngine):
         columns: Any = None,
         **kwargs: Any,
     ) -> LocalBoundedDataFrame:
-        dio = DuckDBIO(self.fs, self.connection)
+        dio = DuckDBIO(self.connection)
         return dio.load_df(path, format_hint, columns, **kwargs)
 
     def save_df(
@@ -504,7 +499,7 @@ class DuckExecutionEngine(ExecutionEngine):
         partition_spec = partition_spec or PartitionSpec()
         if not partition_spec.empty and not force_single:
             kwargs["partition_cols"] = partition_spec.partition_by
-        dio = DuckDBIO(self.fs, self.connection)
+        dio = DuckDBIO(self.connection)
         dio.save_df(_to_duck_df(self, df), path, format_hint, mode, **kwargs)
 
     def convert_yield_dataframe(self, df: DataFrame, as_local: bool) -> DataFrame:

fugue_ibis/dataframe.py

@@ -143,6 +143,19 @@ class IbisDataFrame(DataFrame):
                 type_safe=type_safe
             )
 
+    def as_dicts(self, columns: Optional[List[str]] = None) -> List[Dict[str, Any]]:
+        if columns is not None:
+            return self[columns].as_dicts()
+        return self.as_local().as_dicts()
+
+    def as_dict_iterable(
+        self, columns: Optional[List[str]] = None
+    ) -> Iterable[Dict[str, Any]]:
+        if columns is not None:
+            yield from self[columns].as_dict_iterable()
+        else:
+            yield from self._to_iterable_df(self._table).as_dict_iterable()
+
     def head(
         self, n: int, columns: Optional[List[str]] = None
     ) -> LocalBoundedDataFrame:

fugue_ibis/execution_engine.py

@@ -5,7 +5,7 @@ from typing import Any, Callable, Dict, List, Optional, Type
 
 import ibis
 from ibis import BaseBackend
-from triad import FileSystem, assert_or_throw
+from triad import assert_or_throw
 
 from fugue import StructuredRawSQL
 from fugue.bag import Bag, LocalBag
@@ -375,10 +375,6 @@ class IbisExecutionEngine(ExecutionEngine):
     def log(self) -> logging.Logger:
         return self.non_ibis_engine.log
 
-    @property
-    def fs(self) -> FileSystem:
-        return self.non_ibis_engine.fs
-
     def get_current_parallelism(self) -> int:
         return self.non_ibis_engine.get_current_parallelism()
 

fugue_polars/polars_dataframe.py

@@ -14,22 +14,32 @@ from triad.utils.pyarrow import (
 
 from fugue import ArrowDataFrame
 from fugue.api import (
+    as_array,
+    as_array_iterable,
     as_arrow,
+    as_dict_iterable,
+    as_dicts,
     drop_columns,
     get_column_names,
     get_schema,
     is_df,
+    is_empty,
     rename,
     select_columns,
 )
 from fugue.dataframe.dataframe import DataFrame, LocalBoundedDataFrame, _input_schema
+from fugue.dataframe.utils import (
+    pa_table_as_array,
+    pa_table_as_array_iterable,
+    pa_table_as_dict_iterable,
+    pa_table_as_dicts,
+)
 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
@@ -52,7 +62,7 @@ class PolarsDataFrame(LocalBoundedDataFrame):
     ):
         if df is None:
             schema = _input_schema(schema).assert_not_empty()
-            self._native: pa.Table = build_empty_pl(schema)
+            self._native: pl.DataFrame = build_empty_pl(schema)
             super().__init__(schema)
             return
         else:
@@ -73,7 +83,7 @@ class PolarsDataFrame(LocalBoundedDataFrame):
 
     @property
     def empty(self) -> bool:
-        return self._native.shape[0] == 0
+        return self._native.is_empty()
 
     def peek_array(self) -> List[Any]:
         self.assert_not_empty()
@@ -118,26 +128,20 @@ class PolarsDataFrame(LocalBoundedDataFrame):
     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()]
+        return _pl_as_array(self.native, columns=columns)
 
     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
-            )
+        yield from _pl_as_array_iterable(self.native, columns=columns)
+
+    def as_dicts(self, columns: Optional[List[str]] = None) -> List[Dict[str, Any]]:
+        return _pl_as_dicts(self.native, 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
-            )
+        yield from _pl_as_dict_iterable(self.native, columns=columns)
 
 
 @as_local.candidate(lambda df: isinstance(df, pl.DataFrame))
@@ -174,7 +178,7 @@ def _pl_is_bounded(df: pl.DataFrame) -> bool:
 
 @is_empty.candidate(lambda df: isinstance(df, pl.DataFrame))
 def _pl_is_empty(df: pl.DataFrame) -> bool:
-    return df.shape[0] == 0
+    return df.is_empty()
 
 
 @is_local.candidate(lambda df: isinstance(df, pl.DataFrame))
@@ -228,6 +232,39 @@ def _select_pa_columns(df: pl.DataFrame, columns: List[Any]) -> pl.DataFrame:
     return df.select(columns)
 
 
+@as_array.candidate(lambda df, *args, **kwargs: isinstance(df, pl.DataFrame))
+def _pl_as_array(
+    df: pl.DataFrame, columns: Optional[List[str]] = None, type_safe: bool = False
+) -> List[List[Any]]:
+    _df = df if columns is None else _select_pa_columns(df, columns)
+    adf = _pl_as_arrow(_df)
+    return pa_table_as_array(adf, columns=columns)
+
+
+@as_array_iterable.candidate(lambda df, *args, **kwargs: isinstance(df, pl.DataFrame))
+def _pl_as_array_iterable(
+    df: pl.DataFrame, columns: Optional[List[str]] = None, type_safe: bool = False
+) -> Iterable[List[Any]]:
+    _df = df if columns is None else _select_pa_columns(df, columns)
+    yield from pa_table_as_array_iterable(_df.to_arrow(), columns=columns)
+
+
+@as_dicts.candidate(lambda df, *args, **kwargs: isinstance(df, pl.DataFrame))
+def _pl_as_dicts(
+    df: pl.DataFrame, columns: Optional[List[str]] = None
+) -> List[Dict[str, Any]]:
+    _df = df if columns is None else _select_pa_columns(df, columns)
+    return pa_table_as_dicts(_df.to_arrow(), columns=columns)
+
+
+@as_dict_iterable.candidate(lambda df, *args, **kwargs: isinstance(df, pl.DataFrame))
+def _pl_as_dict_iterable(
+    df: pl.DataFrame, columns: Optional[List[str]] = None
+) -> Iterable[Dict[str, Any]]:
+    _df = df if columns is None else _select_pa_columns(df, columns)
+    yield from pa_table_as_dict_iterable(_df.to_arrow(), columns=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:

fugue_ray/_utils/io.py

@@ -4,23 +4,24 @@ 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, save_df
-from fugue.collections.partition import PartitionSpec
-from fugue.dataframe import DataFrame
-from fugue_ray.dataframe import RayDataFrame
 from pyarrow import csv as pacsv
 from pyarrow import json as pajson
 from ray.data.datasource import FileExtensionFilter
 from triad.collections import Schema
 from triad.collections.dict import ParamDict
 from triad.utils.assertion import assert_or_throw
+from triad.utils.io import exists, makedirs, rm
+
+from fugue import ExecutionEngine
+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
 
 
 class RayIO(object):
     def __init__(self, engine: ExecutionEngine):
         self._engine = engine
-        self._fs = engine.fs
         self._logger = engine.log
         self._loads: Dict[str, Callable[..., DataFrame]] = {
             "csv": self._load_csv,
@@ -49,7 +50,7 @@ class RayIO(object):
             len(fmts) == 1, NotImplementedError("can't support multiple formats")
         )
         fmt = fmts[0]
-        files = [f.uri for f in fp]
+        files = [f.path for f in fp]
         return self._loads[fmt](files, columns, **kwargs)
 
     def save_df(
@@ -63,24 +64,21 @@ class RayIO(object):
         **kwargs: Any,
     ) -> None:
         partition_spec = partition_spec or PartitionSpec()
-        if self._fs.exists(uri):
+        if exists(uri):
             assert_or_throw(mode == "overwrite", FileExistsError(uri))
             try:
-                self._fs.remove(uri)
-            except Exception:
-                try:
-                    self._fs.removetree(uri)
-                except Exception:  # pragma: no cover
-                    pass
+                rm(uri, recursive=True)
+            except Exception:  # pragma: no cover
+                pass
         p = FileParser(uri, format_hint)
         if not force_single:
             df = self._prepartition(df, partition_spec=partition_spec)
 
-            self._saves[p.file_format](df=df, uri=p.uri, **kwargs)
+            self._saves[p.file_format](df=df, uri=p.path, **kwargs)
         else:
             ldf = df.as_local()
-            self._fs.makedirs(os.path.dirname(uri), recreate=True)
-            save_df(ldf, uri, format_hint=format_hint, mode=mode, fs=self._fs, **kwargs)
+            makedirs(os.path.dirname(uri), exist_ok=True)
+            save_df(ldf, uri, format_hint=format_hint, mode=mode, **kwargs)
 
     def _save_parquet(
         self,

fugue_ray/dataframe.py

@@ -4,14 +4,22 @@ import pandas as pd
 import pyarrow as pa
 import ray
 import ray.data as rd
+from triad import assert_or_throw
 from triad.collections.schema import Schema
 from triad.utils.pyarrow import cast_pa_table
 
 from fugue.dataframe import ArrowDataFrame, DataFrame, LocalBoundedDataFrame
 from fugue.dataframe.dataframe import _input_schema
+from fugue.dataframe.utils import pa_table_as_array, pa_table_as_dicts
 from fugue.exceptions import FugueDataFrameOperationError, FugueDatasetEmptyError
 from fugue.plugins import (
+    as_array,
+    as_array_iterable,
+    as_arrow,
+    as_dict_iterable,
+    as_dicts,
     as_local_bounded,
+    as_pandas,
     get_column_names,
     get_num_partitions,
     is_df,
@@ -141,13 +149,11 @@ class RayDataFrame(DataFrame):
     def _select_cols(self, cols: List[Any]) -> DataFrame:
         if cols == self.columns:
             return self
-        rdf = self.native.map_batches(
-            lambda b: b.select(cols),
-            batch_format="pyarrow",
-            **_ZERO_COPY,
-            **self._remote_args(),
+        return RayDataFrame(
+            self.native.select_columns(cols),
+            self.schema.extract(cols),
+            internal_schema=True,
         )
-        return RayDataFrame(rdf, self.schema.extract(cols), internal_schema=True)
 
     def peek_array(self) -> List[Any]:
         data = self.native.limit(1).to_pandas().values.tolist()
@@ -164,10 +170,10 @@ class RayDataFrame(DataFrame):
         return self.native.count()
 
     def as_arrow(self, type_safe: bool = False) -> pa.Table:
-        return pa.concat_tables(_get_arrow_tables(self.native))
+        return _rd_as_arrow(self.native)
 
     def as_pandas(self) -> pd.DataFrame:
-        return self.as_arrow().to_pandas()
+        return _rd_as_pandas(self.native)
 
     def rename(self, columns: Dict[str, str]) -> DataFrame:
         try:
@@ -201,18 +207,20 @@ class RayDataFrame(DataFrame):
     def as_array(
         self, columns: Optional[List[str]] = None, type_safe: bool = False
     ) -> List[Any]:
-        df: DataFrame = self
-        if columns is not None:
-            df = df[columns]
-        adf = df.as_arrow()
-        if adf.shape[0] == 0:
-            return []
-        return ArrowDataFrame(adf).as_array(type_safe=type_safe)
+        return _rd_as_array(self.native, columns, type_safe)
 
     def as_array_iterable(
         self, columns: Optional[List[str]] = None, type_safe: bool = False
     ) -> Iterable[Any]:
-        yield from self.as_array(columns=columns, type_safe=type_safe)
+        yield from _rd_as_array_iterable(self.native, columns, type_safe)
+
+    def as_dicts(self, columns: Optional[List[str]] = None) -> List[Dict[str, Any]]:
+        return _rd_as_dicts(self.native, columns)
+
+    def as_dict_iterable(
+        self, columns: Optional[List[str]] = None
+    ) -> Iterable[Dict[str, Any]]:
+        yield from _rd_as_dict_iterable(self.native, columns)
 
     def head(
         self, n: int, columns: Optional[List[str]] = None
@@ -259,8 +267,8 @@ def _rd_num_partitions(df: rd.Dataset) -> int:
 
 
 @as_local_bounded.candidate(lambda df: isinstance(df, rd.Dataset))
-def _rd_as_local(df: rd.Dataset) -> bool:
-    return pa.concat_tables(_get_arrow_tables(df))
+def _rd_as_local(df: rd.Dataset) -> pa.Table:
+    return _rd_as_arrow(df)
 
 
 @get_column_names.candidate(lambda df: isinstance(df, rd.Dataset))
@@ -290,10 +298,54 @@ def _rename_ray_dataframe(df: rd.Dataset, columns: Dict[str, Any]) -> rd.Dataset
     )
 
 
+@as_pandas.candidate(lambda df: isinstance(df, rd.Dataset))
+def _rd_as_pandas(df: rd.Dataset) -> pd.DataFrame:
+    return _rd_as_arrow(df).to_pandas()
+
+
+@as_arrow.candidate(lambda df: isinstance(df, rd.Dataset))
+def _rd_as_arrow(df: rd.Dataset) -> pa.Table:
+    return pa.concat_tables(_get_arrow_tables(df))
+
+
+@as_array.candidate(lambda df, *args, **kwargs: isinstance(df, rd.Dataset))
+def _rd_as_array(
+    df: rd.Dataset, columns: Optional[List[str]] = None, type_safe: bool = False
+) -> List[Any]:
+    assert_or_throw(columns is None or len(columns) > 0, ValueError("empty columns"))
+    _df = df if columns is None or len(columns) == 0 else df.select_columns(columns)
+    adf = _rd_as_arrow(_df)
+    return pa_table_as_array(adf)
+
+
+@as_array_iterable.candidate(lambda df, *args, **kwargs: isinstance(df, rd.Dataset))
+def _rd_as_array_iterable(
+    df: rd.Dataset, columns: Optional[List[str]] = None, type_safe: bool = False
+) -> Iterable[Any]:
+    yield from _rd_as_array(df, columns, type_safe)
+
+
+@as_dicts.candidate(lambda df, *args, **kwargs: isinstance(df, rd.Dataset))
+def _rd_as_dicts(
+    df: rd.Dataset, columns: Optional[List[str]] = None, type_safe: bool = False
+) -> List[Dict[str, Any]]:
+    assert_or_throw(columns is None or len(columns) > 0, ValueError("empty columns"))
+    _df = df if columns is None or len(columns) == 0 else df.select_columns(columns)
+    adf = _rd_as_arrow(_df)
+    return pa_table_as_dicts(adf)
+
+
+@as_dict_iterable.candidate(lambda df, *args, **kwargs: isinstance(df, rd.Dataset))
+def _rd_as_dict_iterable(
+    df: rd.Dataset, columns: Optional[List[str]] = None, type_safe: bool = False
+) -> Iterable[Dict[str, Any]]:
+    yield from _rd_as_dicts(df, columns, type_safe)
+
+
 def _get_arrow_tables(df: rd.Dataset) -> Iterable[pa.Table]:
     last_empty: Any = None
     empty = True
-    for block in df.get_internal_block_refs():
+    for block in df.to_arrow_refs():
         tb = ray.get(block)
         if tb.shape[0] > 0:
             yield tb

fugue_spark/_utils/io.py

@@ -4,7 +4,6 @@ import pyspark.sql as ps
 from pyspark.sql import SparkSession
 from triad.collections import Schema
 from triad.collections.dict import ParamDict
-from triad.collections.fs import FileSystem
 from triad.utils.assertion import assert_or_throw
 
 from fugue._utils.io import FileParser, save_df
@@ -16,9 +15,8 @@ from .convert import to_schema, to_spark_schema
 
 
 class SparkIO(object):
-    def __init__(self, spark_session: SparkSession, fs: FileSystem):
+    def __init__(self, spark_session: SparkSession):
         self._session = spark_session
-        self._fs = fs
         self._loads: Dict[str, Callable[..., DataFrame]] = {
             "csv": self._load_csv,
             "parquet": self._load_parquet,
@@ -41,7 +39,7 @@ class SparkIO(object):
             len(fmts) == 1, NotImplementedError("can't support multiple formats")
         )
         fmt = fmts[0]
-        files = [f.uri for f in fp]
+        files = [f.path for f in fp]
         return self._loads[fmt](files, columns, **kwargs)
 
     def save_df(
@@ -64,7 +62,7 @@ class SparkIO(object):
             ldf = df.as_local()
             if isinstance(ldf, PandasDataFrame) and hasattr(ldf.native, "attrs"):
                 ldf.native.attrs = {}  # pragma: no cover
-            save_df(ldf, uri, format_hint=format_hint, mode=mode, fs=self._fs, **kwargs)
+            save_df(ldf, uri, format_hint=format_hint, mode=mode, **kwargs)
 
     def _get_writer(
         self, sdf: ps.DataFrame, partition_spec: PartitionSpec