fugue 0.8.6.dev1__tar.gz → 0.8.6.dev2__tar.gz

{fugue-0.8.6.dev1 → fugue-0.8.6.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fugue
-Version: 0.8.6.dev1
+Version: 0.8.6.dev2
 Summary: An abstraction layer for distributed computation
 Home-page: http://github.com/fugue-project/fugue
 Author: The Fugue Development Team

{fugue-0.8.6.dev1 → fugue-0.8.6.dev2}/fugue/dataframe/dataframe.py

@@ -113,6 +113,13 @@ class DataFrame(Dataset):
     def as_pandas(self) -> pd.DataFrame:
         """Convert to pandas DataFrame"""
         pdf = pd.DataFrame(self.as_array(), columns=self.columns)
+        if len(pdf) == 0:  # TODO: move to triad
+            return pd.DataFrame(
+                {
+                    k: pd.Series(dtype=v.type.to_pandas_dtype())
+                    for k, v in self.schema.items()
+                }
+            )
         return PD_UTILS.enforce_type(pdf, self.schema.pa_schema, null_safe=True)
 
     def as_arrow(self, type_safe: bool = False) -> pa.Table:

{fugue-0.8.6.dev1 → fugue-0.8.6.dev2}/fugue/dataframe/dataframe_iterable_dataframe.py

@@ -165,7 +165,7 @@ class LocalDataFrameIterableDataFrame(LocalUnboundedDataFrame):
 
     def as_pandas(self) -> pd.DataFrame:
         if self.empty:
-            return ArrayDataFrame([], self.schema).as_pandas()
+            return PandasDataFrame(schema=self.schema).as_pandas()
 
         return pd.concat(df.as_pandas() for df in self.native)
 

{fugue-0.8.6.dev1 → fugue-0.8.6.dev2}/fugue/dataframe/function_wrapper.py

@@ -23,9 +23,10 @@ from triad.collections.function_wrapper import (
 from triad.utils.iter import EmptyAwareIterable, make_empty_aware
 
 from ..constants import FUGUE_ENTRYPOINT
+from ..dataset.api import count as df_count
 from .array_dataframe import ArrayDataFrame
 from .arrow_dataframe import ArrowDataFrame
-from .dataframe import DataFrame, LocalDataFrame
+from .dataframe import AnyDataFrame, DataFrame, LocalDataFrame, as_fugue_df
 from .dataframe_iterable_dataframe import (
     IterableArrowDataFrame,
     IterablePandasDataFrame,
@@ -172,6 +173,19 @@ class DataFrameParam(_DataFrameParamBase):
             return sum(1 for _ in df.as_array_iterable())
 
 
+@fugue_annotated_param(AnyDataFrame)
+class _AnyDataFrameParam(DataFrameParam):
+    def to_output_df(self, output: AnyDataFrame, schema: Any, ctx: Any) -> DataFrame:
+        return (
+            as_fugue_df(output)
+            if schema is None
+            else as_fugue_df(output, schema=schema)
+        )
+
+    def count(self, df: Any) -> int:
+        return df_count(df)
+
+
 @fugue_annotated_param(LocalDataFrame, "l", child_can_reuse_code=True)
 class LocalDataFrameParam(DataFrameParam):
     def to_input_data(self, df: DataFrame, ctx: Any) -> LocalDataFrame:
@@ -333,6 +347,9 @@ class _PandasParam(LocalDataFrameParam):
 
     @no_type_check
     def to_output_df(self, output: pd.DataFrame, schema: Any, ctx: Any) -> DataFrame:
+        _schema: Optional[Schema] = None if schema is None else Schema(schema)
+        if _schema is not None and _schema.names != list(output.columns):
+            output = output[_schema.names]
         return PandasDataFrame(output, schema)
 
     @no_type_check
@@ -361,8 +378,15 @@ class _IterablePandasParam(LocalDataFrameParam):
         self, output: Iterable[pd.DataFrame], schema: Any, ctx: Any
     ) -> DataFrame:
         def dfs():
+            _schema: Optional[Schema] = None if schema is None else Schema(schema)
+            has_return = False
             for df in output:
-                yield PandasDataFrame(df, schema)
+                if _schema is not None and _schema.names != list(df.columns):
+                    df = df[_schema.names]
+                yield PandasDataFrame(df, _schema)
+                has_return = True
+            if not has_return and _schema is not None:
+                yield PandasDataFrame(schema=_schema)
 
         return IterablePandasDataFrame(dfs())
 
@@ -381,7 +405,12 @@ class _PyArrowTableParam(LocalDataFrameParam):
 
     def to_output_df(self, output: Any, schema: Any, ctx: Any) -> DataFrame:
         assert isinstance(output, pa.Table)
-        return ArrowDataFrame(output, schema=schema)
+        adf: DataFrame = ArrowDataFrame(output)
+        if schema is not None:
+            _schema = Schema(schema)
+            if adf.schema != _schema:
+                adf = adf[_schema.names].alter_columns(_schema)
+        return adf
 
     def count(self, df: Any) -> int:  # pragma: no cover
         return df.count()
@@ -409,13 +438,15 @@ class _IterableArrowParam(LocalDataFrameParam):
     ) -> DataFrame:
         def dfs():
             _schema: Optional[Schema] = None if schema is None else Schema(schema)
+            has_return = False
             for df in output:
-                adf = ArrowDataFrame(df)
-                if _schema is not None and not (  # pylint: disable-all
-                    adf.schema == schema
-                ):
+                adf: DataFrame = ArrowDataFrame(df)
+                if _schema is not None and adf.schema != _schema:
                     adf = adf[_schema.names].alter_columns(_schema)
                 yield adf
+                has_return = True
+            if not has_return and _schema is not None:
+                yield ArrowDataFrame(schema=_schema)
 
         return IterableArrowDataFrame(dfs())
 

{fugue-0.8.6.dev1 → fugue-0.8.6.dev2}/fugue.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: fugue
-Version: 0.8.6.dev1
+Version: 0.8.6.dev2
 Summary: An abstraction layer for distributed computation
 Home-page: http://github.com/fugue-project/fugue
 Author: The Fugue Development Team

{fugue-0.8.6.dev1 → fugue-0.8.6.dev2}/fugue.egg-info/entry_points.txt

@@ -7,5 +7,5 @@ ibis = fugue_ibis [ibis]
 polars = fugue_polars.registry [polars]
 ray = fugue_ray.registry [ray]
 spark = fugue_spark.registry [spark]
-spark_ibis = fugue_spark.ibis_engine [spark,ibis]
+spark_ibis = fugue_spark.ibis_engine [ibis,spark]
 

{fugue-0.8.6.dev1 → fugue-0.8.6.dev2}/fugue.egg-info/requires.txt

@@ -26,7 +26,7 @@ dask[dataframe,distributed]
 ibis-framework>=2.1.1
 
 [all:python_version >= "3.8"]
-dask[dataframe,distributed]>=2022.9.0
+dask[dataframe,distributed]<2023.7.1,>=2022.9.0
 ibis-framework<6,>=3.2.0
 
 [cpp_sql_parser]
@@ -39,7 +39,7 @@ qpd[dask]>=0.4.4
 dask[dataframe,distributed]
 
 [dask:python_version >= "3.8"]
-dask[dataframe,distributed]>=2022.9.0
+dask[dataframe,distributed]<2023.7.1,>=2022.9.0
 
 [duckdb]
 duckdb>=0.5.0

{fugue-0.8.6.dev1 → fugue-0.8.6.dev2}/fugue_spark/execution_engine.py

@@ -134,9 +134,8 @@ class SparkMapEngine(MapEngine):
     ) -> DataFrame:
         output_schema = Schema(output_schema)
         if self._should_use_pandas_udf(output_schema):
-            # pandas udf can only be used for pyspark > 3
             if len(partition_spec.partition_by) > 0:
-                if partition_spec.algo == "coarse":
+                if partition_spec.algo in ["coarse", "even"]:
                     return self._map_by_pandas_udf(
                         df,
                         map_func=map_func,
@@ -145,7 +144,18 @@ class SparkMapEngine(MapEngine):
                         on_init=on_init,
                         map_func_format_hint=map_func_format_hint,
                     )
-                elif partition_spec.algo != "even" or self.is_spark_connect:
+                else:
+                    if (  # not simple partitioning
+                        partition_spec.algo != "hash"
+                        or partition_spec.num_partitions != "0"
+                    ):
+                        # TODO: not sure if presort should be done
+                        # on physical partition level
+                        df = self.to_df(
+                            self.execution_engine.repartition(
+                                df, PartitionSpec(partition_spec, presort=[])
+                            )
+                        )
                     return self._group_map_by_pandas_udf(
                         df,
                         map_func=map_func,
@@ -154,7 +164,7 @@ class SparkMapEngine(MapEngine):
                         on_init=on_init,
                         map_func_format_hint=map_func_format_hint,
                     )
-            elif len(partition_spec.partition_by) == 0:
+            else:
                 return self._map_by_pandas_udf(
                     df,
                     map_func=map_func,

{fugue-0.8.6.dev1 → fugue-0.8.6.dev2}/fugue_test/builtin_suite.py

@@ -22,6 +22,7 @@ from triad import SerializableRLock
 
 import fugue.api as fa
 from fugue import (
+    AnyDataFrame,
     ArrayDataFrame,
     CoTransformer,
     DataFrame,
@@ -365,6 +366,12 @@ class BuiltInTests(object):
                 dag.output(dict(df=a), using=mock_outputter2)
                 a.partition(num=3).output(MockOutputter3)
                 dag.output(dict(aa=a, bb=b), using=MockOutputter4)
+
+                a = dag.create(mock_creator2, params=dict(p=2))
+                b = dag.create(mock_creator2, params=dict(p=2))
+                c = dag.process(a, b, using=mock_processor4)
+                c.assert_eq(ArrayDataFrame([[2]], "a:int"))
+                dag.output(a, b, using=mock_outputter4)
             dag.run(self.engine)
 
         def test_zip(self):
@@ -435,9 +442,14 @@ class BuiltInTests(object):
             # this test is important for using mapInPandas in spark
 
             # schema: *,c:int
-            def mt_pandas(dfs: Iterable[pd.DataFrame]) -> Iterator[pd.DataFrame]:
+            def mt_pandas(
+                dfs: Iterable[pd.DataFrame], empty: bool = False
+            ) -> Iterator[pd.DataFrame]:
                 for df in dfs:
-                    yield df.assign(c=2)
+                    if not empty:
+                        df = df.assign(c=2)
+                        df = df[reversed(list(df.columns))]
+                        yield df
 
             with FugueWorkflow() as dag:
                 a = dag.df([[1, 2], [3, 4]], "a:int,b:int")
@@ -445,10 +457,25 @@ class BuiltInTests(object):
                 dag.df([[1, 2, 2], [3, 4, 2]], "a:int,b:int,c:int").assert_eq(b)
             dag.run(self.engine)
 
+            # when iterable returns nothing
+            with FugueWorkflow() as dag:
+                a = dag.df([[1, 2], [3, 4]], "a:int,b:int")
+                # without partitioning
+                b = a.transform(mt_pandas, params=dict(empty=True))
+                dag.df([], "a:int,b:int,c:int").assert_eq(b)
+                # with partitioning
+                b = a.partition_by("a").transform(mt_pandas, params=dict(empty=True))
+                dag.df([], "a:int,b:int,c:int").assert_eq(b)
+            dag.run(self.engine)
+
             # schema: *
-            def mt_arrow(dfs: Iterable[pa.Table]) -> Iterator[pa.Table]:
+            def mt_arrow(
+                dfs: Iterable[pa.Table], empty: bool = False
+            ) -> Iterator[pa.Table]:
                 for df in dfs:
-                    yield df
+                    if not empty:
+                        df = df.select(reversed(df.schema.names))
+                        yield df
 
             # schema: a:long
             def mt_arrow_2(dfs: Iterable[pa.Table]) -> Iterator[pa.Table]:
@@ -463,6 +490,17 @@ class BuiltInTests(object):
                 dag.df([[1], [3]], "a:long").assert_eq(b)
             dag.run(self.engine)
 
+            # when iterable returns nothing
+            with FugueWorkflow() as dag:
+                a = dag.df([[1, 2], [3, 4]], "a:int,b:int")
+                # without partitioning
+                b = a.transform(mt_arrow, params=dict(empty=True))
+                dag.df([], "a:int,b:int").assert_eq(b)
+                # with partitioning
+                b = a.partition_by("a").transform(mt_arrow, params=dict(empty=True))
+                dag.df([], "a:int,b:int").assert_eq(b)
+            dag.run(self.engine)
+
         def test_transform_binary(self):
             with FugueWorkflow() as dag:
                 a = dag.df([[1, pickle.dumps([0, "a"])]], "a:int,b:bytes")
@@ -1829,6 +1867,10 @@ def mock_creator(p: int) -> DataFrame:
     return ArrayDataFrame([[p]], "a:int")
 
 
+def mock_creator2(p: int) -> AnyDataFrame:
+    return fa.as_fugue_df([[p]], schema="a:int")
+
+
 def mock_processor(df1: List[List[Any]], df2: List[List[Any]]) -> DataFrame:
     return ArrayDataFrame([[len(df1) + len(df2)]], "a:int")
 
@@ -1844,6 +1886,10 @@ class MockProcessor3(Processor):
         return ArrayDataFrame([[sum(s.count() for s in dfs.values())]], "a:int")
 
 
+def mock_processor4(df1: AnyDataFrame, df2: AnyDataFrame) -> AnyDataFrame:
+    return ArrayDataFrame([[fa.count(df1) + fa.count(df2)]], "a:int")
+
+
 def mock_outputter(df1: List[List[Any]], df2: List[List[Any]]) -> None:
     assert len(df1) == len(df2)
 
@@ -1857,6 +1903,10 @@ class MockOutputter3(Outputter):
         assert "3" == self.partition_spec.num_partitions
 
 
+def mock_outputter4(df1: AnyDataFrame, df2: AnyDataFrame) -> None:
+    assert fa.count(df1) == fa.count(df2)
+
+
 class MockOutputter4(Outputter):
     def process(self, dfs):
         for k, v in dfs.items():
@@ -1895,8 +1945,8 @@ def mock_tf0(df: pd.DataFrame, p=1, col="p") -> pd.DataFrame:
 
 # schema: *,ct:int,p:int
 def mock_tf1(df: pd.DataFrame, p=1) -> pd.DataFrame:
-    df["ct"] = df.shape[0]
     df["p"] = p
+    df["ct"] = df.shape[0]
     return df
 
 

{fugue-0.8.6.dev1 → fugue-0.8.6.dev2}/setup.py

@@ -52,7 +52,7 @@ setup(
         "spark": ["pyspark>=3.1.1"],
         "dask": [
             "dask[distributed,dataframe]; python_version < '3.8'",
-            "dask[distributed,dataframe]>=2022.9.0; python_version >= '3.8'",
+            "dask[distributed,dataframe]>=2022.9.0,<2023.7.1; python_version >= '3.8'",
             "qpd[dask]>=0.4.4",
         ],
         "ray": ["ray[data]>=2.1.0", "duckdb>=0.5.0", "pyarrow>=6.0.1"],
@@ -73,7 +73,7 @@ setup(
             "fugue-sql-antlr[cpp]>=0.1.6",
             "pyspark>=3.1.1",
             "dask[distributed,dataframe]; python_version < '3.8'",
-            "dask[distributed,dataframe]>=2022.9.0; python_version >= '3.8'",
+            "dask[distributed,dataframe]>=2022.9.0,<2023.7.1; python_version >= '3.8'",
             "ray[data]>=2.1.0",
             "qpd[dask]>=0.4.4",
             "notebook",