Pandas 2.0 Support

.github/workflows/ci.yml

@@ -25,43 +25,73 @@ jobs:
       matrix:
         numfocus_nightly: [false]
         os: ["ubuntu-latest"]
-        pyarrow: ["3.0.0", "4.0.1", "nightly"]
+        pandas: [""]
+        pyarrow: ["4.0.1", "nightly"]
         python: ["3.8"]
         include:
           - numfocus_nightly: true
             os: "ubuntu-latest"
+            pandas: ""
             pyarrow: "4.0.1"
             python: "3.10"
           - numfocus_nightly: false
             os: "ubuntu-latest"
+            pandas: "1.5.3"
+            pyarrow: "4.0.1"
+            python: "3.11"
+          - numfocus_nightly: false
+            os: "ubuntu-latest"
+            pandas: "1.5.3"
+            pyarrow: "13.0.0"
+            python: "3.11"
+          - numfocus_nightly: false
+            os: "ubuntu-latest"
+            pandas: ""
             pyarrow: "5.0.0"
             python: "3.9"
           - numfocus_nightly: false
             os: "ubuntu-latest"
+            pandas: ""
             pyarrow: "6.0.1"
             python: "3.9"
           - numfocus_nightly: false
             os: "ubuntu-latest"
+            pandas: ""
             pyarrow: "7.0.0"
             python: "3.10"
           - numfocus_nightly: false
             os: "ubuntu-latest"
+            pandas: ""
             pyarrow: "8.0.1"
             python: "3.10"
           - numfocus_nightly: false
             os: "ubuntu-latest"
+            pandas: ""
             pyarrow: "9.0.0"
             python: "3.10"
           - numfocus_nightly: false
             os: "ubuntu-latest"
+            pandas: ""
             pyarrow: "10.0.1"
             python: "3.11"
           - numfocus_nightly: false
             os: "ubuntu-latest"
+            pandas: ""
             pyarrow: "11.0.0"
             python: "3.11"
+          - numfocus_nightly: false
+            os: "ubuntu-latest"
+            pandas: ""
+            pyarrow: "12.0.0"
+            python: "3.11"
+          - numfocus_nightly: false
+            os: "ubuntu-latest"
+            pandas: ""
+            pyarrow: "13.0.0"
+            python: "3.11"
           - numfocus_nightly: false
             os: "macos-latest"
+            pandas: ""
             pyarrow: "4.0.1"
             python: "3.8"
     continue-on-error: ${{ matrix.numfocus_nightly || matrix.pyarrow == 'nightly' }}
@@ -89,22 +119,28 @@ jobs:
           cache-env: true
           extra-specs: |
             python=${{ matrix.PYTHON_VERSION }}
-      - name: Install repository
-        run: python -m pip install --no-build-isolation --no-deps --disable-pip-version-check -e .
       - name: Install Pyarrow (non-nightly)
-        run: micromamba install pyarrow==${{ matrix.pyarrow }}
-        if: matrix.pyarrow != 'nightly'
+        # Don't pin python as older versions of pyarrow require older versions of python
+        # Pin asv so it doesn't get updated before the benchmarks are run
+        run: micromamba install -y --no-py-pin pyarrow==${{ matrix.pyarrow }} "pandas<2.1.0" "asv<0.6"
+        if: matrix.pyarrow != 'nightly' && matrix.pandas == ''
       - name: Install Pyarrow (nightly)
         # Install both arrow-cpp and pyarrow to make sure that we have the
         # latest nightly of both packages. It is sadly not guaranteed that the
         # nightlies and the latest release would otherwise work together.
         run: micromamba update -c arrow-nightlies -c conda-forge arrow-cpp pyarrow
         if: matrix.pyarrow == 'nightly'
-      - name: Pip Instal NumFOCUS nightly
+      - name: Install Pyarrow (downgrade pandas)
+        run: micromamba install -y --no-py-pin pyarrow==${{ matrix.pyarrow }} pandas==${{ matrix.pandas }}
+        if: matrix.pyarrow != 'nightly' && matrix.pandas != ''
+      - name: Pip Install NumFOCUS nightly
         # NumFOCUS nightly wheels, contains numpy and pandas
         # TODO(gh-45): Re-add numpy
-        run: python -m pip install --pre --upgrade --timeout=60 --extra-index-url https://pypi.anaconda.org/scipy-wheels-nightly/simple pandas
+        # TODO: Remove pandas version pin once https://github.com/pandas-dev/pandas/issues/55014 is fixed
+        run: python -m pip install --pre --upgrade --timeout=60 --extra-index-url https://pypi.anaconda.org/scipy-wheels-nightly/simple "pandas<2.1.0"
         if: matrix.numfocus_nightly
+      - name: Install repository
+        run: python -m pip install --no-build-isolation --no-deps --disable-pip-version-check -e .
       - name: Test import
         run: |
           python -c "import plateau"

CHANGES.rst

@@ -2,6 +2,14 @@
 Changelog
 =========
 
+Plateau 4.2.0 (unreleased)
+==========================
+
+* Support pandas 2
+* Test pyarrow 12 and 13
+* Prevent dask from casting all object dtypes to strings
+* Remove tests for pyarrow<=3 as they fail with pandas>=2
+
 Plateau 4.1.5 (2023-03-14)
 ==========================
 

docs/environment-docs.yml

@@ -3,13 +3,13 @@ channels:
   - conda-forge
 dependencies:
   - python>=3.8
-  - dask[dataframe]
+  - dask<2023.9.2
   - decorator
   - msgpack-python>=0.5.2
   # Currently dask and numpy==1.16.0 clash
   - numpy!=1.15.0,!=1.16.0
-  - pandas>=0.23.0, !=1.0.0
-  - pyarrow>=0.17.1,!=1.0.0
+  - pandas>=0.23.0,!=1.0.0,<2.1.0
+  - pyarrow>=4
   - simplejson
   - minimalkv
   - toolz

environment.yml

@@ -3,14 +3,15 @@ channels:
   - conda-forge
   - nodefaults
 dependencies:
-  - dask!=2021.5.1,!=2021.6.0  # gh475 - 2021.5.1 and 2021.6.0 broke ci, omit those versions
+  # TODO: Investigate issue with dask 2023.9.2
+  - dask!=2021.5.1,!=2021.6.0,<2023.9.2  # gh475 - 2021.5.1 and 2021.6.0 broke ci, omit those versions
   - decorator
   - msgpack-python>=0.5.2
   # Currently dask and numpy==1.16.0 clash
   # TODO: add support for numpy>=1.23
   - numpy!=1.15.0,!=1.16.0
-  - pandas>=0.23.0,!=1.0.0
-  - pyarrow>=0.17.1,!=1.0.0
+  - pandas>=0.23.0,!=1.0.0,<2.1.0
+  - pyarrow>=4
   - simplejson
   - minimalkv>=1.4.2
   - toolz
@@ -36,6 +37,6 @@ dependencies:
   # CLI
   - ipython
   # ASV // Benchmark
-  - asv
+  - asv<0.6
   # Packaging infrastructure
   - python-build

plateau/core/common_metadata.py

@@ -10,12 +10,14 @@
 import pyarrow.parquet as pq
 import simplejson
 from minimalkv import KeyValueStore
+from packaging import version
 
 from plateau.core import naming
 from plateau.core._compat import load_json
 from plateau.core.naming import SINGLE_TABLE
 from plateau.core.utils import ensure_string_type
 from plateau.serialization._parquet import PARQUET_VERSION
+from plateau.serialization._util import schema_metadata_bytes_to_object
 
 _logger = logging.getLogger()
 
@@ -28,6 +30,8 @@
     "normalize_column_order",
 )
 
+PYARROW_LT_13 = version.parse(pa.__version__) < version.parse("13")
+
 
 class SchemaWrapper:
     """Wrapper object for pyarrow.Schema to handle forwards and backwards
@@ -736,7 +740,9 @@ def _dict_to_binary(dct):
     return simplejson.dumps(dct, sort_keys=True).encode("utf8")
 
 
-def empty_dataframe_from_schema(schema, columns=None, date_as_object=False):
+def empty_dataframe_from_schema(
+    schema, columns=None, date_as_object=False, coerce_temporal_nanoseconds=True
+):
     """Create an empty DataFrame from provided schema.
 
     Parameters
@@ -746,14 +752,26 @@ def empty_dataframe_from_schema(schema, columns=None, date_as_object=False):
     columns: Union[None, List[str]]
         Optional list of columns that should be part of the resulting DataFrame. All columns in that list must also be
         part of the provided schema.
+    date_as_object: bool
+        Cast dates to objects.
+    coerce_temporal_nanoseconds: bool
+        Coerce date32, date64, duration and timestamp units to nanoseconds to retain behaviour of pandas 1.x.
+        Only applicable to pandas version >= 2.0 and PyArrow version >= 13.0.0.
 
     Returns
     -------
     DataFrame
         Empty DataFrame with requested columns and types.
     """
-
-    df = schema.internal().empty_table().to_pandas(date_as_object=date_as_object)
+    # HACK: Cast bytes to object in metadata until Pandas bug is fixed: https://github.com/pandas-dev/pandas/issues/50127
+    schema = schema_metadata_bytes_to_object(schema.internal())
+
+    # Prior to pyarrow 13.0.0 coerce_temporal_nanoseconds didn't exist
+    # as it was introduced for backwards compatibility with pandas 1.x
+    _coerce = {}
+    if not PYARROW_LT_13:
+        _coerce["coerce_temporal_nanoseconds"] = coerce_temporal_nanoseconds
+    df = schema.empty_table().to_pandas(date_as_object=date_as_object, **_coerce)
 
     df.columns = df.columns.map(ensure_string_type)
     if columns is not None:

plateau/core/index.py

@@ -6,6 +6,7 @@
 import pandas as pd
 import pyarrow as pa
 import pyarrow.parquet as pq
+from packaging import version
 from toolz.itertoolz import partition_all
 
 import plateau.core._time
@@ -37,6 +38,8 @@
     "PartitionIndex",
 )
 
+PYARROW_LT_13 = version.parse(pa.__version__) < version.parse("13")
+
 
 class IndexBase(CopyMixin):
     """Initialize an IndexBase.
@@ -136,11 +139,21 @@ def __repr__(self) -> str:
             class_=type(self).__name__, attrs=", ".join(repr_str)
         )
 
-    def observed_values(self, date_as_object=True) -> np.ndarray:
+    def observed_values(
+        self, date_as_object=True, coerce_temporal_nanoseconds=True
+    ) -> np.ndarray:
         """Return an array of all observed values."""
         keys = np.array(list(self.index_dct.keys()))
         labeled_array = pa.array(keys, type=self.dtype)
-        return np.array(labeled_array.to_pandas(date_as_object=date_as_object))
+
+        # Prior to pyarrow 13.0.0 coerce_temporal_nanoseconds didn't exist
+        # as it was introduced for backwards compatibility with pandas 1.x
+        _coerce = {}
+        if not PYARROW_LT_13:
+            _coerce["coerce_temporal_nanoseconds"] = coerce_temporal_nanoseconds
+        return np.array(
+            labeled_array.to_pandas(date_as_object=date_as_object, **_coerce)
+        )
 
     @staticmethod
     def normalize_value(dtype: pa.DataType, value: Any) -> Any:
@@ -476,7 +489,10 @@ def as_flat_series(
         table = _index_dct_to_table(
             self.index_dct, column=self.column, dtype=self.dtype
         )
-        df = table.to_pandas(date_as_object=date_as_object)
+        # Prior to pyarrow 13.0.0 coerce_temporal_nanoseconds didn't exist
+        # as it was introduced for backwards compatibility with pandas 1.x
+        _coerce = {} if PYARROW_LT_13 else {"coerce_temporal_nanoseconds": True}
+        df = table.to_pandas(date_as_object=date_as_object, **_coerce)
 
         if predicates is not None:
             # If there is a conjunction without any reference to the index
@@ -862,7 +878,10 @@ def _parquet_bytes_to_dict(column: str, index_buffer: bytes):
     if column_type == pa.timestamp("us"):
         column_type = pa.timestamp("ns")
 
-    df = table.to_pandas()
+    # Prior to pyarrow 13.0.0 coerce_temporal_nanoseconds didn't exist
+    # as it was introduced for backwards compatibility with pandas 1.x
+    _coerce = {} if PYARROW_LT_13 else {"coerce_temporal_nanoseconds": True}
+    df = table.to_pandas(**_coerce)
 
     index_dct = dict(
         zip(df[column].values, (list(x) for x in df[_PARTITION_COLUMN_NAME].values))

plateau/io/dask/compression.py

@@ -2,6 +2,7 @@
 from functools import partial
 from typing import List, Union
 
+import dask
 import dask.dataframe as dd
 import pandas as pd
 
@@ -109,7 +110,8 @@ def pack_payload(df: dd.DataFrame, group_key: Union[List[str], str]) -> dd.DataF
 
     _pack_payload = partial(pack_payload_pandas, group_key=group_key)
 
-    return df.map_partitions(_pack_payload, meta=packed_meta)
+    with dask.config.set({"dataframe.convert-string": False}):
+        return df.map_partitions(_pack_payload, meta=packed_meta)
 
 
 def unpack_payload_pandas(
@@ -154,6 +156,7 @@ def unpack_payload(df: dd.DataFrame, unpack_meta: pd.DataFrame) -> dd.DataFrame:
         )
         return df
 
-    return df.map_partitions(
-        unpack_payload_pandas, unpack_meta=unpack_meta, meta=unpack_meta
-    )
+    with dask.config.set({"dataframe.convert-string": False}):
+        return df.map_partitions(
+            unpack_payload_pandas, unpack_meta=unpack_meta, meta=unpack_meta
+        )