Repository: pandas-dev/pandas
Pull request #57786 (2024-03-08 23:15:11)

PERF: Allow ensure_index_from_sequence to return RangeIndex
Discovered in https://github.com/pandas-dev/pandas/pull/57441

Builds on https://github.com/pandas-dev/pandas/pull/57752



Patch:
diff --git a/pandas/core/indexes/base.py b/pandas/core/indexes/base.py
index c2df773326..3c01778e05 100644
--- a/pandas/core/indexes/base.py
+++ b/pandas/core/indexes/base.py
@@ -7154,6 +7154,43 @@ class Index(IndexOpsMixin, PandasObject):
         return (len(self),)
 
 
+def maybe_sequence_to_range(sequence) -> Any | range:
+    """
+    Convert a 1D, non-pandas sequence to a range if possible.
+
+    Returns the input if not possible.
+
+    Parameters
+    ----------
+    sequence : 1D sequence
+    names : sequence of str
+
+    Returns
+    -------
+    Any : input or range
+    """
+    if isinstance(sequence, (ABCSeries, Index)):
+        return sequence
+    np_sequence = np.asarray(sequence)
+    if np_sequence.dtype.kind != "i" or len(np_sequence) == 1:
+        return sequence
+    elif len(np_sequence) == 0:
+        return range(0)
+    diff = np_sequence[1] - np_sequence[0]
+    if diff == 0:
+        return sequence
+    elif len(np_sequence) == 2:
+        return range(np_sequence[0], np_sequence[1] + diff, diff)
+    maybe_range_indexer, remainder = np.divmod(np_sequence - np_sequence[0], diff)
+    if (
+        lib.is_range_indexer(maybe_range_indexer, len(maybe_range_indexer))
+        and not remainder.any()
+    ):
+        return range(np_sequence[0], np_sequence[-1] + diff, diff)
+    else:
+        return sequence
+
+
 def ensure_index_from_sequences(sequences, names=None) -> Index:
     """
     Construct an index from sequences of data.
@@ -7172,8 +7209,8 @@ def ensure_index_from_sequences(sequences, names=None) -> Index:
 
     Examples
     --------
-    >>> ensure_index_from_sequences([[1, 2, 3]], names=["name"])
-    Index([1, 2, 3], dtype='int64', name='name')
+    >>> ensure_index_from_sequences([[1, 2, 4]], names=["name"])
+    Index([1, 2, 4], dtype='int64', name='name')
 
     >>> ensure_index_from_sequences([["a", "a"], ["a", "b"]], names=["L1", "L2"])
     MultiIndex([('a', 'a'),
@@ -7189,8 +7226,9 @@ def ensure_index_from_sequences(sequences, names=None) -> Index:
     if len(sequences) == 1:
         if names is not None:
             names = names[0]
-        return Index(sequences[0], name=names)
+        return Index(maybe_sequence_to_range(sequences[0]), name=names)
     else:
+        # TODO: Apply maybe_sequence_to_range to sequences?
         return MultiIndex.from_arrays(sequences, names=names)
 
 
diff --git a/pandas/core/indexes/range.py b/pandas/core/indexes/range.py
index 82bf8d7c70..c573828a22 100644
--- a/pandas/core/indexes/range.py
+++ b/pandas/core/indexes/range.py
@@ -29,7 +29,6 @@ from pandas.util._decorators import (
     doc,
 )
 
-from pandas.core.dtypes import missing
 from pandas.core.dtypes.base import ExtensionDtype
 from pandas.core.dtypes.common import (
     ensure_platform_int,
@@ -475,28 +474,13 @@ class RangeIndex(Index):
         if values.dtype.kind == "i" and values.ndim == 1:
             # GH 46675 & 43885: If values is equally spaced, return a
             # more memory-compact RangeIndex instead of Index with 64-bit dtype
-            if len(values) == 0:
-                return type(self)._simple_new(_empty_range, name=name)
-            elif len(values) == 1:
+            if len(values) == 1:
                 start = values[0]
                 new_range = range(start, start + self.step, self.step)
                 return type(self)._simple_new(new_range, name=name)
-            diff = values[1] - values[0]
-            if not missing.isna(diff) and diff != 0:
-                if len(values) == 2:
-                    # Can skip is_range_indexer check
-                    new_range = range(values[0], values[-1] + diff, diff)
-                    return type(self)._simple_new(new_range, name=name)
-                else:
-                    maybe_range_indexer, remainder = np.divmod(values - values[0], diff)
-                    if (
-                        lib.is_range_indexer(
-                            maybe_range_indexer, len(maybe_range_indexer)
-                        )
-                        and not remainder.any()
-                    ):
-                        new_range = range(values[0], values[-1] + diff, diff)
-                        return type(self)._simple_new(new_range, name=name)
+            maybe_range = ibase.maybe_sequence_to_range(values)
+            if isinstance(maybe_range, range):
+                return type(self)._simple_new(maybe_range, name=name)
         return self._constructor._simple_new(values, name=name)
 
     def _view(self) -> Self:


Test patch:
diff --git a/pandas/tests/indexes/test_base.py b/pandas/tests/indexes/test_base.py
index 4c703c3af9..beee14197b 100644
--- a/pandas/tests/indexes/test_base.py
+++ b/pandas/tests/indexes/test_base.py
@@ -1514,8 +1514,10 @@ class TestIndexUtils:
     @pytest.mark.parametrize(
         "data, names, expected",
         [
-            ([[1, 2, 3]], None, Index([1, 2, 3])),
-            ([[1, 2, 3]], ["name"], Index([1, 2, 3], name="name")),
+            ([[1, 2, 4]], None, Index([1, 2, 4])),
+            ([[1, 2, 4]], ["name"], Index([1, 2, 4], name="name")),
+            ([[1, 2, 3]], None, RangeIndex(1, 4)),
+            ([[1, 2, 3]], ["name"], RangeIndex(1, 4, name="name")),
             (
                 [["a", "a"], ["c", "d"]],
                 None,
@@ -1530,7 +1532,7 @@ class TestIndexUtils:
     )
     def test_ensure_index_from_sequences(self, data, names, expected):
         result = ensure_index_from_sequences(data, names)
-        tm.assert_index_equal(result, expected)
+        tm.assert_index_equal(result, expected, exact=True)
 
     def test_ensure_index_mixed_closed_intervals(self):
         # GH27172
diff --git a/pandas/tests/io/xml/conftest.py b/pandas/tests/io/xml/conftest.py
index 40a94f27e9..8900b0dc46 100644
--- a/pandas/tests/io/xml/conftest.py
+++ b/pandas/tests/io/xml/conftest.py
@@ -11,7 +11,7 @@ def xml_data_path():
     Examples
     --------
     >>> def test_read_xml(xml_data_path):
-    ...     read_xml(xml_data_path / "file.xsl")
+    ...     pd.read_xml(xml_data_path / "file.xsl")
     """
     return Path(__file__).parent.parent / "data" / "xml"
 
@@ -24,7 +24,7 @@ def xml_books(xml_data_path, datapath):
     Examples
     --------
     >>> def test_read_xml(xml_books):
-    ...     read_xml(xml_books)
+    ...     pd.read_xml(xml_books)
     """
     return datapath(xml_data_path / "books.xml")
 
@@ -37,7 +37,7 @@ def xml_doc_ch_utf(xml_data_path, datapath):
     Examples
     --------
     >>> def test_read_xml(xml_doc_ch_utf):
-    ...     read_xml(xml_doc_ch_utf)
+    ...     pd.read_xml(xml_doc_ch_utf)
     """
     return datapath(xml_data_path / "doc_ch_utf.xml")
 
@@ -50,7 +50,7 @@ def xml_baby_names(xml_data_path, datapath):
     Examples
     --------
     >>> def test_read_xml(xml_baby_names):
-    ...     read_xml(xml_baby_names)
+    ...     pd.read_xml(xml_baby_names)
     """
     return datapath(xml_data_path / "baby_names.xml")
 
@@ -63,7 +63,7 @@ def kml_cta_rail_lines(xml_data_path, datapath):
     Examples
     --------
     >>> def test_read_xml(kml_cta_rail_lines):
-    ...     read_xml(
+    ...     pd.read_xml(
     ...         kml_cta_rail_lines,
     ...         xpath=".//k:Placemark",
     ...         namespaces={"k": "http://www.opengis.net/kml/2.2"},
@@ -80,7 +80,7 @@ def xsl_flatten_doc(xml_data_path, datapath):
 
     Examples
     --------
-    >>> def test_read_xsl(xsl_flatten_doc):
+    >>> def test_read_xsl(xsl_flatten_doc, mode):
     ...     with open(
     ...         xsl_flatten_doc, mode, encoding="utf-8" if mode == "r" else None
     ...     ) as f:
@@ -96,7 +96,7 @@ def xsl_row_field_output(xml_data_path, datapath):
 
     Examples
     --------
-    >>> def test_read_xsl(xsl_row_field_output):
+    >>> def test_read_xsl(xsl_row_field_output, mode):
     ...     with open(
     ...         xsl_row_field_output, mode, encoding="utf-8" if mode == "r" else None
     ...     ) as f: