Fixes for pandas 2&3 (#976)

Author: martindurant

.github/workflows/main.yaml

@@ -42,6 +42,43 @@ jobs:
         shell: bash -l {0}
         run: |
           pytest --verbose --cov=fastparquet
+
+  linux-pandas3:
+    name: ${{ matrix.CONDA_ENV }}-pandas3-pytest
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        CONDA_ENV: [py312, py313, py314]
+    steps:
+      - name: APT
+        run: sudo apt-get install liblzo2-dev
+
+      - name: Checkout
+        uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Fetch upstream tags
+        run: |
+          git remote add upstream https://github.com/dask/fastparquet.git
+          git fetch upstream --tags
+
+      - name: Setup conda
+        uses: conda-incubator/setup-miniconda@v3
+        with:
+          environment-file: ci/environment-${{ matrix.CONDA_ENV }}-pandas3.yml
+
+      - name: pip-install
+        shell: bash -l {0}
+        run: |
+          pip install -e . --no-deps
+
+      - name: Run Tests
+        shell: bash -l {0}
+        run: |
+          pytest --verbose --cov=fastparquet
+
   v2:
     name: v2-py310
     runs-on: ubuntu-latest
@@ -156,3 +193,4 @@ jobs:
         shell: bash -l {0}
         run: |
           pytest --verbose --cov=fastparquet
+

ci/environment-py312-pandas3.yml

@@ -0,0 +1,20 @@
+name: test_env
+channels:
+  - conda-forge
+dependencies:
+  - python=3.12
+  - bson
+  - lz4
+  - lzo
+  - pytest
+  - pandas>=3.0
+  - dask
+  - pytest-cov
+  - thrift
+  - numpy
+  - cramjam
+  - packaging
+  - orjson
+  - ujson
+  - python-rapidjson
+  - pyarrow

ci/environment-py313-pandas3.yml

@@ -0,0 +1,20 @@
+name: test_env
+channels:
+  - conda-forge
+dependencies:
+  - python=3.13
+  - bson
+  - lz4
+  - lzo
+  - pytest
+  - pandas>=3.0
+  - dask
+  - pytest-cov
+  - thrift
+  - numpy
+  - cramjam
+  - packaging
+  - orjson
+  - ujson
+  - python-rapidjson
+  - pyarrow

ci/environment-py314-pandas3.yml

@@ -0,0 +1,20 @@
+name: test_env
+channels:
+  - conda-forge
+dependencies:
+  - python=3.14
+  - bson
+  - lz4
+  - lzo
+  - pytest
+  - pandas>=3.0
+  - dask
+  - pytest-cov
+  - thrift
+  - numpy
+  - cramjam
+  - packaging
+  - orjson
+  - ujson
+  - python-rapidjson
+  - pyarrow

fastparquet/api.py

@@ -15,6 +15,10 @@
 from fastparquet.util import (default_open, default_remove, ParquetException, val_to_num,
                    ops, ensure_bytes, ensure_str, check_column_names, metadata_from_many,
                    ex_from_sep, _strip_path_tail, get_fs, PANDAS_VERSION, join_path)
+from fastparquet.encoding import _HAVE_ARROW, _USE_ARROW_STRINGS
+
+if _HAVE_ARROW:
+    import pyarrow as _pa
 
 
 # Find in names of partition files the integer matching "**part.*.parquet",
@@ -102,11 +106,16 @@ class ParquetFile(object):
     _pdm = None
     _kvm = None
     _categories = None
+    _user_dtypes = False  # whether _base_dtype was user-supplied
 
     def __init__(self, fn, verify=False, open_with=default_open, root=False,
                  sep=None, fs=None, pandas_nulls=True, dtypes=None):
         self.pandas_nulls = pandas_nulls
         self._base_dtype = dtypes
+        # Track whether _base_dtype was user-supplied (True) or auto-inferred (False).
+        # This is needed so that an auto-inferred dtype('O') for a UTF-8 column
+        # does not prevent the arrow string path from activating.
+        self._user_dtypes = dtypes is not None
         self.tz = None
         self._columns_dtype = None
         if open_with is default_open and fs is None:
@@ -346,7 +355,7 @@ def row_group_filename(self, rg):
 
     def read_row_group_file(self, rg, columns, categories, index=None,
                             assign=None, partition_meta=None, row_filter=False,
-                            infile=None):
+                            infile=None, arrow_string_columns=None):
         """ Open file for reading, and process it as a row-group
 
         assign is None if this method is called directly (not from to_pandas),
@@ -389,7 +398,7 @@ def read_row_group_file(self, rg, columns, categories, index=None,
             f, rg, columns, categories, self.schema, self.cats,
             selfmade=self.selfmade, index=index,
             assign=assign, scheme=self.file_scheme, partition_meta=partition_meta,
-            row_filter=row_filter
+            row_filter=row_filter, arrow_string_columns=arrow_string_columns,
         )
         if ret:
             return df
@@ -665,26 +674,26 @@ def _column_filter(self, df, filters):
                 if name in self.cats:
                     continue
                 if op == 'in':
-                    out |= df[name].isin(val).values
+                    out |= np.asarray(df[name].isin(val), dtype=bool)
                 elif op == "not in":
-                    out |= ~df[name].isin(val).values
+                    out |= ~np.asarray(df[name].isin(val), dtype=bool)
                 elif op in ops:
-                    out |= ops[op](df[name], val).values
+                    out |= np.asarray(ops[op](df[name], val), dtype=bool)
                 elif op == "~":
-                    out |= ~df[name].values
+                    out |= ~np.asarray(df[name], dtype=bool)
             else:
                 and_part = np.ones(len(df), dtype=bool)
                 for name, op, val in or_part:
                     if name in self.cats:
                         continue
                     if op == 'in':
-                        and_part &= df[name].isin(val).values
+                        and_part &= np.asarray(df[name].isin(val), dtype=bool)
                     elif op == "not in":
-                        and_part &= ~df[name].isin(val).values
+                        and_part &= ~np.asarray(df[name].isin(val), dtype=bool)
                     elif op in ops:
-                        and_part &= ops[op](df[name].values, val)
+                        and_part &= np.asarray(ops[op](df[name], val), dtype=bool)
                     elif op == "~":
-                        and_part &= ~df[name].values
+                        and_part &= ~np.asarray(df[name], dtype=bool)
                 out |= and_part
         return out
 
@@ -771,6 +780,44 @@ def to_pandas(self, columns=None, categories=None, filters=[],
             import json
             df.attrs = json.loads(self.key_value_metadata["PANDAS_ATTRS"])
 
+        # Build per-column accumulators for arrow string columns.
+        # Object-dtype columns whose parquet schema marks them as UTF-8 will
+        # use the fast Cython path that builds ArrowStringArray directly.
+        # Skip columns whose dtype was explicitly requested as object (via dtypes=)
+        # to honour the caller's preference and to handle schema-evolution cases.
+        arrow_string_columns = {}
+        if _HAVE_ARROW and _USE_ARROW_STRINGS:
+            _cats = categories or {}
+            # Resolve user-supplied dtype overrides only (not auto-inferred ones).
+            # Auto-inferred dtypes always show dtype('O') for UTF-8 columns, so
+            # including them would block the arrow path for all string columns.
+            _user_dtypes = {}
+            if self._user_dtypes and self._base_dtype:
+                _user_dtypes.update(self._base_dtype)
+            if dtypes is not None:
+                _user_dtypes.update(dtypes)
+            # Detect schema-evolution columns: absent from some row groups.
+            # For those, the existing object-array fallback is more reliable.
+            _schema_evo_cols = set()
+            if len(rgs) > 1:
+                for rg in rgs:
+                    rg_cols = {'.'.join(c.meta_data.path_in_schema)
+                               for c in rg.columns}
+                    for col in df.columns:
+                        if col not in rg_cols:
+                            _schema_evo_cols.add(col)
+            for col in df.columns:
+                if (df[col].dtype == np.dtype('O')
+                        and col not in _cats
+                        and col not in _schema_evo_cols
+                        and _user_dtypes.get(col) != np.dtype('O')):
+                    arrow_string_columns[col] = []
+            # Also check index columns that are object dtype
+            if hasattr(df.index, 'dtype') and df.index.dtype == np.dtype('O'):
+                idx_name = df.index.name
+                if idx_name and idx_name in columns and idx_name not in _schema_evo_cols:
+                    arrow_string_columns[idx_name] = []
+
         start = 0
         if self.file_scheme == 'simple':
             infile = self.open(self.fn, 'rb')
@@ -789,8 +836,27 @@ def to_pandas(self, columns=None, categories=None, filters=[],
                      for (name, v) in views.items()}
             self.read_row_group_file(rg, columns, categories, index,
                                      assign=parts, partition_meta=self.partition_meta,
-                                     row_filter=sel, infile=infile)
+                                     row_filter=sel, infile=infile,
+                                     arrow_string_columns=arrow_string_columns)
             start += thislen
+
+        # Finalize arrow string columns: concatenate per-page arrays and assign.
+        if _HAVE_ARROW and _USE_ARROW_STRINGS:
+            for col, parts in arrow_string_columns.items():
+                if not parts:
+                    continue
+                final_arr = _pa.concat_arrays(parts)
+                if len(final_arr) != len(df):
+                    # Schema evolution: column absent in some row groups.  The
+                    # pre-allocated object array already has None for the missing
+                    # rows; skip the arrow conversion so callers see partial data.
+                    continue
+                arrow_str = pd.arrays.ArrowStringArray(_pa.chunked_array([final_arr]))
+                if col in df.columns:
+                    df[col] = pd.Series(arrow_str, index=df.index, name=col)
+                elif col == df.index.name:
+                    df.index = pd.Index(arrow_str, name=col)
+
         return df
 
     def pre_allocate(self, size, columns, categories, index, dtypes=None):

fastparquet/benchmarks/columns.py

@@ -69,6 +69,9 @@ def time_column():
                 d.loc[n//2, col] = np.nan
             elif d.dtypes[col].kind in ['i', 'u']:
                 continue
+            elif d.dtypes[col].kind == "b":
+                d[col] = d[col].astype("boolean")
+                d.loc[n // 2, col] = pd.NA
             else:
                 d.loc[n//2, col] = None
             with measure('%s: write, with null, has_null=True' % d.dtypes[col], result):

fastparquet/converted_types.py

@@ -167,9 +167,13 @@ def convert(data, se, timestamp96=True, dtype=None):
     if ctype is None:
         return data
     if ctype == parquet_thrift.ConvertedType.UTF8:
+        import pandas as pd
+        if isinstance(data.dtype, pd.StringDtype):
+            # pandas 3+: data is already an arrow-backed string array (decoded),
+            # no further conversion needed.
+            return data
         if data.dtype != "O" or (len(data) == 1 and not isinstance(data[0], str)):
-            # fixed string
-            import pandas as pd
+            # fixed-length byte string (e.g. FIXED_LEN_BYTE_ARRAY): decode to str
             return pd.Series(data).str.decode("utf8").values
         # already converted in speedups.unpack_byte_array
         return data
@@ -204,7 +208,7 @@ def convert(data, se, timestamp96=True, dtype=None):
         return data.view('datetime64[ns]')
     elif ctype == parquet_thrift.ConvertedType.TIME_MILLIS:
         # this was not covered by new pandas time units
-        data = data.astype('int64', copy=False)
+        data = data.astype('int64')
         time_shift(data, 1000000)
         return data.view('timedelta64[ns]')
     elif ctype == parquet_thrift.ConvertedType.TIMESTAMP_MILLIS: