Also fix pandas 2 :)

Author: martindurant

fastparquet/api.py

@@ -15,7 +15,7 @@
 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
+from fastparquet.encoding import _HAVE_ARROW, _USE_ARROW_STRINGS
 
 if _HAVE_ARROW:
     import pyarrow as _pa
@@ -786,7 +786,7 @@ def to_pandas(self, columns=None, categories=None, filters=[],
         # 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:
+        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
@@ -841,7 +841,7 @@ def to_pandas(self, columns=None, categories=None, filters=[],
             start += thislen
 
         # Finalize arrow string columns: concatenate per-page arrays and assign.
-        if _HAVE_ARROW:
+        if _HAVE_ARROW and _USE_ARROW_STRINGS:
             for col, parts in arrow_string_columns.items():
                 if not parts:
                     continue

fastparquet/core.py

@@ -2,7 +2,7 @@
 import pandas as pd
 
 from fastparquet import encoding
-from fastparquet.encoding import read_plain, _HAVE_ARROW
+from fastparquet.encoding import read_plain, _HAVE_ARROW, _USE_ARROW_STRINGS
 import fastparquet.cencoding as encoding
 from fastparquet.compression import decompress_data, rev_map, decom_into
 from fastparquet.converted_types import convert, simple, converts_inplace
@@ -308,7 +308,7 @@ def read_data_page_v2(infile, schema_helper, se, data_header2, cmd,
         raw_bytes = decompress_data(np.frombuffer(infile.read(size), "uint8"),
                                     uncompressed_page_size, codec)
         # Arrow path: build ArrowStringArray directly without Python str objects.
-        if (arrow_parts is not None and _HAVE_ARROW
+        if (arrow_parts is not None and _HAVE_ARROW and _USE_ARROW_STRINGS
                 and se.converted_type == parquet_thrift.ConvertedType.UTF8
                 and not use_cat):
             raw_np = np.asarray(raw_bytes, dtype=np.uint8)
@@ -487,9 +487,11 @@ def read_col(column, schema_helper, infile, use_cat=False,
     rows = row_filter.sum() if isinstance(row_filter, np.ndarray) else cmd.num_values
 
     # Determine whether to use the arrow accumulation path for this column.
-    # Conditions: pyarrow available, utf=True schema element, caller provided a list.
+    # Conditions: pyarrow available, arrow strings are the default, utf=True
+    # schema element, and caller provided a list to accumulate into.
     use_arrow = (
         _HAVE_ARROW
+        and _USE_ARROW_STRINGS
         and arrow_parts is not None
         and not use_cat
         and se.converted_type == parquet_thrift.ConvertedType.UTF8

fastparquet/encoding.py

@@ -11,6 +11,19 @@
 except ImportError:
     _HAVE_ARROW = False
 
+# Use the arrow string path only when pandas itself defaults to arrow strings.
+# In pandas 3+, pd.options.future.infer_string is True by default.
+# In pandas 2, it is False (object dtype is the default for strings).
+# This can also be explicitly enabled in pandas 2 with
+#   pd.options.future.infer_string = True
+# so we respect that too.
+try:
+    import pandas as _pd
+    _USE_ARROW_STRINGS = bool(getattr(getattr(_pd.options, 'future', None),
+                                      'infer_string', False))
+except Exception:
+    _USE_ARROW_STRINGS = False
+
 
 def read_plain_boolean(raw_bytes, count, out=None):
     data = np.frombuffer(raw_bytes, dtype='uint8')
@@ -45,9 +58,11 @@ def read_plain(raw_bytes, type_, count, width=0, utf=False, stat=False):
                 return np.array([bytes(raw_bytes).decode()], dtype='O')
             else:
                 return np.array([bytes(raw_bytes)], dtype='O')
-        if utf and _HAVE_ARROW:
+        if utf and _HAVE_ARROW and _USE_ARROW_STRINGS:
             # Build an ArrowStringArray directly from the packed bytes, without
             # creating intermediate Python str objects in the inner loop.
+            # Only used when pandas defaults to arrow strings (pandas 3+, or
+            # pandas 2 with future.infer_string=True).
             raw_np = np.frombuffer(memoryview(raw_bytes), dtype=np.uint8)
             return _unpack_arrow(raw_np, count)
         return unpack_byte_array(raw_bytes, count, utf=utf)

fastparquet/test/test_api.py

@@ -82,12 +82,18 @@ def test_text_schema(tempdir):
     write(fn, df)
     p = ParquetFile(fn)
     t = p.schema.text
-    expected = ('- schema: \n'
-                '| - A: DOUBLE, OPTIONAL\n'
-                '| - B: DOUBLE, OPTIONAL\n'
-                '| - C: BYTE_ARRAY, UTF8, OPTIONAL\n'
-                '  - D: INT64, TIMESTAMP[MICROS], TIMESTAMP_MICROS, OPTIONAL')
-    assert t == expected
+    # pandas 2 uses nanosecond timestamps (NANOS), pandas 3 uses microseconds (MICROS)
+    expected_micros = ('- schema: \n'
+                       '| - A: DOUBLE, OPTIONAL\n'
+                       '| - B: DOUBLE, OPTIONAL\n'
+                       '| - C: BYTE_ARRAY, UTF8, OPTIONAL\n'
+                       '  - D: INT64, TIMESTAMP[MICROS], TIMESTAMP_MICROS, OPTIONAL')
+    expected_nanos = ('- schema: \n'
+                      '| - A: DOUBLE, OPTIONAL\n'
+                      '| - B: DOUBLE, OPTIONAL\n'
+                      '| - C: BYTE_ARRAY, UTF8, OPTIONAL\n'
+                      '  - D: INT64, TIMESTAMP[NANOS], OPTIONAL')
+    assert t in (expected_micros, expected_nanos)
     assert repr(p.schema) == "<Parquet Schema with 5 entries>"