Update _standardize_muli_index_columns to be faster for dataframe del…

…egates.
lsst · Apr 19, 2023 · 9707e0e · 9707e0e
1 parent 220d528
commit 9707e0e
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 9 deletions.
diff --git a/python/lsst/daf/butler/delegates/dataframe.py b/python/lsst/daf/butler/delegates/dataframe.py
@@ -23,10 +23,10 @@
 from __future__ import annotations
 
 import collections.abc
-from typing import Any, Mapping, Optional
+import itertools
+from typing import Any, Mapping, Optional, Sequence
 
 import pandas
-import pyarrow as pa
 from lsst.daf.butler import StorageClassDelegate
 from lsst.daf.butler.formatters.parquet import DataFrameSchema
 from lsst.utils.introspection import get_full_type_name
@@ -112,9 +112,8 @@ def handleParameters(
 
             if isinstance(inMemoryDataset.columns, pandas.MultiIndex):
                 # We have a multi-index dataframe which needs special handling.
-                arrow_table = pa.Table.from_pandas(inMemoryDataset)
                 readColumns = _standardize_multi_index_columns(
-                    arrow_table.schema,
+                    inMemoryDataset.columns,
                     parameters["columns"],
                     stringify=False,
                 )

diff --git a/python/lsst/daf/butler/formatters/parquet.py b/python/lsst/daf/butler/formatters/parquet.py
@@ -107,7 +107,10 @@ def read(self, component: Optional[str] = None) -> Any:
                                 f"Column {par_column} specified in parameters not available in parquet file."
                             )
                 else:
-                    par_columns = _standardize_multi_index_columns(schema, par_columns)
+                    par_columns = _standardize_multi_index_columns(
+                        arrow_schema_to_pandas_index(schema),
+                        par_columns,
+                    )
 
             if len(self.fileDescriptor.parameters):
                 raise ValueError(
@@ -895,7 +898,7 @@ def _split_multi_index_column_names(n: int, names: Iterable[str]) -> List[Sequen
 
 
 def _standardize_multi_index_columns(
-    schema: pa.Schema,
+    pd_index: pd.MultiIndex,
     columns: Any,
     stringify: bool = True,
 ) -> list[str | Sequence[Any]]:
@@ -904,8 +907,8 @@ def _standardize_multi_index_columns(
 
     Parameters
     ----------
-    schema : `pyarrow.Schema`
-        Pyarrow schema.
+    pd_index : `pandas.MultiIndex`
+        Pandas multi-index.
     columns : `list` [`tuple`] or `dict` [`str`, `str` or `list` [`str`]]
         Columns to standardize.
     stringify : `bool`, optional
@@ -916,7 +919,6 @@ def _standardize_multi_index_columns(
     names : `list` [`str`]
         Stringified representation of a multi-index column name.
     """
-    pd_index = arrow_schema_to_pandas_index(schema)
     index_level_names = tuple(pd_index.names)
 
     names: list[str | Sequence[Any]] = []