fugue 0.9.7.dev2__tar.gz → 0.9.9.dev0__tar.gz

{fugue-0.9.7.dev2 → fugue-0.9.9.dev0}/PKG-INFO

@@ -1,28 +1,30 @@
 Metadata-Version: 2.4
 Name: fugue
-Version: 0.9.7.dev2
+Version: 0.9.9.dev0
 Summary: An abstraction layer for distributed computing
 Author-email: The Fugue Development Team <hello@fugue.ai>
 License: Apache-2.0
 Project-URL: Homepage, http://github.com/fugue-project/fugue
 Project-URL: Repository, http://github.com/fugue-project/fugue
-Keywords: distributed,spark,dask,ray,sql,dsl,domain specific language
+Keywords: distributed,spark,dask,ray,duckdb,sql,dsl,domain specific language
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Developers
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
+Classifier: Programming Language :: Python :: 3.15
 Classifier: Programming Language :: Python :: 3 :: Only
-Requires-Python: >=3.10
+Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: triad>=1.0.1
+Requires-Dist: triad>=1.0.4
 Requires-Dist: adagio>=0.2.6
-Requires-Dist: pandas<3
+Requires-Dist: pandas; python_version < "3.15"
+Requires-Dist: pandas>=3; python_version >= "3.15"
+Requires-Dist: packaging
 Provides-Extra: sql
 Requires-Dist: duckdb>=0.5.0; extra == "sql"
 Requires-Dist: fugue-sql-antlr>=0.2.4; extra == "sql"
@@ -31,10 +33,11 @@ Requires-Dist: jinja2; extra == "sql"
 Provides-Extra: cpp-sql-parser
 Requires-Dist: fugue-sql-antlr[cpp]>=0.2.4; extra == "cpp-sql-parser"
 Provides-Extra: spark
-Requires-Dist: pyspark>=3.1.1; extra == "spark"
+Requires-Dist: pyspark>=3.5; python_version < "3.12" and extra == "spark"
+Requires-Dist: pyspark>=4.0; python_version >= "3.12" and extra == "spark"
 Requires-Dist: zstandard>=0.25.0; extra == "spark"
 Provides-Extra: dask
-Requires-Dist: dask[dataframe,distributed]>=2024.4.0; extra == "dask"
+Requires-Dist: dask[dataframe,distributed]>=2025.12.0; extra == "dask"
 Requires-Dist: pyarrow>=7.0.0; extra == "dask"
 Requires-Dist: pandas>=2.0.2; extra == "dask"
 Provides-Extra: ray
@@ -54,7 +57,7 @@ Provides-Extra: ibis
 Requires-Dist: fugue-sql-antlr>=0.2.4; extra == "ibis"
 Requires-Dist: sqlglot<28; extra == "ibis"
 Requires-Dist: jinja2; extra == "ibis"
-Requires-Dist: ibis-framework[pandas]; extra == "ibis"
+Requires-Dist: ibis-framework; extra == "ibis"
 Provides-Extra: notebook
 Requires-Dist: notebook; extra == "notebook"
 Requires-Dist: jupyterlab; extra == "notebook"
@@ -63,8 +66,9 @@ Provides-Extra: all
 Requires-Dist: fugue-sql-antlr>=0.2.4; extra == "all"
 Requires-Dist: sqlglot<28; extra == "all"
 Requires-Dist: jinja2; extra == "all"
-Requires-Dist: pyspark>=3.1.1; extra == "all"
-Requires-Dist: dask[dataframe,distributed]>=2024.4.0; extra == "all"
+Requires-Dist: pyspark>=3.5; python_version < "3.12" and extra == "all"
+Requires-Dist: pyspark>=4.0; python_version >= "3.12" and extra == "all"
+Requires-Dist: dask[dataframe,distributed]>=2025.12.0; extra == "all"
 Requires-Dist: dask-sql; extra == "all"
 Requires-Dist: ray[data]>=2.30.0; python_version < "3.14" and extra == "all"
 Requires-Dist: notebook; extra == "all"
@@ -73,7 +77,7 @@ Requires-Dist: ipython>=7.10.0; extra == "all"
 Requires-Dist: duckdb>=0.5.0; extra == "all"
 Requires-Dist: pyarrow>=6.0.1; extra == "all"
 Requires-Dist: pandas>=2.0.2; extra == "all"
-Requires-Dist: ibis-framework[duckdb,pandas]; extra == "all"
+Requires-Dist: ibis-framework[duckdb]; extra == "all"
 Requires-Dist: polars; extra == "all"
 Dynamic: license-file
 
@@ -82,7 +86,7 @@ Dynamic: license-file
 [![PyPI version](https://badge.fury.io/py/fugue.svg)](https://pypi.python.org/pypi/fugue/)
 [![PyPI pyversions](https://img.shields.io/pypi/pyversions/fugue.svg)](https://pypi.python.org/pypi/fugue/)
 [![PyPI license](https://img.shields.io/pypi/l/fugue.svg)](https://pypi.python.org/pypi/fugue/)
-[![codecov](https://codecov.io/gh/fugue-project/fugue/branch/main/graph/badge.svg?token=ZO9YD5N3IA)](https://codecov.io/gh/fugue-project/fugue)
+[![codecov](https://codecov.io/gh/fugue-project/fugue/graph/badge.svg?token=ZO9YD5N3IA)](https://codecov.io/gh/fugue-project/fugue)
 [![Codacy Badge](https://app.codacy.com/project/badge/Grade/4fa5f2f53e6f48aaa1218a89f4808b91)](https://www.codacy.com/gh/fugue-project/fugue/dashboard?utm_source=github.com&utm_medium=referral&utm_content=fugue-project/fugue&utm_campaign=Badge_Grade)
 [![Downloads](https://static.pepy.tech/badge/fugue)](https://pepy.tech/project/fugue)
 

{fugue-0.9.7.dev2 → fugue-0.9.9.dev0}/README.md

@@ -3,7 +3,7 @@
 [![PyPI version](https://badge.fury.io/py/fugue.svg)](https://pypi.python.org/pypi/fugue/)
 [![PyPI pyversions](https://img.shields.io/pypi/pyversions/fugue.svg)](https://pypi.python.org/pypi/fugue/)
 [![PyPI license](https://img.shields.io/pypi/l/fugue.svg)](https://pypi.python.org/pypi/fugue/)
-[![codecov](https://codecov.io/gh/fugue-project/fugue/branch/main/graph/badge.svg?token=ZO9YD5N3IA)](https://codecov.io/gh/fugue-project/fugue)
+[![codecov](https://codecov.io/gh/fugue-project/fugue/graph/badge.svg?token=ZO9YD5N3IA)](https://codecov.io/gh/fugue-project/fugue)
 [![Codacy Badge](https://app.codacy.com/project/badge/Grade/4fa5f2f53e6f48aaa1218a89f4808b91)](https://www.codacy.com/gh/fugue-project/fugue/dashboard?utm_source=github.com&utm_medium=referral&utm_content=fugue-project/fugue&utm_campaign=Badge_Grade)
 [![Downloads](https://static.pepy.tech/badge/fugue)](https://pepy.tech/project/fugue)
 

{fugue-0.9.7.dev2 → fugue-0.9.9.dev0}/fugue/dataframe/arrow_dataframe.py

@@ -5,7 +5,12 @@ import pyarrow as pa
 from triad.collections.schema import Schema
 from triad.exceptions import InvalidOperationError
 from triad.utils.assertion import assert_or_throw
-from triad.utils.pyarrow import cast_pa_table, pa_table_to_pandas
+from triad.utils.pyarrow import (
+    cast_pa_table,
+    normalize_large_types,
+    normalize_large_types_in_table,
+    pa_table_to_pandas,
+)
 
 from fugue.dataset.api import (
     as_fugue_dataset,
@@ -73,6 +78,7 @@ class ArrowDataFrame(LocalBoundedDataFrame):
                 schema is None,
                 InvalidOperationError("can't reset schema for pa.Table"),
             )
+            df = normalize_large_types_in_table(df)
             self._native = df
             super().__init__(Schema(df.schema))
             return
@@ -253,6 +259,7 @@ def _pa_table_as_dict_iterable(
 def _pa_table_alter_columns(
     df: pa.Table, columns: Any, as_fugue: bool = False
 ) -> pa.Table:
+    df = normalize_large_types_in_table(df)
     schema = Schema(df.schema)
     new_schema = schema.alter(columns)
     if schema != new_schema:
@@ -302,7 +309,7 @@ def _get_pyarrow_table_columns(df: pa.Table) -> List[Any]:
 
 @get_schema.candidate(lambda df: isinstance(df, pa.Table))
 def _get_pyarrow_table_schema(df: pa.Table) -> Schema:
-    return Schema(df.schema)
+    return Schema(normalize_large_types(df.schema))
 
 
 @rename.candidate(lambda df, *args, **kwargs: isinstance(df, pa.Table))

{fugue-0.9.7.dev2 → fugue-0.9.9.dev0}/fugue/dataframe/pandas_dataframe.py

@@ -1,3 +1,4 @@
+import math
 from typing import Any, Dict, Iterable, List, Optional, Tuple
 
 import pandas as pd
@@ -109,7 +110,17 @@ class PandasDataFrame(LocalBoundedDataFrame):
 
     def peek_array(self) -> List[Any]:
         self.assert_not_empty()
-        return self.native.iloc[0].values.tolist()
+        arr = self.native.iloc[0].values.tolist()
+        # missing non-float values may come back as nan / pd.NA
+        return [
+            (
+                None
+                if (v is pd.NA or (isinstance(v, float) and math.isnan(v)))
+                and not pa.types.is_floating(t)
+                else v
+            )
+            for v, t in zip(arr, self.schema.types)
+        ]
 
     def count(self) -> int:
         return self.native.shape[0]

{fugue-0.9.7.dev2 → fugue-0.9.9.dev0}/fugue/execution/native_execution_engine.py

@@ -2,7 +2,6 @@ import logging
 import os
 from typing import Any, Callable, Dict, List, Optional, Type, Union
 
-import numpy as np
 import pandas as pd
 from triad import Schema
 from triad.collections.dict import IndexedOrderedDict
@@ -149,7 +148,7 @@ class PandasMapEngine(MapEngine):
                 )
                 dfs: List[DataFrame] = []
                 for p, subdf in enumerate(
-                    np.array_split(input_df.as_pandas(), partitions)
+                    _split_pandas(input_df.as_pandas(), partitions)
                 ):
                     if len(subdf) > 0:
                         tdf = PandasDataFrame(subdf, df.schema, pandas_df_wrapper=True)
@@ -443,3 +442,15 @@ class _NativeExecutionEngineParam(ExecutionEngineParam):
 def _to_native_execution_engine_df(df: AnyDataFrame, schema: Any = None) -> DataFrame:
     fdf = as_fugue_df(df) if schema is None else as_fugue_df(df, schema=schema)
     return fdf.as_local_bounded()
+
+
+def _split_pandas(df: pd.DataFrame, n: int) -> List[pd.DataFrame]:
+    # same chunk sizes as np.array_split, which returns ndarrays on pandas 3
+    size, extra = divmod(len(df), n)
+    res: List[pd.DataFrame] = []
+    start = 0
+    for i in range(n):
+        end = start + size + (1 if i < extra else 0)
+        res.append(df.iloc[start:end].reset_index(drop=True))
+        start = end
+    return res

{fugue-0.9.7.dev2 → fugue-0.9.9.dev0}/fugue.egg-info/PKG-INFO

@@ -1,28 +1,30 @@
 Metadata-Version: 2.4
 Name: fugue
-Version: 0.9.7.dev2
+Version: 0.9.9.dev0
 Summary: An abstraction layer for distributed computing
 Author-email: The Fugue Development Team <hello@fugue.ai>
 License: Apache-2.0
 Project-URL: Homepage, http://github.com/fugue-project/fugue
 Project-URL: Repository, http://github.com/fugue-project/fugue
-Keywords: distributed,spark,dask,ray,sql,dsl,domain specific language
+Keywords: distributed,spark,dask,ray,duckdb,sql,dsl,domain specific language
 Classifier: Development Status :: 5 - Production/Stable
 Classifier: Intended Audience :: Developers
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Classifier: Programming Language :: Python :: 3
-Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
+Classifier: Programming Language :: Python :: 3.15
 Classifier: Programming Language :: Python :: 3 :: Only
-Requires-Python: >=3.10
+Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: triad>=1.0.1
+Requires-Dist: triad>=1.0.4
 Requires-Dist: adagio>=0.2.6
-Requires-Dist: pandas<3
+Requires-Dist: pandas; python_version < "3.15"
+Requires-Dist: pandas>=3; python_version >= "3.15"
+Requires-Dist: packaging
 Provides-Extra: sql
 Requires-Dist: duckdb>=0.5.0; extra == "sql"
 Requires-Dist: fugue-sql-antlr>=0.2.4; extra == "sql"
@@ -31,10 +33,11 @@ Requires-Dist: jinja2; extra == "sql"
 Provides-Extra: cpp-sql-parser
 Requires-Dist: fugue-sql-antlr[cpp]>=0.2.4; extra == "cpp-sql-parser"
 Provides-Extra: spark
-Requires-Dist: pyspark>=3.1.1; extra == "spark"
+Requires-Dist: pyspark>=3.5; python_version < "3.12" and extra == "spark"
+Requires-Dist: pyspark>=4.0; python_version >= "3.12" and extra == "spark"
 Requires-Dist: zstandard>=0.25.0; extra == "spark"
 Provides-Extra: dask
-Requires-Dist: dask[dataframe,distributed]>=2024.4.0; extra == "dask"
+Requires-Dist: dask[dataframe,distributed]>=2025.12.0; extra == "dask"
 Requires-Dist: pyarrow>=7.0.0; extra == "dask"
 Requires-Dist: pandas>=2.0.2; extra == "dask"
 Provides-Extra: ray
@@ -54,7 +57,7 @@ Provides-Extra: ibis
 Requires-Dist: fugue-sql-antlr>=0.2.4; extra == "ibis"
 Requires-Dist: sqlglot<28; extra == "ibis"
 Requires-Dist: jinja2; extra == "ibis"
-Requires-Dist: ibis-framework[pandas]; extra == "ibis"
+Requires-Dist: ibis-framework; extra == "ibis"
 Provides-Extra: notebook
 Requires-Dist: notebook; extra == "notebook"
 Requires-Dist: jupyterlab; extra == "notebook"
@@ -63,8 +66,9 @@ Provides-Extra: all
 Requires-Dist: fugue-sql-antlr>=0.2.4; extra == "all"
 Requires-Dist: sqlglot<28; extra == "all"
 Requires-Dist: jinja2; extra == "all"
-Requires-Dist: pyspark>=3.1.1; extra == "all"
-Requires-Dist: dask[dataframe,distributed]>=2024.4.0; extra == "all"
+Requires-Dist: pyspark>=3.5; python_version < "3.12" and extra == "all"
+Requires-Dist: pyspark>=4.0; python_version >= "3.12" and extra == "all"
+Requires-Dist: dask[dataframe,distributed]>=2025.12.0; extra == "all"
 Requires-Dist: dask-sql; extra == "all"
 Requires-Dist: ray[data]>=2.30.0; python_version < "3.14" and extra == "all"
 Requires-Dist: notebook; extra == "all"
@@ -73,7 +77,7 @@ Requires-Dist: ipython>=7.10.0; extra == "all"
 Requires-Dist: duckdb>=0.5.0; extra == "all"
 Requires-Dist: pyarrow>=6.0.1; extra == "all"
 Requires-Dist: pandas>=2.0.2; extra == "all"
-Requires-Dist: ibis-framework[duckdb,pandas]; extra == "all"
+Requires-Dist: ibis-framework[duckdb]; extra == "all"
 Requires-Dist: polars; extra == "all"
 Dynamic: license-file
 
@@ -82,7 +86,7 @@ Dynamic: license-file
 [![PyPI version](https://badge.fury.io/py/fugue.svg)](https://pypi.python.org/pypi/fugue/)
 [![PyPI pyversions](https://img.shields.io/pypi/pyversions/fugue.svg)](https://pypi.python.org/pypi/fugue/)
 [![PyPI license](https://img.shields.io/pypi/l/fugue.svg)](https://pypi.python.org/pypi/fugue/)
-[![codecov](https://codecov.io/gh/fugue-project/fugue/branch/main/graph/badge.svg?token=ZO9YD5N3IA)](https://codecov.io/gh/fugue-project/fugue)
+[![codecov](https://codecov.io/gh/fugue-project/fugue/graph/badge.svg?token=ZO9YD5N3IA)](https://codecov.io/gh/fugue-project/fugue)
 [![Codacy Badge](https://app.codacy.com/project/badge/Grade/4fa5f2f53e6f48aaa1218a89f4808b91)](https://www.codacy.com/gh/fugue-project/fugue/dashboard?utm_source=github.com&utm_medium=referral&utm_content=fugue-project/fugue&utm_campaign=Badge_Grade)
 [![Downloads](https://static.pepy.tech/badge/fugue)](https://pepy.tech/project/fugue)
 

{fugue-0.9.7.dev2 → fugue-0.9.9.dev0}/fugue.egg-info/requires.txt

@@ -1,13 +1,18 @@
-triad>=1.0.1
+triad>=1.0.4
 adagio>=0.2.6
-pandas<3
+packaging
+
+[:python_version < "3.15"]
+pandas
+
+[:python_version >= "3.15"]
+pandas>=3
 
 [all]
 fugue-sql-antlr>=0.2.4
 sqlglot<28
 jinja2
-pyspark>=3.1.1
-dask[dataframe,distributed]>=2024.4.0
+dask[dataframe,distributed]>=2025.12.0
 dask-sql
 notebook
 jupyterlab
@@ -15,17 +20,23 @@ ipython>=7.10.0
 duckdb>=0.5.0
 pyarrow>=6.0.1
 pandas>=2.0.2
-ibis-framework[duckdb,pandas]
+ibis-framework[duckdb]
 polars
 
+[all:python_version < "3.12"]
+pyspark>=3.5
+
 [all:python_version < "3.14"]
 ray[data]>=2.30.0
 
+[all:python_version >= "3.12"]
+pyspark>=4.0
+
 [cpp_sql_parser]
 fugue-sql-antlr[cpp]>=0.2.4
 
 [dask]
-dask[dataframe,distributed]>=2024.4.0
+dask[dataframe,distributed]>=2025.12.0
 pyarrow>=7.0.0
 pandas>=2.0.2
 
@@ -40,7 +51,7 @@ numpy
 fugue-sql-antlr>=0.2.4
 sqlglot<28
 jinja2
-ibis-framework[pandas]
+ibis-framework
 
 [notebook]
 notebook
@@ -59,9 +70,14 @@ pandas
 ray[data]>=2.30.0
 
 [spark]
-pyspark>=3.1.1
 zstandard>=0.25.0
 
+[spark:python_version < "3.12"]
+pyspark>=3.5
+
+[spark:python_version >= "3.12"]
+pyspark>=4.0
+
 [sql]
 duckdb>=0.5.0
 fugue-sql-antlr>=0.2.4

{fugue-0.9.7.dev2 → fugue-0.9.9.dev0}/fugue_dask/execution_engine.py

@@ -169,7 +169,7 @@ class DaskMapEngine(MapEngine):
                     sort=False,
                     group_keys=False,
                     dropna=False,
-                ).apply(_map, meta=output_dtypes)
+                )[list(df.native.columns)].apply(_map, meta=output_dtypes)
             elif partition_spec.algo == "coarse":
                 result = pdf.native.map_partitions(  # type: ignore
                     _map, meta=output_dtypes
@@ -515,7 +515,7 @@ class DaskExecutionEngine(ExecutionEngine):
                 )
 
             d = (
-                d.groupby(partition_spec.partition_by, dropna=False)
+                d.groupby(partition_spec.partition_by, dropna=False)[list(d.columns)]
                 .apply(_partition_take, n=n, presort=_presort, meta=meta)
                 .reset_index(drop=True)
             )

{fugue-0.9.7.dev2 → fugue-0.9.9.dev0}/fugue_duckdb/dataframe.py

@@ -4,7 +4,7 @@ import pandas as pd
 import pyarrow as pa
 from duckdb import DuckDBPyRelation
 from triad import Schema, assert_or_throw
-from triad.utils.pyarrow import LARGE_TYPES_REPLACEMENT, replace_types_in_table
+from triad.utils.pyarrow import normalize_large_types_in_table
 
 from fugue import ArrowDataFrame, DataFrame, LocalBoundedDataFrame
 from fugue.dataframe.arrow_dataframe import _pa_table_as_pandas
@@ -166,8 +166,7 @@ def _duck_as_local(df: DuckDBPyRelation) -> DuckDBPyRelation:
 @as_arrow.candidate(lambda df: isinstance(df, DuckDBPyRelation))
 def _duck_as_arrow(df: DuckDBPyRelation) -> pa.Table:
     _df = df.fetch_arrow_table()
-    _df = replace_types_in_table(_df, LARGE_TYPES_REPLACEMENT, recursive=True)
-    return _df
+    return normalize_large_types_in_table(_df)
 
 
 @as_pandas.candidate(lambda df: isinstance(df, DuckDBPyRelation))

{fugue-0.9.7.dev2 → fugue-0.9.9.dev0}/fugue_polars/polars_dataframe.py

@@ -6,11 +6,7 @@ import pyarrow as pa
 from triad.collections.schema import Schema
 from triad.exceptions import InvalidOperationError
 from triad.utils.assertion import assert_or_throw
-from triad.utils.pyarrow import (
-    LARGE_TYPES_REPLACEMENT,
-    replace_types_in_schema,
-    replace_types_in_table,
-)
+from triad.utils.pyarrow import normalize_large_types, normalize_large_types_in_table
 
 from fugue import ArrowDataFrame
 from fugue.api import (
@@ -157,8 +153,7 @@ def _pl_as_local_bounded(df: pl.DataFrame) -> pl.DataFrame:
 @as_arrow.candidate(lambda df: isinstance(df, pl.DataFrame))
 def _pl_as_arrow(df: pl.DataFrame) -> pa.Table:
     adf = df.to_arrow()
-    adf = replace_types_in_table(adf, LARGE_TYPES_REPLACEMENT)
-    return adf
+    return normalize_large_types_in_table(adf)
 
 
 @is_df.candidate(lambda df: isinstance(df, pl.DataFrame))
@@ -199,8 +194,7 @@ def _get_pl_columns(df: pl.DataFrame) -> List[Any]:
 @get_schema.candidate(lambda df: isinstance(df, pl.DataFrame))
 def _get_pl_schema(df: pl.DataFrame) -> Schema:
     adf = df.to_arrow()
-    schema = replace_types_in_schema(adf.schema, LARGE_TYPES_REPLACEMENT)
-    return Schema(schema)
+    return Schema(normalize_large_types(adf.schema))
 
 
 @rename.candidate(lambda df, *args, **kwargs: isinstance(df, pl.DataFrame))

{fugue-0.9.7.dev2 → fugue-0.9.9.dev0}/fugue_ray/_utils/dataframe.py

@@ -5,6 +5,7 @@ import pandas as pd
 import pyarrow as pa
 import ray.data as rd
 from triad import Schema
+from triad.utils.pyarrow import normalize_large_types
 
 from .._constants import _ZERO_COPY
 
@@ -40,10 +41,10 @@ def get_dataset_format(df: rd.Dataset) -> Tuple[Optional[str], rd.Dataset]:
 
 def to_schema(schema: Any) -> Schema:  # pragma: no cover
     if isinstance(schema, pa.Schema):
-        return Schema(schema)
+        return Schema(normalize_large_types(schema))
     if isinstance(schema, rd.Schema):
         if hasattr(schema, "base_schema") and isinstance(schema.base_schema, pa.Schema):
-            return Schema(schema.base_schema)
+            return Schema(normalize_large_types(schema.base_schema))
         return Schema(list(zip(schema.names, schema.types)))
     raise ValueError(f"{schema} is not supported")
 

{fugue-0.9.7.dev2 → fugue-0.9.9.dev0}/fugue_ray/_utils/io.py

@@ -1,17 +1,17 @@
 import os
 import pathlib
-from typing import Any, Callable, Dict, Iterable, List, Optional, Union
+from collections.abc import Callable, Iterable
+from typing import Any
 
 import pyarrow as pa
 import ray.data as rd
 from packaging import version
 from pyarrow import csv as pacsv
 from pyarrow import json as pajson
-
 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, isfile
+from triad.utils.io import exists, isfile, makedirs, rm
 
 from fugue import ExecutionEngine
 from fugue._utils.io import FileParser, save_df
@@ -25,9 +25,7 @@ try:
     from ray.data.datasource import FileExtensionFilter
 
     class _FileFiler(FileExtensionFilter):  # pragma: no cover
-        def __init__(
-            self, file_extensions: Union[str, List[str]], exclude: Iterable[str]
-        ):
+        def __init__(self, file_extensions: str | list[str], exclude: Iterable[str]):
             super().__init__(file_extensions, allow_if_no_extension=True)
             self._exclude = set(exclude)
 
@@ -36,23 +34,23 @@ try:
                 path
             ).name not in self._exclude and self._file_has_extension(path)
 
-        def __call__(self, paths: List[str]) -> List[str]:
+        def __call__(self, paths: list[str]) -> list[str]:
             return [path for path in paths if self._is_valid(path)]
 
 except ImportError:  # pragma: no cover
     pass  # ray >=2.10
 
 
-class RayIO(object):
+class RayIO:
     def __init__(self, engine: ExecutionEngine):
         self._engine = engine
         self._logger = engine.log
-        self._loads: Dict[str, Callable[..., DataFrame]] = {
+        self._loads: dict[str, Callable[..., DataFrame]] = {
             "csv": self._load_csv,
             "parquet": self._load_parquet,
             "json": self._load_json,
         }
-        self._saves: Dict[str, Callable[..., None]] = {
+        self._saves: dict[str, Callable[..., None]] = {
             "csv": self._save_csv,
             "parquet": self._save_parquet,
             "json": self._save_json,
@@ -60,8 +58,8 @@ class RayIO(object):
 
     def load_df(
         self,
-        uri: Union[str, List[str]],
-        format_hint: Optional[str] = None,
+        uri: str | list[str],
+        format_hint: str | None = None,
         columns: Any = None,
         **kwargs: Any,
     ) -> DataFrame:
@@ -81,8 +79,8 @@ class RayIO(object):
         self,
         df: RayDataFrame,
         uri: str,
-        format_hint: Optional[str] = None,
-        partition_spec: Optional[PartitionSpec] = None,
+        format_hint: str | None = None,
+        partition_spec: PartitionSpec | None = None,
         mode: str = "overwrite",
         force_single: bool = False,
         **kwargs: Any,
@@ -122,7 +120,7 @@ class RayIO(object):
         if "header" in kw:
             kw["include_header"] = kw.pop("header")
 
-        def _fn() -> Dict[str, Any]:  # pragma: no cover
+        def _fn() -> dict[str, Any]:  # pragma: no cover
             return dict(write_options=pacsv.WriteOptions(**kw))
 
         df.native.write_csv(
@@ -152,13 +150,25 @@ class RayIO(object):
         )
 
     def _load_parquet(
-        self, p: List[str], columns: Any = None, **kwargs: Any
+        self, p: list[str], columns: Any = None, **kwargs: Any
     ) -> DataFrame:
         # in 2.52.0 the default changes to ["parquet"]
         if "file_extensions" not in kwargs:
             kwargs = kwargs.copy()
             kwargs["file_extensions"] = None
-        sdf = rd.read_parquet(p, ray_remote_args=self._remote_args(), **kwargs)
+        # ray>=2.60 routes read_parquet to the V2 datasource by default, which
+        # (1) forces file_extensions=None to ["parquet"], so single files without
+        # extension can't be read, and (2) may raise ImportError from a threaded
+        # fsspec HTTPFileSystem import when listing multiple paths. Use V1.
+        ctx = rd.DataContext.get_current()
+        use_v2 = getattr(ctx, "use_datasource_v2", None)
+        if use_v2:  # pragma: no cover
+            ctx.use_datasource_v2 = False
+        try:
+            sdf = rd.read_parquet(p, ray_remote_args=self._remote_args(), **kwargs)
+        finally:
+            if use_v2:  # pragma: no cover
+                ctx.use_datasource_v2 = use_v2
         if columns is None:
             return RayDataFrame(sdf)
         if isinstance(columns, list):  # column names
@@ -166,20 +176,18 @@ class RayIO(object):
         schema = Schema(columns)
         return RayDataFrame(sdf)[schema.names].alter_columns(schema)
 
-    def _load_csv(  # noqa: C901
-        self, p: List[str], columns: Any = None, **kwargs: Any
-    ) -> DataFrame:
+    def _load_csv(self, p: list[str], columns: Any = None, **kwargs: Any) -> DataFrame:
         kw = ParamDict(kwargs)
         infer_schema = kw.get("infer_schema", False)
-        read_options: Dict[str, Any] = {"use_threads": False}
-        parse_options: Dict[str, Any] = {}
-        convert_options: Dict[str, Any] = {}
+        read_options: dict[str, Any] = {"use_threads": False}
+        parse_options: dict[str, Any] = {}
+        convert_options: dict[str, Any] = {}
         if infer_schema and columns is not None and not isinstance(columns, list):
             raise ValueError("can't set columns as a schema when infer schema is true")
 
         if RAY_VERSION >= version.parse("2.10"):
             if len(p) == 1 and isfile(p[0]):  # TODO: very hacky
-                params: Dict[str, Any] = {}
+                params: dict[str, Any] = {}
             else:
                 params = {"file_extensions": ["csv"]}
         else:  # pragma: no cover
@@ -230,14 +238,14 @@ class RayIO(object):
         else:  # pragma: no cover
             raise NotImplementedError(f"{header} is not supported")
 
-    def _load_json(self, p: List[str], columns: Any = None, **kwargs: Any) -> DataFrame:
-        read_options: Dict[str, Any] = {"use_threads": False}
-        parse_options: Dict[str, Any] = {}
+    def _load_json(self, p: list[str], columns: Any = None, **kwargs: Any) -> DataFrame:
+        read_options: dict[str, Any] = {"use_threads": False}
+        parse_options: dict[str, Any] = {}
 
         def _read_json() -> RayDataFrame:  # pragma: no cover
             if RAY_VERSION >= version.parse("2.10"):
                 if len(p) == 1 and isfile(p[0]):  # TODO: very hacky
-                    params: Dict[str, Any] = {"file_extensions": None}
+                    params: dict[str, Any] = {"file_extensions": None}
                 else:
                     params = {"file_extensions": ["json"]}
             elif RAY_VERSION >= version.parse("2.9"):  # pragma: no cover
@@ -272,5 +280,5 @@ class RayIO(object):
             schema = Schema(columns)
             return _read_json()[schema.names].alter_columns(schema)
 
-    def _remote_args(self) -> Dict[str, Any]:
+    def _remote_args(self) -> dict[str, Any]:
         return {"num_cpus": 1}

{fugue-0.9.7.dev2 → fugue-0.9.9.dev0}/fugue_ray/dataframe.py

@@ -6,11 +6,14 @@ 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 triad.utils.pyarrow import cast_pa_table, normalize_large_types_in_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.dataframe.utils import (
+    pa_table_as_array,
+    pa_table_as_dicts,
+)
 from fugue.exceptions import FugueDataFrameOperationError, FugueDatasetEmptyError
 from fugue.plugins import (
     as_array,
@@ -140,7 +143,8 @@ class RayDataFrame(DataFrame):
 
     @property
     def num_partitions(self) -> int:
-        return _rd_num_partitions(self.native)
+        self._native = materialize(self._native)
+        return self._native.num_blocks()
 
     def _drop_cols(self, cols: List[str]) -> DataFrame:
         cols = (self.schema - cols).names
@@ -263,7 +267,7 @@ def _rd_is_df(df: rd.Dataset) -> bool:
 
 @get_num_partitions.candidate(lambda df: isinstance(df, rd.Dataset))
 def _rd_num_partitions(df: rd.Dataset) -> int:
-    return df.num_blocks()
+    return materialize(df).num_blocks()
 
 
 @as_local_bounded.candidate(lambda df: isinstance(df, rd.Dataset))
@@ -305,7 +309,7 @@ def _rd_as_pandas(df: rd.Dataset) -> pd.DataFrame:
 
 @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))
+    return normalize_large_types_in_table(pa.concat_tables(_get_arrow_tables(df)))
 
 
 @as_array.candidate(lambda df, *args, **kwargs: isinstance(df, rd.Dataset))