Coverage for python/lsst/daf/butler/formatters/parquet.py: 12%
542 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-14 01:18 -0800
1# This file is part of daf_butler.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
28from __future__ import annotations
30__all__ = (
31 "ArrowAstropySchema",
32 "ArrowNumpySchema",
33 "DataFrameSchema",
34 "ParquetFormatter",
35 "add_pandas_index_to_astropy",
36 "arrow_schema_to_pandas_index",
37 "arrow_to_astropy",
38 "arrow_to_numpy",
39 "arrow_to_numpy_dict",
40 "arrow_to_pandas",
41 "astropy_to_arrow",
42 "astropy_to_pandas",
43 "compute_row_group_size",
44 "numpy_dict_to_arrow",
45 "numpy_to_arrow",
46 "numpy_to_astropy",
47 "pandas_to_arrow",
48 "pandas_to_astropy",
49)
51import collections.abc
52import itertools
53import json
54import logging
55import re
56from collections.abc import Iterable, Sequence
57from fnmatch import fnmatchcase
58from typing import TYPE_CHECKING, Any, cast
60import pyarrow as pa
61import pyarrow.parquet as pq
63from lsst.daf.butler import DatasetProvenance, FormatterV2
64from lsst.daf.butler.delegates.arrowtable import _add_arrow_provenance, _checkArrowCompatibleType
65from lsst.resources import ResourcePath
66from lsst.utils.introspection import get_full_type_name
67from lsst.utils.iteration import ensure_iterable
69log = logging.getLogger(__name__)
71if TYPE_CHECKING:
72 import astropy.table as atable
73 import numpy as np
74 import pandas as pd
76 try:
77 import fsspec
78 from fsspec.spec import AbstractFileSystem
79 except ImportError:
80 fsspec = None
81 AbstractFileSystem = type
83TARGET_ROW_GROUP_BYTES = 1_000_000_000
84ASTROPY_PANDAS_INDEX_KEY = "lsst::arrow::astropy_pandas_index"
87class ParquetFormatter(FormatterV2):
88 """Interface for reading and writing Arrow Table objects to and from
89 Parquet files.
90 """
92 default_extension = ".parq"
93 can_read_from_uri = True
94 can_read_from_local_file = True
96 def can_accept(self, in_memory_dataset: Any) -> bool:
97 # Docstring inherited.
98 return _checkArrowCompatibleType(in_memory_dataset) is not None
100 def read_from_uri(self, uri: ResourcePath, component: str | None = None, expected_size: int = -1) -> Any:
101 # Docstring inherited from Formatter.read.
102 try:
103 fs, path = uri.to_fsspec()
104 except ImportError:
105 log.debug("fsspec not available; falling back to local file access.")
106 # This signals to the formatter to use the read_from_local_file
107 # code path.
108 return NotImplemented
110 return self._read_parquet(path=path, fs=fs, component=component, expected_size=expected_size)
112 def read_from_local_file(self, path: str, component: str | None = None, expected_size: int = -1) -> Any:
113 # Docstring inherited from Formatter.read.
114 return self._read_parquet(path=path, component=component, expected_size=expected_size)
116 def _read_parquet(
117 self,
118 path: str,
119 fs: AbstractFileSystem | None = None,
120 component: str | None = None,
121 expected_size: int = -1,
122 ) -> Any:
123 schema = pq.read_schema(path, filesystem=fs)
125 schema_names = ["ArrowSchema", "DataFrameSchema", "ArrowAstropySchema", "ArrowNumpySchema"]
127 if component in ("columns", "schema") or self.file_descriptor.readStorageClass.name in schema_names:
128 # The schema will be translated to column format
129 # depending on the input type.
130 return schema
131 elif component == "rowcount":
132 # Get the rowcount from the metadata if possible, otherwise count.
133 if b"lsst::arrow::rowcount" in schema.metadata:
134 return int(schema.metadata[b"lsst::arrow::rowcount"])
136 temp_table = pq.read_table(
137 path,
138 filesystem=fs,
139 columns=[schema.names[0]],
140 use_threads=False,
141 use_pandas_metadata=False,
142 )
144 return len(temp_table[schema.names[0]])
146 par_columns = None
147 if self.file_descriptor.parameters:
148 par_columns = self.file_descriptor.parameters.pop("columns", None)
149 if par_columns:
150 has_pandas_multi_index = False
151 if b"pandas" in schema.metadata:
152 md = json.loads(schema.metadata[b"pandas"])
153 if len(md["column_indexes"]) > 1:
154 has_pandas_multi_index = True
156 if not has_pandas_multi_index:
157 # Ensure uniqueness, keeping order.
158 par_columns_in = list(dict.fromkeys(ensure_iterable(par_columns)))
159 file_columns = [name for name in schema.names if not name.startswith("__")]
161 # Do case-sensitive glob-style matching, again ensuring
162 # uniqueness and ordering.
163 par_columns = {}
164 for par_column in par_columns_in:
165 found = False
166 for file_column in file_columns:
167 if fnmatchcase(file_column, par_column):
168 found = True
169 par_columns[file_column] = True
170 if not found:
171 raise ValueError(
172 f"Column {par_column} specified in parameters not available in parquet file."
173 )
174 par_columns = list(par_columns.keys())
175 else:
176 par_columns = _standardize_multi_index_columns(
177 arrow_schema_to_pandas_index(schema),
178 par_columns,
179 )
181 if len(self.file_descriptor.parameters):
182 raise ValueError(
183 f"Unsupported parameters {self.file_descriptor.parameters} in ArrowTable read."
184 )
186 metadata = schema.metadata if schema.metadata is not None else {}
187 arrow_table = pq.read_table(
188 path,
189 filesystem=fs,
190 columns=par_columns,
191 use_threads=False,
192 use_pandas_metadata=(b"pandas" in metadata),
193 )
195 return arrow_table
197 def add_provenance(self, in_memory_dataset: Any, provenance: DatasetProvenance | None = None) -> Any:
198 return _add_arrow_provenance(in_memory_dataset, self.dataset_ref, provenance)
200 def write_local_file(self, in_memory_dataset: Any, uri: ResourcePath) -> None:
201 if isinstance(in_memory_dataset, pa.Schema):
202 pq.write_metadata(in_memory_dataset, uri.ospath)
203 return
205 type_string = _checkArrowCompatibleType(in_memory_dataset)
207 if type_string is None:
208 raise ValueError(
209 f"Unsupported type {get_full_type_name(in_memory_dataset)} of "
210 "inMemoryDataset for ParquetFormatter."
211 )
213 if type_string == "arrow":
214 arrow_table = in_memory_dataset
215 elif type_string == "astropy":
216 arrow_table = astropy_to_arrow(in_memory_dataset)
217 elif type_string == "numpy":
218 arrow_table = numpy_to_arrow(in_memory_dataset)
219 elif type_string == "numpydict":
220 arrow_table = numpy_dict_to_arrow(in_memory_dataset)
221 else:
222 arrow_table = pandas_to_arrow(in_memory_dataset)
224 row_group_size = compute_row_group_size(arrow_table.schema)
226 pq.write_table(arrow_table, uri.ospath, row_group_size=row_group_size)
229def arrow_to_pandas(arrow_table: pa.Table) -> pd.DataFrame:
230 """Convert a pyarrow table to a pandas DataFrame.
232 Parameters
233 ----------
234 arrow_table : `pyarrow.Table`
235 Input arrow table to convert. If the table has ``pandas`` metadata
236 in the schema it will be used in the construction of the
237 ``DataFrame``.
239 Returns
240 -------
241 dataframe : `pandas.DataFrame`
242 Converted pandas dataframe.
243 """
244 dataframe = arrow_table.to_pandas(use_threads=False, integer_object_nulls=True)
246 metadata = arrow_table.schema.metadata if arrow_table.schema.metadata is not None else {}
247 if (key := ASTROPY_PANDAS_INDEX_KEY.encode()) in metadata:
248 pandas_index = metadata[key].decode("UTF8")
249 if pandas_index in arrow_table.schema.names:
250 dataframe.set_index(pandas_index, inplace=True)
251 else:
252 log.warning(
253 "Index column ``%s`` not available for arrow table conversion to DataFrame",
254 pandas_index,
255 )
257 return dataframe
260def arrow_to_astropy(arrow_table: pa.Table) -> atable.Table:
261 """Convert a pyarrow table to an `astropy.Table`.
263 Parameters
264 ----------
265 arrow_table : `pyarrow.Table`
266 Input arrow table to convert. If the table has astropy unit
267 metadata in the schema it will be used in the construction
268 of the ``astropy.Table``.
270 Returns
271 -------
272 table : `astropy.Table`
273 Converted astropy table.
274 """
275 from astropy.table import Table
277 astropy_table = Table(arrow_to_numpy_dict(arrow_table))
279 _apply_astropy_metadata(astropy_table, arrow_table.schema)
281 if (key := ASTROPY_PANDAS_INDEX_KEY) in astropy_table.meta:
282 if astropy_table.meta[key] not in astropy_table.columns:
283 astropy_table.meta.pop(key)
285 return astropy_table
288def arrow_to_numpy(arrow_table: pa.Table) -> np.ndarray | np.ma.MaskedArray:
289 """Convert a pyarrow table to a structured numpy array.
291 Parameters
292 ----------
293 arrow_table : `pyarrow.Table`
294 Input arrow table.
296 Returns
297 -------
298 array : `numpy.ndarray` or `numpy.ma.MaskedArray` (N,)
299 Numpy array table with N rows and the same column names
300 as the input arrow table. Will be masked records if any values
301 in the table are null.
302 """
303 import numpy as np
305 numpy_dict = arrow_to_numpy_dict(arrow_table)
307 has_mask = False
308 dtype: list[tuple] = []
309 for name, col in numpy_dict.items():
310 if len(shape := numpy_dict[name].shape) <= 1:
311 dtype.append((name, col.dtype))
312 else:
313 dtype.append((name, (col.dtype, shape[1:])))
315 if not has_mask and isinstance(col, np.ma.MaskedArray):
316 has_mask = True
318 if has_mask:
319 array = np.ma.mrecords.fromarrays(numpy_dict.values(), dtype=dtype)
320 else:
321 array = np.rec.fromarrays(numpy_dict.values(), dtype=dtype)
322 return array
325def arrow_to_numpy_dict(arrow_table: pa.Table) -> dict[str, np.ndarray]:
326 """Convert a pyarrow table to a dict of numpy arrays.
328 Parameters
329 ----------
330 arrow_table : `pyarrow.Table`
331 Input arrow table.
333 Returns
334 -------
335 numpy_dict : `dict` [`str`, `numpy.ndarray`]
336 Dict with keys as the column names, values as the arrays.
337 """
338 import numpy as np
340 schema = arrow_table.schema
341 metadata = schema.metadata if schema.metadata is not None else {}
343 numpy_dict = {}
345 for name in schema.names:
346 t = schema.field(name).type
348 if arrow_table[name].null_count == 0:
349 # Regular non-masked column
350 col = arrow_table[name].to_numpy()
351 else:
352 # For a masked column, we need to ask arrow to fill the null
353 # values with an appropriately typed value before conversion.
354 # Then we apply the mask to get a masked array of the correct type.
355 null_value: Any
356 match t:
357 case t if t in (pa.float64(), pa.float32(), pa.float16()):
358 null_value = np.nan
359 case t if t in (pa.int64(), pa.int32(), pa.int16(), pa.int8()):
360 null_value = -1
361 case t if t in (pa.bool_(),):
362 null_value = True
363 case t if t in (pa.string(), pa.binary()):
364 null_value = ""
365 case _:
366 # This is the fallback for unsigned ints in particular.
367 null_value = 0
369 col = np.ma.masked_array(
370 data=arrow_table[name].fill_null(null_value).to_numpy(),
371 mask=arrow_table[name].is_null().to_numpy(),
372 fill_value=null_value,
373 )
375 if t in (pa.string(), pa.binary()):
376 col = col.astype(_arrow_string_to_numpy_dtype(schema, name, col))
377 elif isinstance(t, pa.FixedSizeListType):
378 if len(col) > 0:
379 col = np.stack(col)
380 else:
381 # this is an empty column, and needs to be coerced to type.
382 col = col.astype(t.value_type.to_pandas_dtype())
384 shape = _multidim_shape_from_metadata(metadata, t.list_size, name)
385 col = col.reshape((len(arrow_table), *shape))
387 numpy_dict[name] = col
389 return numpy_dict
392def _numpy_dict_to_numpy(numpy_dict: dict[str, np.ndarray]) -> np.ndarray:
393 """Convert a dict of numpy arrays to a structured numpy array.
395 Parameters
396 ----------
397 numpy_dict : `dict` [`str`, `numpy.ndarray`]
398 Dict with keys as the column names, values as the arrays.
400 Returns
401 -------
402 array : `numpy.ndarray` (N,)
403 Numpy array table with N rows and columns names from the dict keys.
404 """
405 return arrow_to_numpy(numpy_dict_to_arrow(numpy_dict))
408def _numpy_to_numpy_dict(np_array: np.ndarray) -> dict[str, np.ndarray]:
409 """Convert a structured numpy array to a dict of numpy arrays.
411 Parameters
412 ----------
413 np_array : `numpy.ndarray`
414 Input numpy array with multiple fields.
416 Returns
417 -------
418 numpy_dict : `dict` [`str`, `numpy.ndarray`]
419 Dict with keys as the column names, values as the arrays.
420 """
421 return arrow_to_numpy_dict(numpy_to_arrow(np_array))
424def numpy_to_arrow(np_array: np.ndarray) -> pa.Table:
425 """Convert a numpy array table to an arrow table.
427 Parameters
428 ----------
429 np_array : `numpy.ndarray`
430 Input numpy array with multiple fields.
432 Returns
433 -------
434 arrow_table : `pyarrow.Table`
435 Converted arrow table.
436 """
437 type_list = _numpy_dtype_to_arrow_types(np_array.dtype)
439 md = {}
440 md[b"lsst::arrow::rowcount"] = str(len(np_array))
442 names = np_array.dtype.names
443 if names is None:
444 names = ()
446 for name in names:
447 _append_numpy_string_metadata(md, name, np_array.dtype[name])
448 _append_numpy_multidim_metadata(md, name, np_array.dtype[name])
450 schema = pa.schema(type_list, metadata=md)
452 arrays = _numpy_style_arrays_to_arrow_arrays(
453 np_array.dtype,
454 len(np_array),
455 np_array,
456 schema,
457 )
459 arrow_table = pa.Table.from_arrays(arrays, schema=schema)
461 return arrow_table
464def numpy_dict_to_arrow(numpy_dict: dict[str, np.ndarray]) -> pa.Table:
465 """Convert a dict of numpy arrays to an arrow table.
467 Parameters
468 ----------
469 numpy_dict : `dict` [`str`, `numpy.ndarray`]
470 Dict with keys as the column names, values as the arrays.
472 Returns
473 -------
474 arrow_table : `pyarrow.Table`
475 Converted arrow table.
477 Raises
478 ------
479 ValueError if columns in numpy_dict have unequal numbers of rows.
480 """
481 dtype, rowcount = _numpy_dict_to_dtype(numpy_dict)
482 type_list = _numpy_dtype_to_arrow_types(dtype)
484 md = {}
485 md[b"lsst::arrow::rowcount"] = str(rowcount)
487 if dtype.names is not None:
488 for name in dtype.names:
489 _append_numpy_string_metadata(md, name, dtype[name])
490 _append_numpy_multidim_metadata(md, name, dtype[name])
492 schema = pa.schema(type_list, metadata=md)
494 arrays = _numpy_style_arrays_to_arrow_arrays(
495 dtype,
496 rowcount,
497 numpy_dict,
498 schema,
499 )
501 arrow_table = pa.Table.from_arrays(arrays, schema=schema)
503 return arrow_table
506def astropy_to_arrow(astropy_table: atable.Table) -> pa.Table:
507 """Convert an astropy table to an arrow table.
509 Parameters
510 ----------
511 astropy_table : `astropy.Table`
512 Input astropy table.
514 Returns
515 -------
516 arrow_table : `pyarrow.Table`
517 Converted arrow table.
518 """
519 from astropy.table import meta
521 type_list = _numpy_dtype_to_arrow_types(astropy_table.dtype)
523 md = {}
524 md[b"lsst::arrow::rowcount"] = str(len(astropy_table))
526 if (key := ASTROPY_PANDAS_INDEX_KEY) in astropy_table.meta:
527 md[key.encode()] = astropy_table.meta[key]
529 for name in astropy_table.dtype.names:
530 _append_numpy_string_metadata(md, name, astropy_table.dtype[name])
531 _append_numpy_multidim_metadata(md, name, astropy_table.dtype[name])
533 meta_yaml = meta.get_yaml_from_table(astropy_table)
534 meta_yaml_str = "\n".join(meta_yaml)
535 md[b"table_meta_yaml"] = meta_yaml_str
537 # Convert type list to fields with metadata.
538 fields = []
539 for name, pa_type in type_list:
540 field_metadata = {}
541 if description := astropy_table[name].description:
542 field_metadata["description"] = description
543 if unit := astropy_table[name].unit:
544 field_metadata["unit"] = str(unit)
545 fields.append(
546 pa.field(
547 name,
548 pa_type,
549 metadata=field_metadata,
550 )
551 )
553 schema = pa.schema(fields, metadata=md)
555 arrays = _numpy_style_arrays_to_arrow_arrays(
556 astropy_table.dtype,
557 len(astropy_table),
558 astropy_table,
559 schema,
560 )
562 arrow_table = pa.Table.from_arrays(arrays, schema=schema)
564 return arrow_table
567def astropy_to_pandas(astropy_table: atable.Table, index: str | None = None) -> pd.DataFrame:
568 """Convert an astropy table to a pandas dataframe via arrow.
570 By going via arrow we avoid pandas masked column bugs (e.g.
571 https://github.com/pandas-dev/pandas/issues/58173)
573 Parameters
574 ----------
575 astropy_table : `astropy.Table`
576 Input astropy table.
577 index : `str`, optional
578 Name of column to set as index.
580 Returns
581 -------
582 dataframe : `pandas.DataFrame`
583 Output pandas dataframe.
584 """
585 index_requested = False
586 if (key := ASTROPY_PANDAS_INDEX_KEY) in astropy_table.meta:
587 _index = astropy_table.meta[key]
588 if _index not in astropy_table.columns:
589 log.warning(
590 "Index column ``%s`` not available for astropy table conversion to DataFrame",
591 _index,
592 )
593 _index = None
594 else:
595 index_requested = True
596 _index = index
598 dataframe = arrow_to_pandas(astropy_to_arrow(astropy_table))
600 # Set the index if we have a valid index name, and either the
601 # index was requested in the call to the function or the dataframe
602 # was not previously indexed with the call to arrow_to_pandas.
603 if isinstance(_index, str) and (index_requested or dataframe.index.name is None):
604 dataframe.set_index(_index, inplace=True)
605 elif _index and index_requested:
606 raise RuntimeError("index must be a string or None.")
608 return dataframe
611def add_pandas_index_to_astropy(astropy_table: atable.Table, index: str) -> None:
612 """Add special metadata to an astropy table to indicate a pandas index.
614 Parameters
615 ----------
616 astropy_table : `astropy.table.Table`
617 Input astropy table.
618 index : `str`
619 Name of column for pandas to set as index, if read as DataFrame.
620 """
621 if index not in astropy_table.columns:
622 raise ValueError("Column ``%s`` not in astropy table columns to use as pandas index.", index)
623 astropy_table.meta[ASTROPY_PANDAS_INDEX_KEY] = index
626def _astropy_to_numpy_dict(astropy_table: atable.Table) -> dict[str, np.ndarray]:
627 """Convert an astropy table to an arrow table.
629 Parameters
630 ----------
631 astropy_table : `astropy.Table`
632 Input astropy table.
634 Returns
635 -------
636 numpy_dict : `dict` [`str`, `numpy.ndarray`]
637 Dict with keys as the column names, values as the arrays.
638 """
639 return arrow_to_numpy_dict(astropy_to_arrow(astropy_table))
642def pandas_to_arrow(dataframe: pd.DataFrame, default_length: int = 10) -> pa.Table:
643 """Convert a pandas dataframe to an arrow table.
645 Parameters
646 ----------
647 dataframe : `pandas.DataFrame`
648 Input pandas dataframe.
649 default_length : `int`, optional
650 Default string length when not in metadata or can be inferred
651 from column.
653 Returns
654 -------
655 arrow_table : `pyarrow.Table`
656 Converted arrow table.
657 """
658 try:
659 arrow_table = pa.Table.from_pandas(dataframe)
660 except pa.ArrowInvalid as e:
661 msg = "; ".join(e.args)
662 msg += "; This is usually because the column is mixed type or has uneven length rows."
663 e.add_note(msg)
664 raise
666 # Update the metadata
667 md = arrow_table.schema.metadata
669 md[b"lsst::arrow::rowcount"] = str(arrow_table.num_rows)
671 # We loop through the arrow table columns because the datatypes have
672 # been checked and converted from pandas objects.
673 for name in arrow_table.column_names:
674 if not name.startswith("__") and arrow_table[name].type == pa.string():
675 if len(arrow_table[name]) > 0:
676 strlen = max(len(row.as_py()) for row in arrow_table[name] if row.is_valid)
677 else:
678 strlen = default_length
679 md[f"lsst::arrow::len::{name}".encode()] = str(strlen)
681 arrow_table = arrow_table.replace_schema_metadata(md)
683 return arrow_table
686def pandas_to_astropy(dataframe: pd.DataFrame) -> atable.Table:
687 """Convert a pandas dataframe to an astropy table, preserving indexes.
689 Parameters
690 ----------
691 dataframe : `pandas.DataFrame`
692 Input pandas dataframe.
694 Returns
695 -------
696 astropy_table : `astropy.table.Table`
697 Converted astropy table.
698 """
699 import pandas as pd
701 if isinstance(dataframe.columns, pd.MultiIndex):
702 raise ValueError("Cannot convert a multi-index dataframe to an astropy table.")
704 return arrow_to_astropy(pandas_to_arrow(dataframe))
707def _pandas_to_numpy_dict(dataframe: pd.DataFrame) -> dict[str, np.ndarray]:
708 """Convert a pandas dataframe to an dict of numpy arrays.
710 Parameters
711 ----------
712 dataframe : `pandas.DataFrame`
713 Input pandas dataframe.
715 Returns
716 -------
717 numpy_dict : `dict` [`str`, `numpy.ndarray`]
718 Dict with keys as the column names, values as the arrays.
719 """
720 return arrow_to_numpy_dict(pandas_to_arrow(dataframe))
723def numpy_to_astropy(np_array: np.ndarray) -> atable.Table:
724 """Convert a numpy table to an astropy table.
726 Parameters
727 ----------
728 np_array : `numpy.ndarray`
729 Input numpy array with multiple fields.
731 Returns
732 -------
733 astropy_table : `astropy.table.Table`
734 Converted astropy table.
735 """
736 from astropy.table import Table
738 return Table(data=np_array, copy=False)
741def arrow_schema_to_pandas_index(schema: pa.Schema) -> pd.Index | pd.MultiIndex:
742 """Convert an arrow schema to a pandas index/multiindex.
744 Parameters
745 ----------
746 schema : `pyarrow.Schema`
747 Input pyarrow schema.
749 Returns
750 -------
751 index : `pandas.Index` or `pandas.MultiIndex`
752 Converted pandas index.
753 """
754 import pandas as pd
756 if b"pandas" in schema.metadata:
757 md = json.loads(schema.metadata[b"pandas"])
758 indexes = md["column_indexes"]
759 len_indexes = len(indexes)
760 else:
761 len_indexes = 0
763 if len_indexes <= 1:
764 return pd.Index(name for name in schema.names if not name.startswith("__"))
765 else:
766 raw_columns = _split_multi_index_column_names(len(indexes), schema.names)
767 return pd.MultiIndex.from_tuples(raw_columns, names=[f["name"] for f in indexes])
770def arrow_schema_to_column_list(schema: pa.Schema) -> list[str]:
771 """Convert an arrow schema to a list of string column names.
773 Parameters
774 ----------
775 schema : `pyarrow.Schema`
776 Input pyarrow schema.
778 Returns
779 -------
780 column_list : `list` [`str`]
781 Converted list of column names.
782 """
783 return list(schema.names)
786class DataFrameSchema:
787 """Wrapper class for a schema for a pandas DataFrame.
789 Parameters
790 ----------
791 dataframe : `pandas.DataFrame`
792 Dataframe to turn into a schema.
793 """
795 def __init__(self, dataframe: pd.DataFrame) -> None:
796 self._schema = dataframe.loc[[False] * len(dataframe)]
798 @classmethod
799 def from_arrow(cls, schema: pa.Schema) -> DataFrameSchema:
800 """Convert an arrow schema into a `DataFrameSchema`.
802 Parameters
803 ----------
804 schema : `pyarrow.Schema`
805 The pyarrow schema to convert.
807 Returns
808 -------
809 dataframe_schema : `DataFrameSchema`
810 Converted dataframe schema.
811 """
812 empty_table = pa.Table.from_pylist([] * len(schema.names), schema=schema)
814 return cls(empty_table.to_pandas())
816 def to_arrow_schema(self) -> pa.Schema:
817 """Convert to an arrow schema.
819 Returns
820 -------
821 arrow_schema : `pyarrow.Schema`
822 Converted pyarrow schema.
823 """
824 arrow_table = pa.Table.from_pandas(self._schema)
826 return arrow_table.schema
828 def to_arrow_numpy_schema(self) -> ArrowNumpySchema:
829 """Convert to an `ArrowNumpySchema`.
831 Returns
832 -------
833 arrow_numpy_schema : `ArrowNumpySchema`
834 Converted arrow numpy schema.
835 """
836 return ArrowNumpySchema.from_arrow(self.to_arrow_schema())
838 def to_arrow_astropy_schema(self) -> ArrowAstropySchema:
839 """Convert to an ArrowAstropySchema.
841 Returns
842 -------
843 arrow_astropy_schema : `ArrowAstropySchema`
844 Converted arrow astropy schema.
845 """
846 return ArrowAstropySchema.from_arrow(self.to_arrow_schema())
848 @property
849 def schema(self) -> np.dtype:
850 return self._schema
852 def __repr__(self) -> str:
853 return repr(self._schema)
855 def __eq__(self, other: object) -> bool:
856 if not isinstance(other, DataFrameSchema):
857 return NotImplemented
859 return self._schema.equals(other._schema)
862class ArrowAstropySchema:
863 """Wrapper class for a schema for an astropy table.
865 Parameters
866 ----------
867 astropy_table : `astropy.table.Table`
868 Input astropy table.
869 """
871 def __init__(self, astropy_table: atable.Table) -> None:
872 self._schema = astropy_table[:0]
874 @classmethod
875 def from_arrow(cls, schema: pa.Schema) -> ArrowAstropySchema:
876 """Convert an arrow schema into a ArrowAstropySchema.
878 Parameters
879 ----------
880 schema : `pyarrow.Schema`
881 Input pyarrow schema.
883 Returns
884 -------
885 astropy_schema : `ArrowAstropySchema`
886 Converted arrow astropy schema.
887 """
888 import numpy as np
889 from astropy.table import Table
891 dtype = _schema_to_dtype_list(schema)
893 data = np.zeros(0, dtype=dtype)
894 astropy_table = Table(data=data)
896 _apply_astropy_metadata(astropy_table, schema)
898 return cls(astropy_table)
900 def to_arrow_schema(self) -> pa.Schema:
901 """Convert to an arrow schema.
903 Returns
904 -------
905 arrow_schema : `pyarrow.Schema`
906 Converted pyarrow schema.
907 """
908 return astropy_to_arrow(self._schema).schema
910 def to_dataframe_schema(self) -> DataFrameSchema:
911 """Convert to a DataFrameSchema.
913 Returns
914 -------
915 dataframe_schema : `DataFrameSchema`
916 Converted dataframe schema.
917 """
918 return DataFrameSchema.from_arrow(astropy_to_arrow(self._schema).schema)
920 def to_arrow_numpy_schema(self) -> ArrowNumpySchema:
921 """Convert to an `ArrowNumpySchema`.
923 Returns
924 -------
925 arrow_numpy_schema : `ArrowNumpySchema`
926 Converted arrow numpy schema.
927 """
928 return ArrowNumpySchema.from_arrow(astropy_to_arrow(self._schema).schema)
930 @property
931 def schema(self) -> atable.Table:
932 return self._schema
934 def __repr__(self) -> str:
935 return repr(self._schema)
937 def __eq__(self, other: object) -> bool:
938 if not isinstance(other, ArrowAstropySchema):
939 return NotImplemented
941 # If this comparison passes then the two tables have the
942 # same column names.
943 if self._schema.dtype != other._schema.dtype:
944 return False
946 for name in self._schema.columns:
947 if not self._schema[name].unit == other._schema[name].unit:
948 return False
949 if not self._schema[name].description == other._schema[name].description:
950 return False
951 if not self._schema[name].format == other._schema[name].format:
952 return False
954 return True
957class ArrowNumpySchema:
958 """Wrapper class for a schema for a numpy ndarray.
960 Parameters
961 ----------
962 numpy_dtype : `numpy.dtype`
963 Numpy dtype to convert.
964 """
966 def __init__(self, numpy_dtype: np.dtype) -> None:
967 self._dtype = numpy_dtype
969 @classmethod
970 def from_arrow(cls, schema: pa.Schema) -> ArrowNumpySchema:
971 """Convert an arrow schema into an `ArrowNumpySchema`.
973 Parameters
974 ----------
975 schema : `pyarrow.Schema`
976 Pyarrow schema to convert.
978 Returns
979 -------
980 numpy_schema : `ArrowNumpySchema`
981 Converted arrow numpy schema.
982 """
983 import numpy as np
985 dtype = _schema_to_dtype_list(schema)
987 return cls(np.dtype(dtype))
989 def to_arrow_astropy_schema(self) -> ArrowAstropySchema:
990 """Convert to an `ArrowAstropySchema`.
992 Returns
993 -------
994 astropy_schema : `ArrowAstropySchema`
995 Converted arrow astropy schema.
996 """
997 import numpy as np
999 return ArrowAstropySchema.from_arrow(numpy_to_arrow(np.zeros(0, dtype=self._dtype)).schema)
1001 def to_dataframe_schema(self) -> DataFrameSchema:
1002 """Convert to a `DataFrameSchema`.
1004 Returns
1005 -------
1006 dataframe_schema : `DataFrameSchema`
1007 Converted dataframe schema.
1008 """
1009 import numpy as np
1011 return DataFrameSchema.from_arrow(numpy_to_arrow(np.zeros(0, dtype=self._dtype)).schema)
1013 def to_arrow_schema(self) -> pa.Schema:
1014 """Convert to a `pyarrow.Schema`.
1016 Returns
1017 -------
1018 arrow_schema : `pyarrow.Schema`
1019 Converted pyarrow schema.
1020 """
1021 import numpy as np
1023 return numpy_to_arrow(np.zeros(0, dtype=self._dtype)).schema
1025 @property
1026 def schema(self) -> np.dtype:
1027 return self._dtype
1029 def __repr__(self) -> str:
1030 return repr(self._dtype)
1032 def __eq__(self, other: object) -> bool:
1033 if not isinstance(other, ArrowNumpySchema):
1034 return NotImplemented
1036 if not self._dtype == other._dtype:
1037 return False
1039 return True
1042def _split_multi_index_column_names(n: int, names: Iterable[str]) -> list[Sequence[str]]:
1043 """Split a string that represents a multi-index column.
1045 PyArrow maps Pandas' multi-index column names (which are tuples in Python)
1046 to flat strings on disk. This routine exists to reconstruct the original
1047 tuple.
1049 Parameters
1050 ----------
1051 n : `int`
1052 Number of levels in the `pandas.MultiIndex` that is being
1053 reconstructed.
1054 names : `~collections.abc.Iterable` [`str`]
1055 Strings to be split.
1057 Returns
1058 -------
1059 column_names : `list` [`tuple` [`str`]]
1060 A list of multi-index column name tuples.
1061 """
1062 column_names: list[Sequence[str]] = []
1064 pattern = re.compile(r"\({}\)".format(", ".join(["'(.*)'"] * n)))
1065 for name in names:
1066 m = re.search(pattern, name)
1067 if m is not None:
1068 column_names.append(m.groups())
1070 return column_names
1073def _standardize_multi_index_columns(
1074 pd_index: pd.MultiIndex,
1075 columns: Any,
1076 stringify: bool = True,
1077) -> list[str | Sequence[Any]]:
1078 """Transform a dictionary/iterable index from a multi-index column list
1079 into a string directly understandable by PyArrow.
1081 Parameters
1082 ----------
1083 pd_index : `pandas.MultiIndex`
1084 Pandas multi-index.
1085 columns : `list` [`tuple`] or `dict` [`str`, `str` or `list` [`str`]]
1086 Columns to standardize.
1087 stringify : `bool`, optional
1088 Should the column names be stringified?
1090 Returns
1091 -------
1092 names : `list` [`str`]
1093 Stringified representation of a multi-index column name.
1094 """
1095 index_level_names = tuple(pd_index.names)
1097 names: list[str | Sequence[Any]] = []
1099 if isinstance(columns, list):
1100 for requested in columns:
1101 if not isinstance(requested, tuple):
1102 raise ValueError(
1103 "Columns parameter for multi-index data frame must be a dictionary or list of tuples. "
1104 f"Instead got a {get_full_type_name(requested)}."
1105 )
1106 if stringify:
1107 names.append(str(requested))
1108 else:
1109 names.append(requested)
1110 else:
1111 if not isinstance(columns, collections.abc.Mapping):
1112 raise ValueError(
1113 "Columns parameter for multi-index data frame must be a dictionary or list of tuples. "
1114 f"Instead got a {get_full_type_name(columns)}."
1115 )
1116 if not set(index_level_names).issuperset(columns.keys()):
1117 raise ValueError(
1118 f"Cannot use dict with keys {set(columns.keys())} to select columns from {index_level_names}."
1119 )
1120 factors = [
1121 ensure_iterable(columns.get(level, pd_index.levels[i]))
1122 for i, level in enumerate(index_level_names)
1123 ]
1124 for requested in itertools.product(*factors):
1125 for i, value in enumerate(requested):
1126 if value not in pd_index.levels[i]:
1127 raise ValueError(f"Unrecognized value {value!r} for index {index_level_names[i]!r}.")
1128 if stringify:
1129 names.append(str(requested))
1130 else:
1131 names.append(requested)
1133 return names
1136def _apply_astropy_metadata(astropy_table: atable.Table, arrow_schema: pa.Schema) -> None:
1137 """Apply any astropy metadata from the schema metadata.
1139 Parameters
1140 ----------
1141 astropy_table : `astropy.table.Table`
1142 Table to apply metadata.
1143 arrow_schema : `pyarrow.Schema`
1144 Arrow schema with metadata.
1145 """
1146 from astropy.table import meta
1148 metadata = arrow_schema.metadata if arrow_schema.metadata is not None else {}
1150 # Check if we have a special astropy metadata header yaml.
1151 meta_yaml = metadata.get(b"table_meta_yaml", None)
1152 if meta_yaml:
1153 meta_yaml = meta_yaml.decode("UTF8").split("\n")
1154 meta_hdr = meta.get_header_from_yaml(meta_yaml)
1156 # Set description, format, unit, meta from the column
1157 # metadata that was serialized with the table.
1158 header_cols = {x["name"]: x for x in meta_hdr["datatype"]}
1159 for col in astropy_table.columns.values():
1160 for attr in ("description", "format", "unit", "meta"):
1161 if attr in header_cols[col.name]:
1162 setattr(col, attr, header_cols[col.name][attr])
1164 if "meta" in meta_hdr:
1165 astropy_table.meta.update(meta_hdr["meta"])
1166 else:
1167 # If we don't have astropy header data, we may have arrow field
1168 # metadata.
1169 for name in arrow_schema.names:
1170 field_metadata = arrow_schema.field(name).metadata
1171 if field_metadata is None:
1172 continue
1173 if (
1174 b"description" in field_metadata
1175 and (description := field_metadata[b"description"].decode("UTF-8")) != ""
1176 ):
1177 astropy_table[name].description = description
1178 if b"unit" in field_metadata and (unit := field_metadata[b"unit"].decode("UTF-8")) != "":
1179 astropy_table[name].unit = unit
1182def _arrow_string_to_numpy_dtype(
1183 schema: pa.Schema, name: str, numpy_column: np.ndarray | None = None, default_length: int = 10
1184) -> str:
1185 """Get the numpy dtype string associated with an arrow column.
1187 Parameters
1188 ----------
1189 schema : `pyarrow.Schema`
1190 Arrow table schema.
1191 name : `str`
1192 Column name.
1193 numpy_column : `numpy.ndarray`, optional
1194 Column to determine numpy string dtype.
1195 default_length : `int`, optional
1196 Default string length when not in metadata or can be inferred
1197 from column.
1199 Returns
1200 -------
1201 dtype_str : `str`
1202 Numpy dtype string.
1203 """
1204 # Special-case for string and binary columns
1205 md_name = f"lsst::arrow::len::{name}"
1206 strlen = default_length
1207 metadata = schema.metadata if schema.metadata is not None else {}
1208 if (encoded := md_name.encode("UTF-8")) in metadata:
1209 # String/bytes length from header.
1210 strlen = int(schema.metadata[encoded])
1211 elif numpy_column is not None and len(numpy_column) > 0:
1212 lengths = [len(row) for row in numpy_column if row]
1213 strlen = max(lengths) if lengths else 0
1215 dtype = f"U{strlen}" if schema.field(name).type == pa.string() else f"|S{strlen}"
1217 return dtype
1220def _append_numpy_string_metadata(metadata: dict[bytes, str], name: str, dtype: np.dtype) -> None:
1221 """Append numpy string length keys to arrow metadata.
1223 All column types are handled, but the metadata is only modified for
1224 string and byte columns.
1226 Parameters
1227 ----------
1228 metadata : `dict` [`bytes`, `str`]
1229 Metadata dictionary; modified in place.
1230 name : `str`
1231 Column name.
1232 dtype : `np.dtype`
1233 Numpy dtype.
1234 """
1235 import numpy as np
1237 if dtype.type is np.str_:
1238 metadata[f"lsst::arrow::len::{name}".encode()] = str(dtype.itemsize // 4)
1239 metadata[f"table::len::{name}".encode()] = str(dtype.itemsize // 4)
1240 elif dtype.type is np.bytes_:
1241 metadata[f"lsst::arrow::len::{name}".encode()] = str(dtype.itemsize)
1242 metadata[f"table::len::{name}".encode()] = str(dtype.itemsize)
1245def _append_numpy_multidim_metadata(metadata: dict[bytes, str], name: str, dtype: np.dtype) -> None:
1246 """Append numpy multi-dimensional shapes to arrow metadata.
1248 All column types are handled, but the metadata is only modified for
1249 multi-dimensional columns.
1251 Parameters
1252 ----------
1253 metadata : `dict` [`bytes`, `str`]
1254 Metadata dictionary; modified in place.
1255 name : `str`
1256 Column name.
1257 dtype : `np.dtype`
1258 Numpy dtype.
1259 """
1260 if len(dtype.shape) > 1:
1261 metadata[f"lsst::arrow::shape::{name}".encode()] = str(dtype.shape)
1264def _multidim_shape_from_metadata(metadata: dict[bytes, bytes], list_size: int, name: str) -> tuple[int, ...]:
1265 """Retrieve the shape from the metadata, if available.
1267 Parameters
1268 ----------
1269 metadata : `dict` [`bytes`, `bytes`]
1270 Metadata dictionary.
1271 list_size : `int`
1272 Size of the list datatype.
1273 name : `str`
1274 Column name.
1276 Returns
1277 -------
1278 shape : `tuple` [`int`]
1279 Shape associated with the column.
1281 Raises
1282 ------
1283 RuntimeError
1284 Raised if metadata is found but has incorrect format.
1285 """
1286 md_name = f"lsst::arrow::shape::{name}"
1287 if (encoded := md_name.encode("UTF-8")) in metadata:
1288 groups = re.search(r"\((.*)\)", metadata[encoded].decode("UTF-8"))
1289 if groups is None:
1290 raise RuntimeError("Illegal value found in metadata.")
1291 shape = tuple(int(x) for x in groups[1].split(",") if x != "")
1292 else:
1293 shape = (list_size,)
1295 return shape
1298def _schema_to_dtype_list(schema: pa.Schema) -> list[tuple[str, tuple[Any] | str]]:
1299 """Convert a pyarrow schema to a numpy dtype.
1301 Parameters
1302 ----------
1303 schema : `pyarrow.Schema`
1304 Input pyarrow schema.
1306 Returns
1307 -------
1308 dtype_list: `list` [`tuple`]
1309 A list with name, type pairs.
1310 """
1311 metadata = schema.metadata if schema.metadata is not None else {}
1313 dtype: list[Any] = []
1314 for name in schema.names:
1315 t = schema.field(name).type
1316 if isinstance(t, pa.FixedSizeListType):
1317 shape = _multidim_shape_from_metadata(metadata, t.list_size, name)
1318 dtype.append((name, (t.value_type.to_pandas_dtype(), shape)))
1319 elif t not in (pa.string(), pa.binary()):
1320 dtype.append((name, t.to_pandas_dtype()))
1321 else:
1322 dtype.append((name, _arrow_string_to_numpy_dtype(schema, name)))
1324 return dtype
1327def _numpy_dtype_to_arrow_types(dtype: np.dtype) -> list[Any]:
1328 """Convert a numpy dtype to a list of arrow types.
1330 Parameters
1331 ----------
1332 dtype : `numpy.dtype`
1333 Numpy dtype to convert.
1335 Returns
1336 -------
1337 type_list : `list` [`object`]
1338 Converted list of arrow types.
1339 """
1340 from math import prod
1342 import numpy as np
1344 type_list: list[Any] = []
1345 if dtype.names is None:
1346 return type_list
1348 for name in dtype.names:
1349 dt = dtype[name]
1350 arrow_type: Any
1351 if len(dt.shape) > 0:
1352 arrow_type = pa.list_(
1353 pa.from_numpy_dtype(cast(tuple[np.dtype, tuple[int, ...]], dt.subdtype)[0].type),
1354 prod(dt.shape),
1355 )
1356 elif dt.type == np.datetime64:
1357 time_unit = "ns" if "ns" in dt.str else "us"
1358 # The pa.timestamp() is the correct datatype to round-trip
1359 # a numpy datetime64[ns] or datetime[us] array.
1360 arrow_type = pa.timestamp(time_unit)
1361 else:
1362 try:
1363 arrow_type = pa.from_numpy_dtype(dt.type)
1364 except pa.ArrowNotImplementedError as e:
1365 msg = f"Could not serialize column {name} (type {str(dt)}) to Parquet."
1366 if dt == np.dtype("O"):
1367 msg += " This is usually because the column is mixed type or has uneven length rows."
1368 e.add_note(msg)
1369 raise
1370 type_list.append((name, arrow_type))
1372 return type_list
1375def _numpy_dict_to_dtype(numpy_dict: dict[str, np.ndarray]) -> tuple[np.dtype, int]:
1376 """Extract equivalent table dtype from dict of numpy arrays.
1378 Parameters
1379 ----------
1380 numpy_dict : `dict` [`str`, `numpy.ndarray`]
1381 Dict with keys as the column names, values as the arrays.
1383 Returns
1384 -------
1385 dtype : `numpy.dtype`
1386 dtype of equivalent table.
1387 rowcount : `int`
1388 Number of rows in the table.
1390 Raises
1391 ------
1392 ValueError if columns in numpy_dict have unequal numbers of rows.
1393 """
1394 import numpy as np
1396 dtype_list: list[tuple] = []
1397 rowcount = 0
1398 for name, col in numpy_dict.items():
1399 if rowcount == 0:
1400 rowcount = len(col)
1401 if len(col) != rowcount:
1402 raise ValueError(f"Column {name} has a different number of rows.")
1403 if len(col.shape) == 1:
1404 dtype_list.append((name, col.dtype))
1405 else:
1406 dtype_list.append((name, (col.dtype, col.shape[1:])))
1407 dtype = np.dtype(dtype_list)
1409 return (dtype, rowcount)
1412def _numpy_style_arrays_to_arrow_arrays(
1413 dtype: np.dtype,
1414 rowcount: int,
1415 np_style_arrays: dict[str, np.ndarray] | np.ndarray | atable.Table,
1416 schema: pa.Schema,
1417) -> list[pa.Array]:
1418 """Convert numpy-style arrays to arrow arrays.
1420 Parameters
1421 ----------
1422 dtype : `numpy.dtype`
1423 Numpy dtype of input table/arrays.
1424 rowcount : `int`
1425 Number of rows in input table/arrays.
1426 np_style_arrays : `dict` [`str`, `np.ndarray`] or `np.ndarray`
1427 or `astropy.table.Table`
1428 Arrays to convert to arrow.
1429 schema : `pyarrow.Schema`
1430 Schema of arrow table.
1432 Returns
1433 -------
1434 arrow_arrays : `list` [`pyarrow.Array`]
1435 List of converted pyarrow arrays.
1436 """
1437 import numpy as np
1439 arrow_arrays: list[pa.Array] = []
1440 if dtype.names is None:
1441 return arrow_arrays
1443 for name in dtype.names:
1444 dt = dtype[name]
1445 val: Any
1446 if len(dt.shape) > 0:
1447 if rowcount > 0:
1448 val = np.split(np_style_arrays[name].ravel(), rowcount)
1449 else:
1450 val = []
1451 else:
1452 val = np_style_arrays[name]
1454 try:
1455 arrow_arrays.append(pa.array(val, type=schema.field(name).type))
1456 except pa.ArrowNotImplementedError as err:
1457 # Check if val is big-endian.
1458 if (np.little_endian and val.dtype.byteorder == ">") or (
1459 not np.little_endian and val.dtype.byteorder == "="
1460 ):
1461 # We need to convert the array to little-endian.
1462 val2 = val.byteswap()
1463 val2.dtype = val2.dtype.newbyteorder("<")
1464 arrow_arrays.append(pa.array(val2, type=schema.field(name).type))
1465 else:
1466 # This failed for some other reason so raise the exception.
1467 raise err
1469 return arrow_arrays
1472def compute_row_group_size(schema: pa.Schema, target_size: int = TARGET_ROW_GROUP_BYTES) -> int:
1473 """Compute approximate row group size for a given arrow schema.
1475 Given a schema, this routine will compute the number of rows in a row group
1476 that targets the persisted size on disk (or smaller). The exact size on
1477 disk depends on the compression settings and ratios; typical binary data
1478 tables will have around 15-20% compression with the pyarrow default
1479 ``snappy`` compression algorithm.
1481 Parameters
1482 ----------
1483 schema : `pyarrow.Schema`
1484 Arrow table schema.
1485 target_size : `int`, optional
1486 The target size (in bytes).
1488 Returns
1489 -------
1490 row_group_size : `int`
1491 Number of rows per row group to hit the target size.
1492 """
1493 bit_width = 0
1495 metadata = schema.metadata if schema.metadata is not None else {}
1497 for name in schema.names:
1498 t = schema.field(name).type
1500 if t in (pa.string(), pa.binary()):
1501 md_name = f"lsst::arrow::len::{name}"
1503 if (encoded := md_name.encode("UTF-8")) in metadata:
1504 # String/bytes length from header.
1505 strlen = int(schema.metadata[encoded])
1506 else:
1507 # We don't know the string width, so guess something.
1508 strlen = 10
1510 # Assuming UTF-8 encoding, and very few wide characters.
1511 t_width = 8 * strlen
1512 elif isinstance(t, pa.FixedSizeListType):
1513 if t.value_type == pa.null():
1514 t_width = 0
1515 else:
1516 t_width = t.list_size * t.value_type.bit_width
1517 elif t == pa.null():
1518 t_width = 0
1519 elif isinstance(t, pa.ListType):
1520 if t.value_type == pa.null():
1521 t_width = 0
1522 else:
1523 # This is a variable length list, just choose
1524 # something arbitrary.
1525 t_width = 10 * t.value_type.bit_width
1526 else:
1527 t_width = t.bit_width
1529 bit_width += t_width
1531 # Insist it is at least 1 byte wide to avoid any divide-by-zero errors.
1532 if bit_width < 8:
1533 bit_width = 8
1535 byte_width = bit_width // 8
1537 return target_size // byte_width