fugue 0.8.7.dev4__py3-none-any.whl → 0.8.7.dev6__py3-none-any.whl

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

@@ -1,5 +1,6 @@
 import pickle
-from typing import Any, Iterable, List, Tuple
+from typing import Any, Iterable, List, Tuple, Optional
+
 
 import pandas as pd
 import pyarrow as pa
@@ -16,7 +17,7 @@ from pyspark.sql.pandas.types import (
 )
 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.pyarrow import TRIAD_DEFAULT_TIMESTAMP, cast_pa_table
 from triad.utils.schema import quote_name
 
 import fugue.api as fa
@@ -41,7 +42,7 @@ def pandas_udf_can_accept(schema: Schema, is_input: bool) -> bool:
             return False
         to_arrow_schema(from_arrow_schema(schema.pa_schema))
         return True
-    except Exception:
+    except Exception:  # pragma: no cover
         return False
 
 
@@ -132,7 +133,7 @@ def to_type_safe_input(rows: Iterable[ps.Row], schema: Schema) -> Iterable[List[
                 if r[i] is not None:
                     r[i] = r[i].asDict(recursive=True)
             yield r
-    else:
+    else:  # pragma: no cover
         for row in rows:
             data = row.asDict(recursive=True)
             r = [data[n] for n in schema.names]
@@ -173,14 +174,14 @@ def pd_to_spark_df(
 
 
 def to_pandas(df: ps.DataFrame) -> pd.DataFrame:
-    if pd.__version__ < "2" or not any(
+    if version.parse(pd.__version__) < version.parse("2.0.0") or not any(
         isinstance(x.dataType, (pt.TimestampType, TimestampNTZType))
         for x in df.schema.fields
     ):
         return df.toPandas()
-    else:
+    else:  # pragma: no cover
 
-        def serialize(dfs):  # pragma: no cover
+        def serialize(dfs):
             for df in dfs:
                 data = pickle.dumps(df)
                 yield pd.DataFrame([[data]], columns=["data"])
@@ -189,6 +190,30 @@ def to_pandas(df: ps.DataFrame) -> pd.DataFrame:
         return pd.concat(pickle.loads(x.data) for x in sdf.collect())
 
 
+def to_arrow(df: ps.DataFrame) -> pa.Table:
+    schema = to_schema(df.schema)
+    destruct: Optional[bool] = None
+    try:
+        jconf = df.sparkSession._jconf
+        if jconf.arrowPySparkEnabled() and pandas_udf_can_accept(
+            schema, is_input=False
+        ):
+            destruct = jconf.arrowPySparkSelfDestructEnabled()
+    except Exception:  # pragma: no cover
+        # older spark does not have this config
+        pass
+    if destruct is not None and hasattr(df, "_collect_as_arrow"):
+        batches = df._collect_as_arrow(split_batches=destruct)
+        if len(batches) == 0:
+            return schema.create_empty_arrow_table()
+        table = pa.Table.from_batches(batches)
+        del batches
+        return cast_pa_table(table, schema.pa_schema)
+    else:  # pragma: no cover
+        # df.toPandas has bugs on nested types
+        return pa.Table.from_pylist(df.collect(), schema=schema.pa_schema)
+
+
 # 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

@@ -9,7 +9,7 @@ from triad.utils.assertion import assert_or_throw
 
 from fugue._utils.io import FileParser, save_df
 from fugue.collections.partition import PartitionSpec
-from fugue.dataframe import DataFrame
+from fugue.dataframe import DataFrame, PandasDataFrame
 from fugue_spark.dataframe import SparkDataFrame
 
 from .convert import to_schema, to_spark_schema
@@ -62,6 +62,8 @@ class SparkIO(object):
             writer.save(uri)
         else:
             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)
 
     def _get_writer(

fugue_spark/dataframe.py

@@ -9,15 +9,21 @@ from triad.collections.schema import SchemaError
 from triad.utils.assertion import assert_or_throw
 
 from fugue.dataframe import (
-    ArrayDataFrame,
+    ArrowDataFrame,
     DataFrame,
     IterableDataFrame,
     LocalBoundedDataFrame,
-    PandasDataFrame,
 )
+from fugue.dataframe.utils import pa_table_as_array, pa_table_as_dicts
 from fugue.exceptions import FugueDataFrameOperationError
 from fugue.plugins import (
+    as_array,
+    as_array_iterable,
+    as_arrow,
+    as_dict_iterable,
+    as_dicts,
     as_local_bounded,
+    as_pandas,
     count,
     drop_columns,
     get_column_names,
@@ -31,7 +37,13 @@ from fugue.plugins import (
     select_columns,
 )
 
-from ._utils.convert import to_cast_expression, to_pandas, to_schema, to_type_safe_input
+from ._utils.convert import (
+    to_arrow,
+    to_cast_expression,
+    to_pandas,
+    to_schema,
+    to_type_safe_input,
+)
 from ._utils.misc import is_spark_connect, is_spark_dataframe
 
 
@@ -92,11 +104,7 @@ class SparkDataFrame(DataFrame):
         return True
 
     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: LocalBoundedDataFrame = ArrayDataFrame(data, self.schema)
-        else:
-            res = PandasDataFrame(self.as_pandas(), self.schema)
+        res = ArrowDataFrame(self.as_arrow())
         if self.has_metadata:
             res.reset_metadata(self.metadata)
         return res
@@ -128,7 +136,10 @@ class SparkDataFrame(DataFrame):
         return SparkDataFrame(self.native[schema.names])
 
     def as_pandas(self) -> pd.DataFrame:
-        return to_pandas(self.native)
+        return _spark_df_as_pandas(self.native)
+
+    def as_arrow(self, type_safe: bool = False) -> pa.Table:
+        return _spark_df_as_arrow(self.native)
 
     def rename(self, columns: Dict[str, str]) -> DataFrame:
         try:
@@ -146,23 +157,22 @@ class SparkDataFrame(DataFrame):
     def as_array(
         self, columns: Optional[List[str]] = None, type_safe: bool = False
     ) -> List[Any]:
-        sdf = self._select_columns(columns)
-        return sdf.as_local().as_array(type_safe=type_safe)
+        return _spark_as_array(self.native, columns=columns, type_safe=type_safe)
 
     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):
-                yield row
-        else:
-            df = IterableDataFrame(sdf.as_array_iterable(type_safe=False), sdf.schema)
-            for row in df.as_array_iterable(type_safe=True):
-                yield row
+        yield from _spark_as_array_iterable(
+            self.native, columns=columns, type_safe=type_safe
+        )
+
+    def as_dicts(self, columns: Optional[List[str]] = None) -> List[Dict[str, Any]]:
+        return _spark_as_dicts(self.native, columns=columns)
+
+    def as_dict_iterable(
+        self, columns: Optional[List[str]] = None
+    ) -> Iterable[Dict[str, Any]]:
+        yield from _spark_as_dict_iterable(self.native, columns=columns)
 
     def head(
         self, n: int, columns: Optional[List[str]] = None
@@ -192,6 +202,16 @@ def _spark_is_df(df: ps.DataFrame) -> bool:
     return True
 
 
+@as_arrow.candidate(lambda df: isinstance(df, ps.DataFrame))
+def _spark_df_as_arrow(df: ps.DataFrame) -> pd.DataFrame:
+    return to_arrow(df)
+
+
+@as_pandas.candidate(lambda df: isinstance(df, ps.DataFrame))
+def _spark_df_as_pandas(df: ps.DataFrame) -> pd.DataFrame:
+    return to_pandas(df)
+
+
 @get_num_partitions.candidate(lambda df: is_spark_dataframe(df))
 def _spark_num_partitions(df: ps.DataFrame) -> int:
     return df.rdd.getNumPartitions()
@@ -272,6 +292,58 @@ def _spark_df_head(
     return SparkDataFrame(res).as_local() if as_fugue else to_pandas(res)
 
 
+@as_array.candidate(lambda df, *args, **kwargs: is_spark_dataframe(df))
+def _spark_as_array(
+    df: ps.DataFrame, 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[columns]
+    return pa_table_as_array(to_arrow(_df), columns)
+
+
+@as_array_iterable.candidate(lambda df, *args, **kwargs: is_spark_dataframe(df))
+def _spark_as_array_iterable(
+    df: ps.DataFrame, columns: Optional[List[str]] = None, type_safe: bool = False
+) -> Iterable[Any]:
+    if is_spark_connect(df):  # pragma: no cover
+        yield from _spark_as_array(df, columns, type_safe=type_safe)
+    else:
+        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[columns]
+        if not type_safe:
+            for row in to_type_safe_input(
+                _df.rdd.toLocalIterator(), to_schema(_df.schema)
+            ):
+                yield list(row)
+        else:
+            tdf = IterableDataFrame(
+                _spark_as_array_iterable(_df, type_safe=False), to_schema(_df.schema)
+            )
+            yield from tdf.as_array_iterable(type_safe=True)
+
+
+@as_dicts.candidate(lambda df, *args, **kwargs: is_spark_dataframe(df))
+def _spark_as_dicts(
+    df: ps.DataFrame, 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[columns]
+    return pa_table_as_dicts(to_arrow(_df), columns)
+
+
+@as_dict_iterable.candidate(lambda df, *args, **kwargs: is_spark_dataframe(df))
+def _spark_as_dict_iterable(
+    df: ps.DataFrame, columns: Optional[List[str]] = None, type_safe: bool = False
+) -> Iterable[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[columns]
+    cols = list(_df.columns)
+    for row in _spark_as_array_iterable(_df, type_safe=type_safe):
+        yield dict(zip(cols, row))
+
+
 def _rename_spark_dataframe(df: ps.DataFrame, names: Dict[str, Any]) -> ps.DataFrame:
     cols: List[ps.Column] = []
     for f in df.schema:

fugue_spark/execution_engine.py

@@ -5,6 +5,7 @@ from uuid import uuid4
 import pandas as pd
 import pyarrow as pa
 import pyspark.sql as ps
+from py4j.protocol import Py4JError
 from pyspark import StorageLevel
 from pyspark.rdd import RDD
 from pyspark.sql import SparkSession
@@ -350,9 +351,12 @@ class SparkExecutionEngine(ExecutionEngine):
         self._spark_session = spark_session
         cf = dict(FUGUE_SPARK_DEFAULT_CONF)
         if not self.is_spark_connect:
-            cf.update(
-                {x[0]: x[1] for x in spark_session.sparkContext.getConf().getAll()}
-            )
+            try:
+                spark_conf = spark_session.sparkContext.getConf()
+                cf.update({x[0]: x[1] for x in spark_conf.getAll()})
+            except Py4JError:  # pragma: no cover:
+                # edge case: https://github.com/fugue-project/fugue/issues/517z
+                pass
         cf.update(ParamDict(conf))
         super().__init__(cf)
         self._lock = SerializableRLock()

fugue_test/builtin_suite.py

@@ -1329,7 +1329,7 @@ class BuiltInTests(object):
                     [[datetime.date(2020, 1, 1), datetime.datetime(2020, 1, 2)]],
                     "a:date,b:datetime",
                 )
-                b.assert_eq(a)
+                b.assert_eq(a, no_pandas=True)
                 c = dag.df([["2020-01-01", "2020-01-01 00:00:00"]], "a:date,b:datetime")
                 c.transform(T2).assert_eq(c)
                 c.partition(by=["a"]).transform(T2).assert_eq(c)

fugue_test/dataframe_suite.py

@@ -208,8 +208,22 @@ class DataFrameTests(object):
         def test_as_dict_iterable(self):
             df = self.df([[pd.NaT, 1]], "a:datetime,b:int")
             assert [dict(a=None, b=1)] == list(fi.as_dict_iterable(df))
+            df = self.df([[pd.NaT, 1]], "a:datetime,b:int")
+            assert [dict(b=1)] == list(fi.as_dict_iterable(df, ["b"]))
             df = self.df([[pd.Timestamp("2020-01-01"), 1]], "a:datetime,b:int")
             assert [dict(a=datetime(2020, 1, 1), b=1)] == list(fi.as_dict_iterable(df))
+            df = self.df([[pd.Timestamp("2020-01-01"), 1]], "a:datetime,b:int")
+            assert [dict(b=1)] == list(fi.as_dict_iterable(df, ["b"]))
+
+        def test_as_dicts(self):
+            df = self.df([[pd.NaT, 1]], "a:datetime,b:int")
+            assert [dict(a=None, b=1)] == fi.as_dicts(df)
+            df = self.df([[pd.NaT, 1]], "a:datetime,b:int")
+            assert [dict(b=1)] == fi.as_dicts(df, ["b"])
+            df = self.df([[pd.Timestamp("2020-01-01"), 1]], "a:datetime,b:int")
+            assert [dict(a=datetime(2020, 1, 1), b=1)] == fi.as_dicts(df)
+            df = self.df([[pd.Timestamp("2020-01-01"), 1]], "a:datetime,b:int")
+            assert [dict(b=1)] == fi.as_dicts(df, ["b"])
 
         def test_list_type(self):
             data = [[[30, 40]]]