fugue 0.8.6.dev1__py3-none-any.whl → 0.8.6.dev2__py3-none-any.whl

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/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/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.dist-info → fugue-0.8.6.dev2.dist-info}/METADATA

@@ -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
@@ -43,14 +43,14 @@ Requires-Dist: pyarrow (>=6.0.1) ; extra == 'all'
 Requires-Dist: polars ; extra == 'all'
 Requires-Dist: dask[dataframe,distributed] ; (python_version < "3.8") and extra == 'all'
 Requires-Dist: ibis-framework (>=2.1.1) ; (python_version < "3.8") and extra == 'all'
-Requires-Dist: dask[dataframe,distributed] (>=2022.9.0) ; (python_version >= "3.8") and extra == 'all'
+Requires-Dist: dask[dataframe,distributed] (<2023.7.1,>=2022.9.0) ; (python_version >= "3.8") and extra == 'all'
 Requires-Dist: ibis-framework (<6,>=3.2.0) ; (python_version >= "3.8") and extra == 'all'
 Provides-Extra: cpp_sql_parser
 Requires-Dist: fugue-sql-antlr[cpp] (>=0.1.6) ; extra == 'cpp_sql_parser'
 Provides-Extra: dask
 Requires-Dist: qpd[dask] (>=0.4.4) ; extra == 'dask'
 Requires-Dist: dask[dataframe,distributed] ; (python_version < "3.8") and extra == 'dask'
-Requires-Dist: dask[dataframe,distributed] (>=2022.9.0) ; (python_version >= "3.8") and extra == 'dask'
+Requires-Dist: dask[dataframe,distributed] (<2023.7.1,>=2022.9.0) ; (python_version >= "3.8") and extra == 'dask'
 Provides-Extra: duckdb
 Requires-Dist: duckdb (>=0.5.0) ; extra == 'duckdb'
 Requires-Dist: pyarrow (>=6.0.1) ; extra == 'duckdb'

{fugue-0.8.6.dev1.dist-info → fugue-0.8.6.dev2.dist-info}/RECORD

@@ -28,10 +28,10 @@ fugue/dataframe/__init__.py,sha256=zm7TbsaJLIvfm7zymWm2LGcuJd3nxfGsFnQiyrSnenM,6
 fugue/dataframe/api.py,sha256=c5Err3c-ayl-k28IUi6kV_ClDWX30NpVNkv97hQKDac,9862
 fugue/dataframe/array_dataframe.py,sha256=oBfN545NTGdYJ5zPIRv7hXRR-R_OW1JieyOfnl296oU,4447
 fugue/dataframe/arrow_dataframe.py,sha256=h0DJH8G0MEgfkyvX4U957iqDXIgvTtrP7YED5iEjizI,12098
-fugue/dataframe/dataframe.py,sha256=rIjaOplyRDt_BfpFwZWMjZALDqa03NePs9tFiLSW6Jg,17247
-fugue/dataframe/dataframe_iterable_dataframe.py,sha256=0gvb12D0s4VntNZ-M8J_Pic2XlFpw7upXf4hfMn2ufY,7255
+fugue/dataframe/dataframe.py,sha256=XOn9x_aKWOLdpach7RHxg2PoP-hpfsyjoRFvLx6xKV0,17496
+fugue/dataframe/dataframe_iterable_dataframe.py,sha256=lx71KfaI4lsVKI-79buc-idaeT20JEMBOq21SQcAiY8,7259
 fugue/dataframe/dataframes.py,sha256=tBSpHsENgbcdOJ0Jgst6PTKbjG7_uoFJch96oTlaQIs,4160
-fugue/dataframe/function_wrapper.py,sha256=DjyIrNBj5Bv9AwIl2I2fG5ClcKe3OoW_eBkyEabYY5Y,13505
+fugue/dataframe/function_wrapper.py,sha256=r6H1SQWaag2eSbJ50327t_bt7MZunbOMOl9OcOcQW2E,14827
 fugue/dataframe/iterable_dataframe.py,sha256=Kn5HZnVU4o1nn9mbbQxaV8rGG869wImZcOCK3AdlA-M,4627
 fugue/dataframe/pandas_dataframe.py,sha256=ZWqI-ZUFiSP7giJ3siRlrZcMedI_fyuoLn227H0YRvw,10453
 fugue/dataframe/utils.py,sha256=nQVU01jspB1NSeRiagE71uzRibDqvyGwi94ZfHwNHD0,10508
@@ -131,7 +131,7 @@ fugue_ray/_utils/io.py,sha256=gHfx70tdXPHmVL6nHxVhmCO5KpWjjyDG8qKT1Lbpav4,8737
 fugue_spark/__init__.py,sha256=rvrMpFs9socMgyH_58gLbnAqmirBf5oidXoO4cekW6U,165
 fugue_spark/_constants.py,sha256=K2uLQfjvMxXk75K-7_Wn47Alpwq5rW57BtECAUrOeqA,177
 fugue_spark/dataframe.py,sha256=xoM2-SwVRFfSyfEEnx4g4b0GO6XgN-DQLvXtUrAMq1Q,9510
-fugue_spark/execution_engine.py,sha256=n_ZdQI4ysbgdDsjyWwbwhyBsgcg_icc9MAmVyWEn2RQ,32029
+fugue_spark/execution_engine.py,sha256=nssgfqt2h1OjGlK5iuGEyF-lt2pEokmsjrqo6K4C1Kg,32440
 fugue_spark/ibis_engine.py,sha256=Yl5xxwROo1idcD2hFaylaI1IpmBUgbvOZRWtcrE0Zjo,1697
 fugue_spark/registry.py,sha256=kyIMk6dAiKRSKCHawQKyXu9DhZ24T6j3gL57TiOAZ8c,4162
 fugue_spark/_utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -143,14 +143,14 @@ fugue_sql/__init__.py,sha256=Cmr7w0Efr7PzoXdQzdJfc4Dgqd69qKqcHZZodENq7EU,287
 fugue_sql/exceptions.py,sha256=ltS0MC8gMnVVrJbQiOZ0kRUWvVQ2LTx33dCW3ugqtb0,260
 fugue_test/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 fugue_test/bag_suite.py,sha256=WbDCFjuAHYoJh4GXSPiSJxOoOwE1VMtYpJ3lQrsUK-Y,2483
-fugue_test/builtin_suite.py,sha256=piY_bVtLylQWrZbencrQwT89IXgvgBmNscddJ6vZieY,76338
+fugue_test/builtin_suite.py,sha256=sT_Btm7cpSsuLnt_PpNFeLZVc2WPmF5NQy4GKVri_-c,78396
 fugue_test/dataframe_suite.py,sha256=6uM7_et2Y55-ePIssG9G_K9mXBYNjPXnpunuCh0xKhw,19082
 fugue_test/execution_suite.py,sha256=HzM_7jUhnt1kD2jGaaPhyr-Q_vtRV4nCQTaLtMnzTuU,50948
 fugue_test/ibis_suite.py,sha256=Dk4AHVD00RcFsNm9VvJ4_4LOyFdGX30OnAtpO2SPruE,3529
 fugue_version/__init__.py,sha256=VpASnrti7EGWxUfSWGgERUfe7NLJltfVXYosOzHbpPg,22
-fugue-0.8.6.dev1.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-fugue-0.8.6.dev1.dist-info/METADATA,sha256=OuDeHU_GR0fSI0jISu0fczkrLCmUgnmr8IYPcEjNbz8,18264
-fugue-0.8.6.dev1.dist-info/WHEEL,sha256=pkctZYzUS4AYVn6dJ-7367OJZivF2e8RA9b_ZBjif18,92
-fugue-0.8.6.dev1.dist-info/entry_points.txt,sha256=xD0_Lj3jz-np_UNVAwsKt9kYWX0CwM7o7IvkCXRmcQY,374
-fugue-0.8.6.dev1.dist-info/top_level.txt,sha256=y1eCfzGdQ1_RkgcShcfbvXs-bopD3DwJcIOxP9EFXno,140
-fugue-0.8.6.dev1.dist-info/RECORD,,
+fugue-0.8.6.dev2.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+fugue-0.8.6.dev2.dist-info/METADATA,sha256=pGSp1pOq2BHvHd3KiRL3YzFHbNKzmKDMUPd8Hcs__xc,18284
+fugue-0.8.6.dev2.dist-info/WHEEL,sha256=AtBG6SXL3KF_v0NxLf0ehyVOh0cold-JbJYXNGorC6Q,92
+fugue-0.8.6.dev2.dist-info/entry_points.txt,sha256=Ta1DD9RIgS_YfhieUvM6PgAzuOYmhlnKe9fWbns9sLc,374
+fugue-0.8.6.dev2.dist-info/top_level.txt,sha256=y1eCfzGdQ1_RkgcShcfbvXs-bopD3DwJcIOxP9EFXno,140
+fugue-0.8.6.dev2.dist-info/RECORD,,

{fugue-0.8.6.dev1.dist-info → fugue-0.8.6.dev2.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.40.0)
+Generator: bdist_wheel (0.41.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

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

@@ -2,10 +2,10 @@
 dask = fugue_dask.registry [dask]
 dask_ibis = fugue_dask.ibis_engine [dask,ibis]
 duckdb = fugue_duckdb.registry [duckdb]
-duckdb_ibis = fugue_duckdb.ibis_engine [ibis,duckdb]
+duckdb_ibis = fugue_duckdb.ibis_engine [duckdb,ibis]
 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_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_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