Fix pandas-dev#58421: Index[timestamp[pyarrow]].union with itself return object type

Author: afonso-antunes

pandas/core/reshape/concat.py

@@ -16,6 +16,8 @@
 
 import numpy as np
 
+import pandas
+
 from pandas._libs import lib
 from pandas.util._decorators import set_module
 from pandas.util._exceptions import find_stack_level
@@ -47,6 +49,10 @@
 )
 from pandas.core.internals import concatenate_managers
 
+from pandas.core.dtypes.common import is_extension_array_dtype
+
+from pandas.core.construction import array
+
 if TYPE_CHECKING:
     from collections.abc import (
         Callable,
@@ -819,7 +825,20 @@ def _get_sample_object(
 
 
 def _concat_indexes(indexes) -> Index:
-    return indexes[0].append(indexes[1:])
+    # try to preserve extension types such as timestamp[pyarrow]
+    values = []
+    for idx in indexes:
+        values.extend(idx._values if hasattr(idx, "_values") else idx)
+
+    # use the first index as a sample to infer the desired dtype
+    sample = indexes[0]
+    try:
+        # this helps preserve extension types like timestamp[pyarrow]
+        arr = array(values, dtype=sample.dtype)
+    except Exception:
+        arr = array(values)  # fallback
+
+    return Index(arr)
 
 
 def validate_unique_levels(levels: list[Index]) -> None:
@@ -876,14 +895,33 @@ def _make_concat_multiindex(indexes, keys, levels=None, names=None) -> MultiInde
 
         concat_index = _concat_indexes(indexes)
 
-        # these go at the end
         if isinstance(concat_index, MultiIndex):
             levels.extend(concat_index.levels)
             codes_list.extend(concat_index.codes)
         else:
-            codes, categories = factorize_from_iterable(concat_index)
-            levels.append(categories)
-            codes_list.append(codes)
+            # handle the case where the resulting index is a flat Index
+            # but contains tuples (i.e., a collapsed MultiIndex)
+            if isinstance(concat_index[0], tuple):
+                # retrieve the original dtypes
+                original_dtypes = [lvl.dtype for lvl in indexes[0].levels]
+
+                unzipped = list(zip(*concat_index))
+                for i, level_values in enumerate(unzipped):
+                    # reconstruct each level using original dtype
+                    arr = array(level_values, dtype=original_dtypes[i])
+                    level_codes, _ = factorize_from_iterable(arr)
+                    levels.append(ensure_index(arr))
+                    codes_list.append(level_codes)
+            else:
+                # simple indexes factorize directly
+                codes, categories = factorize_from_iterable(concat_index)
+                values = getattr(concat_index, "_values", concat_index)
+                if is_extension_array_dtype(values):
+                    levels.append(values)
+                else:
+                    levels.append(categories)
+                codes_list.append(codes)
+
 
         if len(names) == len(levels):
             names = list(names)

pandas/tests/frame/methods/test_concat_arrow_index.py

@@ -0,0 +1,47 @@
+import pandas as pd
+import pytest
+
+
+schema = {
+    "id": "int64[pyarrow]",
+    "time": "timestamp[s][pyarrow]",
+    "value": "float[pyarrow]",
+}
+
+@pytest.mark.parametrize("dtype", ["timestamp[s][pyarrow]"])
+def test_concat_preserves_pyarrow_timestamp(dtype):
+    dfA = (
+        pd.DataFrame(
+            [
+                (0, "2021-01-01 00:00:00", 5.3),
+                (1, "2021-01-01 00:01:00", 5.4),
+                (2, "2021-01-01 00:01:00", 5.4),
+                (3, "2021-01-01 00:02:00", 5.5),
+            ],
+            columns=schema,
+        )
+        .astype(schema)
+        .set_index(["id", "time"])
+    )
+
+    dfB = (
+        pd.DataFrame(
+            [
+                (1, "2022-01-01 08:00:00", 6.3),
+                (2, "2022-01-01 08:01:00", 6.4),
+                (3, "2022-01-01 08:02:00", 6.5),
+            ],
+            columns=schema,
+        )
+        .astype(schema)
+        .set_index(["id", "time"])
+    )
+
+    df = pd.concat([dfA, dfB], keys=[0, 1], names=["run"])
+    
+    # chech whether df.index is multiIndex
+    assert isinstance(df.index, pd.MultiIndex), f"Expected MultiIndex, but received {type(df.index)}"
+    
+    # Verifying special dtype timestamp[s][pyarrow] stays intact after concat
+    assert df.index.levels[2].dtype == dtype, f"Expected {dtype}, but received {df.index.levels[2].dtype}"
+