Coverage for python/lsst/daf/butler/formatters/parquet.py: 96%
578 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-13 08:57 +0000
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 "pandas_to_numpy",
50)
52import collections.abc
53import contextlib
54import itertools
55import json
56import logging
57import re
58from collections.abc import Generator, Iterable, Sequence
59from fnmatch import fnmatchcase
60from typing import IO, TYPE_CHECKING, Any, cast
62import pyarrow as pa
63import pyarrow.parquet as pq
65from lsst.daf.butler import DatasetProvenance, FormatterV2
66from lsst.daf.butler.delegates.arrowtable import _add_arrow_provenance, _checkArrowCompatibleType
67from lsst.resources import ResourcePath
68from lsst.utils.introspection import get_full_type_name
69from lsst.utils.iteration import ensure_iterable
71log = logging.getLogger(__name__)
73if TYPE_CHECKING:
74 import astropy.table as atable
75 import numpy as np
76 import pandas as pd
78 try:
79 import fsspec
80 from fsspec.spec import AbstractFileSystem
81 except ImportError:
82 fsspec = None
83 AbstractFileSystem = type
85TARGET_ROW_GROUP_BYTES = 1_000_000_000
86ASTROPY_PANDAS_INDEX_KEY = "lsst::arrow::astropy_pandas_index"
89@contextlib.contextmanager
90def generic_open(path: str, fs: AbstractFileSystem | None) -> Generator[IO]:
91 if fs is None:
92 with open(path, "rb") as fh:
93 yield fh
94 else:
95 with fs.open(path) as fh:
96 yield fh
99class ParquetFormatter(FormatterV2):
100 """Interface for reading and writing Arrow Table objects to and from
101 Parquet files.
102 """
104 default_extension = ".parq"
105 can_read_from_uri = True
106 can_read_from_local_file = True
108 def can_accept(self, in_memory_dataset: Any) -> bool:
109 # Docstring inherited.
110 return _checkArrowCompatibleType(in_memory_dataset) is not None
112 def read_from_uri(self, uri: ResourcePath, component: str | None = None, expected_size: int = -1) -> Any:
113 # Docstring inherited from Formatter.read.
114 try:
115 fs, path = uri.to_fsspec()
116 except ImportError:
117 log.debug("fsspec not available; falling back to local file access.")
118 # This signals to the formatter to use the read_from_local_file
119 # code path.
120 return NotImplemented
122 return self._read_parquet(path=path, fs=fs, component=component, expected_size=expected_size)
124 def read_from_local_file(self, path: str, component: str | None = None, expected_size: int = -1) -> Any:
125 # Docstring inherited from Formatter.read.
126 return self._read_parquet(path=path, component=component, expected_size=expected_size)
128 def _read_parquet(
129 self,
130 path: str,
131 fs: AbstractFileSystem | None = None,
132 component: str | None = None,
133 expected_size: int = -1,
134 ) -> Any:
135 with generic_open(path, fs) as handle:
136 schema = pq.read_schema(handle)
138 schema_names = ["ArrowSchema", "DataFrameSchema", "ArrowAstropySchema", "ArrowNumpySchema"]
140 if (
141 component in ("columns", "schema")
142 or self.file_descriptor.readStorageClass.name in schema_names
143 ):
144 # The schema will be translated to column format
145 # depending on the input type.
146 return schema
147 elif component == "rowcount":
148 # Get the rowcount from the metadata if possible, otherwise
149 # count.
150 if b"lsst::arrow::rowcount" in schema.metadata:
151 return int(schema.metadata[b"lsst::arrow::rowcount"])
153 temp_table = pq.read_table(
154 handle,
155 columns=[schema.names[0]],
156 use_threads=False,
157 use_pandas_metadata=False,
158 schema=schema,
159 )
160 return temp_table.num_rows
162 par_columns = None
163 strip_astropy_meta_yaml = True
164 if self.file_descriptor.parameters:
165 par_columns = self.file_descriptor.parameters.pop("columns", None)
166 if par_columns:
167 has_pandas_multi_index = False
168 if schema.metadata and b"pandas" in schema.metadata:
169 md = json.loads(schema.metadata[b"pandas"])
170 if len(md["column_indexes"]) > 1:
171 has_pandas_multi_index = True
173 if not has_pandas_multi_index:
174 # Ensure uniqueness, keeping order.
175 par_columns_in = list(dict.fromkeys(ensure_iterable(par_columns)))
176 file_columns = [name for name in schema.names if not name.startswith("__")]
178 # Do case-sensitive glob-style matching, again ensuring
179 # uniqueness and ordering.
180 par_columns = {}
181 for par_column in par_columns_in:
182 found = False
183 for file_column in file_columns:
184 if fnmatchcase(file_column, par_column):
185 found = True
186 par_columns[file_column] = True
187 if not found:
188 raise ValueError(
189 f"""Column {par_column} specified in parameters not available in """
190 """parquet file."""
191 )
192 par_columns = list(par_columns.keys())
193 else:
194 par_columns = _standardize_multi_index_columns(
195 arrow_schema_to_pandas_index(schema),
196 par_columns,
197 )
199 strip_astropy_meta_yaml = self.file_descriptor.parameters.pop(
200 "strip_astropy_meta_yaml",
201 True,
202 )
204 if len(self.file_descriptor.parameters): 204 ↛ 205line 204 didn't jump to line 205 because the condition on line 204 was never true
205 raise ValueError(
206 f"Unsupported parameters {self.file_descriptor.parameters} in ArrowTable read."
207 )
209 metadata = schema.metadata if schema.metadata is not None else {}
210 arrow_table = pq.read_table(
211 handle,
212 columns=par_columns,
213 use_threads=False,
214 use_pandas_metadata=(b"pandas" in metadata),
215 schema=schema,
216 )
218 if strip_astropy_meta_yaml:
219 metadata = arrow_table.schema.metadata
220 # Only strip if (a) we have metadata; (b) it contains
221 # ``table_meta_yaml``; (c) it contains
222 # ``lsst::arrow::rowcount`` to avoid stripping data from pure
223 # astropy tables (not written by the butler).
224 if (
225 metadata
226 and metadata.pop(b"table_meta_yaml", None)
227 and b"lsst::arrow::rowcount" in metadata
228 ):
229 arrow_table = arrow_table.replace_schema_metadata(metadata)
231 return arrow_table
233 def add_provenance(self, in_memory_dataset: Any, provenance: DatasetProvenance | None = None) -> Any:
234 return _add_arrow_provenance(in_memory_dataset, self.dataset_ref, provenance)
236 def write_local_file(self, in_memory_dataset: Any, uri: ResourcePath) -> None:
237 """Serialize the in memory dataset to a local parquet file.
239 Parameters
240 ----------
241 in_memory_dataset : `typing.Any`
242 The Python object to serialize.
243 uri : `lsst.resources.ResourcePath`
244 The location to write the local file.
245 """
246 if isinstance(in_memory_dataset, pa.Schema):
247 pq.write_metadata(in_memory_dataset, uri.ospath)
248 return
250 type_string = _checkArrowCompatibleType(in_memory_dataset)
252 if type_string is None:
253 raise ValueError(
254 f"Unsupported type {get_full_type_name(in_memory_dataset)} of "
255 "inMemoryDataset for ParquetFormatter."
256 )
258 if type_string == "arrow":
259 arrow_table = in_memory_dataset
260 elif type_string == "astropy":
261 arrow_table = astropy_to_arrow(in_memory_dataset)
262 elif type_string == "numpy":
263 arrow_table = numpy_to_arrow(in_memory_dataset)
264 elif type_string == "numpydict":
265 arrow_table = numpy_dict_to_arrow(in_memory_dataset)
266 else:
267 arrow_table = pandas_to_arrow(in_memory_dataset)
269 row_group_size = compute_row_group_size(arrow_table.schema)
271 pq.write_table(arrow_table, uri.ospath, row_group_size=row_group_size)
274def arrow_to_pandas(arrow_table: pa.Table) -> pd.DataFrame:
275 """Convert a pyarrow table to a pandas DataFrame.
277 Parameters
278 ----------
279 arrow_table : `pyarrow.Table`
280 Input arrow table to convert. If the table has ``pandas`` metadata
281 in the schema it will be used in the construction of the
282 ``DataFrame``.
284 Returns
285 -------
286 dataframe : `pandas.DataFrame`
287 Converted pandas dataframe.
288 """
289 dataframe = arrow_table.to_pandas(use_threads=False, integer_object_nulls=True)
291 metadata = arrow_table.schema.metadata if arrow_table.schema.metadata is not None else {}
292 if (key := ASTROPY_PANDAS_INDEX_KEY.encode()) in metadata:
293 pandas_index = metadata[key].decode("UTF8")
294 if pandas_index in arrow_table.schema.names:
295 dataframe.set_index(pandas_index, inplace=True)
296 else:
297 log.warning(
298 "Index column ``%s`` not available for arrow table conversion to DataFrame",
299 pandas_index,
300 )
302 return dataframe
305def arrow_to_astropy(arrow_table: pa.Table) -> atable.Table:
306 """Convert a pyarrow table to an `astropy.table.Table`.
308 Parameters
309 ----------
310 arrow_table : `pyarrow.Table`
311 Input arrow table to convert. If the table has astropy unit
312 metadata in the schema it will be used in the construction
313 of the ``astropy.table.Table``.
315 Returns
316 -------
317 table : `astropy.table.Table`
318 Converted astropy table.
319 """
320 from astropy.table import Table
322 astropy_table = Table(arrow_to_numpy_dict(arrow_table))
324 _apply_astropy_metadata(astropy_table, arrow_table.schema)
326 if (key := ASTROPY_PANDAS_INDEX_KEY) in astropy_table.meta:
327 if astropy_table.meta[key] not in astropy_table.columns:
328 astropy_table.meta.pop(key)
330 return astropy_table
333def arrow_to_numpy(arrow_table: pa.Table) -> np.ndarray | np.ma.MaskedArray:
334 """Convert a pyarrow table to a structured numpy array.
336 Parameters
337 ----------
338 arrow_table : `pyarrow.Table`
339 Input arrow table.
341 Returns
342 -------
343 array : `numpy.ndarray` or `numpy.ma.MaskedArray` (N,)
344 Numpy array table with N rows and the same column names
345 as the input arrow table. Will be masked records if any values
346 in the table are null.
347 """
348 import numpy as np
350 numpy_dict = arrow_to_numpy_dict(arrow_table)
352 has_mask = False
353 dtype: list[tuple] = []
354 for name, col in numpy_dict.items():
355 if len(shape := numpy_dict[name].shape) <= 1:
356 dtype.append((name, col.dtype))
357 else:
358 dtype.append((name, (col.dtype, shape[1:])))
360 if not has_mask and isinstance(col, np.ma.MaskedArray):
361 has_mask = True
363 array: Any
365 if has_mask:
366 import numpy.ma.mrecords as mrecords
368 array = mrecords.fromarrays(list(numpy_dict.values()), dtype=dtype)
369 else:
370 array = np.rec.fromarrays(numpy_dict.values(), dtype=dtype)
371 return array
374def arrow_to_numpy_dict(arrow_table: pa.Table) -> dict[str, np.ndarray]:
375 """Convert a pyarrow table to a dict of numpy arrays.
377 Parameters
378 ----------
379 arrow_table : `pyarrow.Table`
380 Input arrow table.
382 Returns
383 -------
384 numpy_dict : `dict` [`str`, `numpy.ndarray`]
385 Dict with keys as the column names, values as the arrays.
386 """
387 import numpy as np
389 schema = arrow_table.schema
390 metadata = schema.metadata if schema.metadata is not None else {}
392 numpy_dict = {}
394 for name in schema.names:
395 t = schema.field(name).type
397 if arrow_table[name].null_count == 0:
398 # Regular non-masked column
399 col = arrow_table[name].to_numpy()
400 else:
401 # For a masked column, we need to ask arrow to fill the null
402 # values with an appropriately typed value before conversion.
403 # Then we apply the mask to get a masked array of the correct type.
404 null_value: Any
405 match t:
406 case t if t in (pa.float64(), pa.float32(), pa.float16()):
407 null_value = np.nan
408 case t if t in (pa.int64(), pa.int32(), pa.int16(), pa.int8()):
409 null_value = -1
410 case t if t in (pa.bool_(),):
411 null_value = True
412 case t if _is_string(t) or _is_binary(t):
413 null_value = ""
414 case _:
415 # This is the fallback for unsigned ints in particular.
416 null_value = 0
418 col = np.ma.MaskedArray(
419 data=arrow_table[name].fill_null(null_value).to_numpy(),
420 mask=arrow_table[name].is_null().to_numpy(),
421 fill_value=null_value,
422 )
424 if _is_string(t) or _is_binary(t):
425 col = col.astype(_arrow_string_to_numpy_dtype(schema, name, col))
426 elif isinstance(t, pa.FixedSizeListType):
427 if len(col) > 0:
428 col = np.stack(col)
429 else:
430 # this is an empty column, and needs to be coerced to type.
431 col = col.astype(t.value_type.to_pandas_dtype())
433 shape = _multidim_shape_from_metadata(metadata, t.list_size, name)
434 col = col.reshape((len(arrow_table), *shape))
436 numpy_dict[name] = col
438 return numpy_dict
441def _numpy_dict_to_numpy(numpy_dict: dict[str, np.ndarray]) -> np.ndarray:
442 """Convert a dict of numpy arrays to a structured numpy array.
444 Parameters
445 ----------
446 numpy_dict : `dict` [`str`, `numpy.ndarray`]
447 Dict with keys as the column names, values as the arrays.
449 Returns
450 -------
451 array : `numpy.ndarray` (N,)
452 Numpy array table with N rows and columns names from the dict keys.
453 """
454 return arrow_to_numpy(numpy_dict_to_arrow(numpy_dict))
457def _numpy_to_numpy_dict(np_array: np.ndarray) -> dict[str, np.ndarray]:
458 """Convert a structured numpy array to a dict of numpy arrays.
460 Parameters
461 ----------
462 np_array : `numpy.ndarray`
463 Input numpy array with multiple fields.
465 Returns
466 -------
467 numpy_dict : `dict` [`str`, `numpy.ndarray`]
468 Dict with keys as the column names, values as the arrays.
469 """
470 return arrow_to_numpy_dict(numpy_to_arrow(np_array))
473def numpy_to_arrow(np_array: np.ndarray) -> pa.Table:
474 """Convert a numpy array table to an arrow table.
476 Parameters
477 ----------
478 np_array : `numpy.ndarray`
479 Input numpy array with multiple fields.
481 Returns
482 -------
483 arrow_table : `pyarrow.Table`
484 Converted arrow table.
485 """
486 type_list = _numpy_dtype_to_arrow_types(np_array.dtype)
488 md = {}
489 md[b"lsst::arrow::rowcount"] = str(len(np_array))
491 names = np_array.dtype.names
492 if names is None: 492 ↛ 493line 492 didn't jump to line 493 because the condition on line 492 was never true
493 names = ()
495 for name in names:
496 _append_numpy_string_metadata(md, name, np_array.dtype[name])
497 _append_numpy_multidim_metadata(md, name, np_array.dtype[name])
499 schema = pa.schema(type_list, metadata=md)
501 arrays = _numpy_style_arrays_to_arrow_arrays(
502 np_array.dtype,
503 len(np_array),
504 np_array,
505 schema,
506 )
508 arrow_table = pa.Table.from_arrays(arrays, schema=schema)
510 return arrow_table
513def numpy_dict_to_arrow(numpy_dict: dict[str, np.ndarray]) -> pa.Table:
514 """Convert a dict of numpy arrays to an arrow table.
516 Parameters
517 ----------
518 numpy_dict : `dict` [`str`, `numpy.ndarray`]
519 Dict with keys as the column names, values as the arrays.
521 Returns
522 -------
523 arrow_table : `pyarrow.Table`
524 Converted arrow table.
526 Raises
527 ------
528 ValueError
529 Raised if columns in ``numpy_dict`` have unequal numbers of
530 rows.
531 """
532 dtype, rowcount = _numpy_dict_to_dtype(numpy_dict)
533 type_list = _numpy_dtype_to_arrow_types(dtype)
535 md = {}
536 md[b"lsst::arrow::rowcount"] = str(rowcount)
538 if dtype.names is not None: 538 ↛ 543line 538 didn't jump to line 543 because the condition on line 538 was always true
539 for name in dtype.names:
540 _append_numpy_string_metadata(md, name, dtype[name])
541 _append_numpy_multidim_metadata(md, name, dtype[name])
543 schema = pa.schema(type_list, metadata=md)
545 arrays = _numpy_style_arrays_to_arrow_arrays(
546 dtype,
547 rowcount,
548 numpy_dict,
549 schema,
550 )
552 arrow_table = pa.Table.from_arrays(arrays, schema=schema)
554 return arrow_table
557def astropy_to_arrow(astropy_table: atable.Table) -> pa.Table:
558 """Convert an astropy table to an arrow table.
560 Parameters
561 ----------
562 astropy_table : `astropy.table.Table`
563 Input astropy table.
565 Returns
566 -------
567 arrow_table : `pyarrow.Table`
568 Converted arrow table.
569 """
570 from astropy.table import meta
572 type_list = _numpy_dtype_to_arrow_types(astropy_table.dtype)
574 md = {}
575 md[b"lsst::arrow::rowcount"] = str(len(astropy_table))
577 if (key := ASTROPY_PANDAS_INDEX_KEY) in astropy_table.meta:
578 md[key.encode()] = astropy_table.meta[key]
580 for name in astropy_table.dtype.names:
581 _append_numpy_string_metadata(md, name, astropy_table.dtype[name])
582 _append_numpy_multidim_metadata(md, name, astropy_table.dtype[name])
584 meta_yaml = meta.get_yaml_from_table(astropy_table)
585 meta_yaml_str = "\n".join(meta_yaml)
586 md[b"table_meta_yaml"] = meta_yaml_str
588 # Convert type list to fields with metadata.
589 fields = []
590 for name, pa_type in type_list:
591 field_metadata = {}
592 if description := astropy_table[name].description:
593 field_metadata["description"] = description
594 if unit := astropy_table[name].unit:
595 field_metadata["unit"] = str(unit)
596 fields.append(
597 pa.field(
598 name,
599 pa_type,
600 metadata=field_metadata,
601 )
602 )
604 schema = pa.schema(fields, metadata=md)
606 arrays = _numpy_style_arrays_to_arrow_arrays(
607 astropy_table.dtype,
608 len(astropy_table),
609 astropy_table,
610 schema,
611 )
613 arrow_table = pa.Table.from_arrays(arrays, schema=schema)
615 return arrow_table
618def astropy_to_pandas(astropy_table: atable.Table, index: str | None = None) -> pd.DataFrame:
619 """Convert an astropy table to a pandas dataframe via arrow.
621 By going via arrow we avoid pandas masked column bugs (e.g.
622 https://github.com/pandas-dev/pandas/issues/58173)
624 Parameters
625 ----------
626 astropy_table : `astropy.table.Table`
627 Input astropy table.
628 index : `str`, optional
629 Name of column to set as index.
631 Returns
632 -------
633 dataframe : `pandas.DataFrame`
634 Output pandas dataframe.
635 """
636 index_requested = False
637 if (key := ASTROPY_PANDAS_INDEX_KEY) in astropy_table.meta:
638 _index = astropy_table.meta[key]
639 if _index not in astropy_table.columns:
640 log.warning(
641 "Index column ``%s`` not available for astropy table conversion to DataFrame",
642 _index,
643 )
644 _index = None
645 else:
646 index_requested = True
647 _index = index
649 dataframe = arrow_to_pandas(astropy_to_arrow(astropy_table))
651 # Set the index if we have a valid index name, and either the
652 # index was requested in the call to the function or the dataframe
653 # was not previously indexed with the call to arrow_to_pandas.
654 if isinstance(_index, str) and (index_requested or dataframe.index.name is None):
655 dataframe.set_index(_index, inplace=True)
656 elif _index and index_requested: 656 ↛ 657line 656 didn't jump to line 657 because the condition on line 656 was never true
657 raise RuntimeError("index must be a string or None.")
659 return dataframe
662def add_pandas_index_to_astropy(astropy_table: atable.Table, index: str) -> None:
663 """Add special metadata to an astropy table to indicate a pandas index.
665 Parameters
666 ----------
667 astropy_table : `astropy.table.Table`
668 Input astropy table.
669 index : `str`
670 Name of column for pandas to set as index, if read as DataFrame.
671 """
672 if index not in astropy_table.columns:
673 raise ValueError("Column ``%s`` not in astropy table columns to use as pandas index.", index)
674 astropy_table.meta[ASTROPY_PANDAS_INDEX_KEY] = index
677def _astropy_to_numpy_dict(astropy_table: atable.Table) -> dict[str, np.ndarray]:
678 """Convert an astropy table to an arrow table.
680 Parameters
681 ----------
682 astropy_table : `astropy.table.Table`
683 Input astropy table.
685 Returns
686 -------
687 numpy_dict : `dict` [`str`, `numpy.ndarray`]
688 Dict with keys as the column names, values as the arrays.
689 """
690 return arrow_to_numpy_dict(astropy_to_arrow(astropy_table))
693def pandas_to_arrow(dataframe: pd.DataFrame, default_length: int = 10) -> pa.Table:
694 """Convert a pandas dataframe to an arrow table.
696 Parameters
697 ----------
698 dataframe : `pandas.DataFrame`
699 Input pandas dataframe.
700 default_length : `int`, optional
701 Default string length when not in metadata or can be inferred
702 from column.
704 Returns
705 -------
706 arrow_table : `pyarrow.Table`
707 Converted arrow table.
708 """
709 import pandas as pd
711 old_index = None
713 if isinstance(dataframe.index, pd.RangeIndex) and dataframe.index.name is not None: 713 ↛ 716line 713 didn't jump to line 716 because the condition on line 713 was never true
714 # Turn the RangeIndex into a regular index, or it won't serialize
715 # interoperably via arrow.
716 old_index = dataframe.index
718 dataframe.index = pd.Index(dataframe.index.to_numpy(), name=dataframe.index.name)
720 try:
721 arrow_table = pa.Table.from_pandas(dataframe)
722 except pa.ArrowInvalid as e:
723 msg = "; ".join(e.args)
724 msg += "; This is usually because the column is mixed type or has uneven length rows."
725 e.add_note(msg)
726 raise
727 finally:
728 if old_index is not None: 728 ↛ 729line 728 didn't jump to line 729 because the condition on line 728 was never true
729 dataframe.index = old_index
731 # Update the metadata
732 md = arrow_table.schema.metadata
734 md[b"lsst::arrow::rowcount"] = str(arrow_table.num_rows)
736 # We loop through the arrow table columns because the datatypes have
737 # been checked and converted from pandas objects.
738 for name in arrow_table.column_names:
739 if not name.startswith("__") and arrow_table[name].type == pa.string():
740 if len(arrow_table[name]) > 0: 740 ↛ 743line 740 didn't jump to line 743 because the condition on line 740 was always true
741 strlen = max(len(row.as_py()) for row in arrow_table[name] if row.is_valid)
742 else:
743 strlen = default_length
744 md[f"lsst::arrow::len::{name}".encode()] = str(strlen)
746 arrow_table = arrow_table.replace_schema_metadata(md)
748 return arrow_table
751def pandas_to_astropy(dataframe: pd.DataFrame) -> atable.Table:
752 """Convert a pandas dataframe to an astropy table, preserving indexes.
754 Parameters
755 ----------
756 dataframe : `pandas.DataFrame`
757 Input pandas dataframe.
759 Returns
760 -------
761 astropy_table : `astropy.table.Table`
762 Converted astropy table.
763 """
764 import pandas as pd
766 if isinstance(dataframe.columns, pd.MultiIndex):
767 raise ValueError("Cannot convert a multi-index dataframe to an astropy table.")
769 return arrow_to_astropy(pandas_to_arrow(dataframe))
772def pandas_to_numpy(dataframe: pd.DataFrame) -> np.ndarray | np.ma.MaskedArray:
773 """Convert a pandas dataframe to a numpy recarray.
775 Parameters
776 ----------
777 dataframe : `pandas.DataFrame`
778 Input pandas dataframe.
780 Returns
781 -------
782 array : `numpy.ndarray` or `numpy.ma.MaskedArray` (N,)
783 Numpy array table with N rows and the same column names
784 as the input dataframe. Will be masked records if any values
785 in the table are null.
786 """
787 # This conversion ensures strings are handled properly.
788 return arrow_to_numpy(pandas_to_arrow(dataframe))
791def _pandas_to_numpy_dict(dataframe: pd.DataFrame) -> dict[str, np.ndarray]:
792 """Convert a pandas dataframe to an dict of numpy arrays.
794 Parameters
795 ----------
796 dataframe : `pandas.DataFrame`
797 Input pandas dataframe.
799 Returns
800 -------
801 numpy_dict : `dict` [`str`, `numpy.ndarray`]
802 Dict with keys as the column names, values as the arrays.
803 """
804 return arrow_to_numpy_dict(pandas_to_arrow(dataframe))
807def numpy_to_astropy(np_array: np.ndarray) -> atable.Table:
808 """Convert a numpy table to an astropy table.
810 Parameters
811 ----------
812 np_array : `numpy.ndarray`
813 Input numpy array with multiple fields.
815 Returns
816 -------
817 astropy_table : `astropy.table.Table`
818 Converted astropy table.
819 """
820 from astropy.table import Table
822 return Table(data=np_array, copy=False)
825def arrow_schema_to_pandas_index(schema: pa.Schema) -> pd.Index | pd.MultiIndex:
826 """Convert an arrow schema to a pandas index/multiindex.
828 Parameters
829 ----------
830 schema : `pyarrow.Schema`
831 Input pyarrow schema.
833 Returns
834 -------
835 index : `pandas.Index` or `pandas.MultiIndex`
836 Converted pandas index.
837 """
838 import pandas as pd
840 if b"pandas" in schema.metadata:
841 md = json.loads(schema.metadata[b"pandas"])
842 indexes = md["column_indexes"]
843 len_indexes = len(indexes)
844 else:
845 len_indexes = 0
847 if len_indexes <= 1:
848 return pd.Index(name for name in schema.names if not name.startswith("__"))
849 else:
850 raw_columns = _split_multi_index_column_names(len(indexes), schema.names)
851 return pd.MultiIndex.from_tuples(raw_columns, names=[f["name"] for f in indexes])
854def arrow_schema_to_column_list(schema: pa.Schema) -> list[str]:
855 """Convert an arrow schema to a list of string column names.
857 Parameters
858 ----------
859 schema : `pyarrow.Schema`
860 Input pyarrow schema.
862 Returns
863 -------
864 column_list : `list` [`str`]
865 Converted list of column names.
866 """
867 return list(schema.names)
870class DataFrameSchema:
871 """Wrapper class for a schema for a pandas DataFrame.
873 Parameters
874 ----------
875 dataframe : `pandas.DataFrame`
876 Dataframe to turn into a schema.
877 """
879 def __init__(self, dataframe: pd.DataFrame) -> None:
880 import pandas as pd
882 self._schema = dataframe.loc[[False] * len(dataframe)]
884 if isinstance(self._schema.index, pd.RangeIndex) and self._schema.index.name is not None: 884 ↛ 887line 884 didn't jump to line 887 because the condition on line 884 was never true
885 # Turn the RangeIndex into a regular index or it won't
886 # give us all the columns via arrow.
887 self._schema.index = pd.Index(self._schema.index.to_numpy(), name=self._schema.index.name)
889 @classmethod
890 def from_arrow(cls, schema: pa.Schema) -> DataFrameSchema:
891 """Convert an arrow schema into a `DataFrameSchema`.
893 Parameters
894 ----------
895 schema : `pyarrow.Schema`
896 The pyarrow schema to convert.
898 Returns
899 -------
900 dataframe_schema : `DataFrameSchema`
901 Converted dataframe schema.
902 """
903 empty_table = pa.Table.from_pylist([] * len(schema.names), schema=schema)
905 return cls(empty_table.to_pandas())
907 def to_arrow_schema(self) -> pa.Schema:
908 """Convert to an arrow schema.
910 Returns
911 -------
912 arrow_schema : `pyarrow.Schema`
913 Converted pyarrow schema.
914 """
915 arrow_table = pa.Table.from_pandas(self._schema)
917 return arrow_table.schema
919 def to_arrow_numpy_schema(self) -> ArrowNumpySchema:
920 """Convert to an `ArrowNumpySchema`.
922 Returns
923 -------
924 arrow_numpy_schema : `ArrowNumpySchema`
925 Converted arrow numpy schema.
926 """
927 return ArrowNumpySchema.from_arrow(self.to_arrow_schema())
929 def to_arrow_astropy_schema(self) -> ArrowAstropySchema:
930 """Convert to an ArrowAstropySchema.
932 Returns
933 -------
934 arrow_astropy_schema : `ArrowAstropySchema`
935 Converted arrow astropy schema.
936 """
937 return ArrowAstropySchema.from_arrow(self.to_arrow_schema())
939 @property
940 def schema(self) -> np.dtype:
941 return self._schema
943 def __repr__(self) -> str:
944 return repr(self._schema)
946 def __eq__(self, other: object) -> bool:
947 if not isinstance(other, DataFrameSchema):
948 return NotImplemented
950 return self._schema.equals(other._schema)
953class ArrowAstropySchema:
954 """Wrapper class for a schema for an astropy table.
956 Parameters
957 ----------
958 astropy_table : `astropy.table.Table`
959 Input astropy table.
960 """
962 def __init__(self, astropy_table: atable.Table) -> None:
963 self._schema = astropy_table[:0]
965 @classmethod
966 def from_arrow(cls, schema: pa.Schema) -> ArrowAstropySchema:
967 """Convert an arrow schema into a ArrowAstropySchema.
969 Parameters
970 ----------
971 schema : `pyarrow.Schema`
972 Input pyarrow schema.
974 Returns
975 -------
976 astropy_schema : `ArrowAstropySchema`
977 Converted arrow astropy schema.
978 """
979 import numpy as np
980 from astropy.table import Table
982 dtype = _schema_to_dtype_list(schema)
984 data = np.zeros(0, dtype=dtype)
985 astropy_table = Table(data=data)
987 _apply_astropy_metadata(astropy_table, schema)
989 return cls(astropy_table)
991 def to_arrow_schema(self) -> pa.Schema:
992 """Convert to an arrow schema.
994 Returns
995 -------
996 arrow_schema : `pyarrow.Schema`
997 Converted pyarrow schema.
998 """
999 return astropy_to_arrow(self._schema).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 return DataFrameSchema.from_arrow(astropy_to_arrow(self._schema).schema)
1011 def to_arrow_numpy_schema(self) -> ArrowNumpySchema:
1012 """Convert to an `ArrowNumpySchema`.
1014 Returns
1015 -------
1016 arrow_numpy_schema : `ArrowNumpySchema`
1017 Converted arrow numpy schema.
1018 """
1019 return ArrowNumpySchema.from_arrow(astropy_to_arrow(self._schema).schema)
1021 @property
1022 def schema(self) -> atable.Table:
1023 return self._schema
1025 def __repr__(self) -> str:
1026 return repr(self._schema)
1028 def __eq__(self, other: object) -> bool:
1029 if not isinstance(other, ArrowAstropySchema):
1030 return NotImplemented
1032 # If this comparison passes then the two tables have the
1033 # same column names.
1034 if self._schema.dtype != other._schema.dtype:
1035 return False
1037 for name in self._schema.columns:
1038 if not self._schema[name].unit == other._schema[name].unit:
1039 return False
1040 if not self._schema[name].description == other._schema[name].description:
1041 return False
1042 if not self._schema[name].format == other._schema[name].format:
1043 return False
1045 return True
1048class ArrowNumpySchema:
1049 """Wrapper class for a schema for a numpy ndarray.
1051 Parameters
1052 ----------
1053 numpy_dtype : `numpy.dtype`
1054 Numpy dtype to convert.
1055 """
1057 def __init__(self, numpy_dtype: np.dtype) -> None:
1058 self._dtype = numpy_dtype
1060 @classmethod
1061 def from_arrow(cls, schema: pa.Schema) -> ArrowNumpySchema:
1062 """Convert an arrow schema into an `ArrowNumpySchema`.
1064 Parameters
1065 ----------
1066 schema : `pyarrow.Schema`
1067 Pyarrow schema to convert.
1069 Returns
1070 -------
1071 numpy_schema : `ArrowNumpySchema`
1072 Converted arrow numpy schema.
1073 """
1074 import numpy as np
1076 dtype = _schema_to_dtype_list(schema)
1078 return cls(np.dtype(dtype))
1080 def to_arrow_astropy_schema(self) -> ArrowAstropySchema:
1081 """Convert to an `ArrowAstropySchema`.
1083 Returns
1084 -------
1085 astropy_schema : `ArrowAstropySchema`
1086 Converted arrow astropy schema.
1087 """
1088 import numpy as np
1090 return ArrowAstropySchema.from_arrow(numpy_to_arrow(np.zeros(0, dtype=self._dtype)).schema)
1092 def to_dataframe_schema(self) -> DataFrameSchema:
1093 """Convert to a `DataFrameSchema`.
1095 Returns
1096 -------
1097 dataframe_schema : `DataFrameSchema`
1098 Converted dataframe schema.
1099 """
1100 import numpy as np
1102 return DataFrameSchema.from_arrow(numpy_to_arrow(np.zeros(0, dtype=self._dtype)).schema)
1104 def to_arrow_schema(self) -> pa.Schema:
1105 """Convert to a `pyarrow.Schema`.
1107 Returns
1108 -------
1109 arrow_schema : `pyarrow.Schema`
1110 Converted pyarrow schema.
1111 """
1112 import numpy as np
1114 return numpy_to_arrow(np.zeros(0, dtype=self._dtype)).schema
1116 @property
1117 def schema(self) -> np.dtype:
1118 return self._dtype
1120 def __repr__(self) -> str:
1121 return repr(self._dtype)
1123 def __eq__(self, other: object) -> bool:
1124 if not isinstance(other, ArrowNumpySchema):
1125 return NotImplemented
1127 if not self._dtype == other._dtype:
1128 return False
1130 return True
1133def _split_multi_index_column_names(n: int, names: Iterable[str]) -> list[Sequence[str]]:
1134 """Split a string that represents a multi-index column.
1136 PyArrow maps Pandas' multi-index column names (which are tuples in Python)
1137 to flat strings on disk. This routine exists to reconstruct the original
1138 tuple.
1140 Parameters
1141 ----------
1142 n : `int`
1143 Number of levels in the `pandas.MultiIndex` that is being
1144 reconstructed.
1145 names : `~collections.abc.Iterable` [`str`]
1146 Strings to be split.
1148 Returns
1149 -------
1150 column_names : `list` [`tuple` [`str`]]
1151 A list of multi-index column name tuples.
1152 """
1153 column_names: list[Sequence[str]] = []
1155 pattern = re.compile(r"\({}\)".format(", ".join(["'(.*)'"] * n)))
1156 for name in names:
1157 m = re.search(pattern, name)
1158 if m is not None:
1159 column_names.append(m.groups())
1161 return column_names
1164def _standardize_multi_index_columns(
1165 pd_index: pd.MultiIndex,
1166 columns: Any,
1167 stringify: bool = True,
1168) -> list[str | Sequence[Any]]:
1169 """Transform a dictionary/iterable index from a multi-index column list
1170 into a string directly understandable by PyArrow.
1172 Parameters
1173 ----------
1174 pd_index : `pandas.MultiIndex`
1175 Pandas multi-index.
1176 columns : `list` [`tuple`] or `dict` [`str`, `str` or `list` [`str`]]
1177 Columns to standardize.
1178 stringify : `bool`, optional
1179 Should the column names be stringified?
1181 Returns
1182 -------
1183 names : `list` [`str`]
1184 Stringified representation of a multi-index column name.
1185 """
1186 index_level_names = tuple(pd_index.names)
1188 names: list[str | Sequence[Any]] = []
1190 if isinstance(columns, list):
1191 for requested in columns:
1192 if not isinstance(requested, tuple):
1193 raise ValueError(
1194 "Columns parameter for multi-index data frame must be a dictionary or list of tuples. "
1195 f"Instead got a {get_full_type_name(requested)}."
1196 )
1197 if stringify:
1198 names.append(str(requested))
1199 else:
1200 names.append(requested)
1201 else:
1202 if not isinstance(columns, collections.abc.Mapping): 1202 ↛ 1203line 1202 didn't jump to line 1203 because the condition on line 1202 was never true
1203 raise ValueError(
1204 "Columns parameter for multi-index data frame must be a dictionary or list of tuples. "
1205 f"Instead got a {get_full_type_name(columns)}."
1206 )
1207 if not set(index_level_names).issuperset(columns.keys()): 1207 ↛ 1208line 1207 didn't jump to line 1208 because the condition on line 1207 was never true
1208 raise ValueError(
1209 f"Cannot use dict with keys {set(columns.keys())} to select columns from {index_level_names}."
1210 )
1211 factors = [
1212 ensure_iterable(columns.get(level, pd_index.levels[i]))
1213 for i, level in enumerate(index_level_names)
1214 ]
1215 for requested in itertools.product(*factors):
1216 for i, value in enumerate(requested):
1217 if value not in pd_index.levels[i]: 1217 ↛ 1218line 1217 didn't jump to line 1218 because the condition on line 1217 was never true
1218 raise ValueError(f"Unrecognized value {value!r} for index {index_level_names[i]!r}.")
1219 if stringify:
1220 names.append(str(requested))
1221 else:
1222 names.append(requested)
1224 return names
1227def _apply_astropy_metadata(astropy_table: atable.Table, arrow_schema: pa.Schema) -> None:
1228 """Apply any astropy metadata from the schema metadata.
1230 Parameters
1231 ----------
1232 astropy_table : `astropy.table.Table`
1233 Table to apply metadata.
1234 arrow_schema : `pyarrow.Schema`
1235 Arrow schema with metadata.
1236 """
1237 from astropy.table import meta
1239 metadata = arrow_schema.metadata if arrow_schema.metadata is not None else {}
1241 # Check if we have a special astropy metadata header yaml.
1242 meta_yaml = metadata.get(b"table_meta_yaml", None)
1243 if meta_yaml:
1244 meta_yaml = meta_yaml.decode("UTF8").split("\n")
1245 meta_hdr = meta.get_header_from_yaml(meta_yaml)
1247 # Set description, format, unit, meta from the column
1248 # metadata that was serialized with the table.
1249 header_cols = {x["name"]: x for x in meta_hdr["datatype"]}
1250 for col in astropy_table.columns.values():
1251 for attr in ("description", "format", "unit", "meta"):
1252 if attr in header_cols[col.name]:
1253 setattr(col, attr, header_cols[col.name][attr])
1255 if "meta" in meta_hdr:
1256 astropy_table.meta.update(meta_hdr["meta"])
1257 else:
1258 # If we don't have astropy header data, we may have arrow field
1259 # metadata.
1260 for name in arrow_schema.names:
1261 field_metadata = arrow_schema.field(name).metadata
1262 if field_metadata is None:
1263 continue
1264 if (
1265 b"description" in field_metadata
1266 and (description := field_metadata[b"description"].decode("UTF-8")) != ""
1267 ):
1268 astropy_table[name].description = description
1269 if b"unit" in field_metadata and (unit := field_metadata[b"unit"].decode("UTF-8")) != "":
1270 astropy_table[name].unit = unit
1272 # Ensure that the special ASTROPY_PANDAS_INDEX_KEY is propagated to
1273 # the table metadata.
1274 if index_key := metadata.get(ASTROPY_PANDAS_INDEX_KEY.encode(), None):
1275 astropy_table.meta[ASTROPY_PANDAS_INDEX_KEY] = index_key.decode("UTF-8")
1278def _arrow_string_to_numpy_dtype(
1279 schema: pa.Schema, name: str, numpy_column: np.ndarray | None = None, default_length: int = 10
1280) -> str:
1281 """Get the numpy dtype string associated with an arrow column.
1283 Parameters
1284 ----------
1285 schema : `pyarrow.Schema`
1286 Arrow table schema.
1287 name : `str`
1288 Column name.
1289 numpy_column : `numpy.ndarray`, optional
1290 Column to determine numpy string dtype.
1291 default_length : `int`, optional
1292 Default string length when not in metadata or can be inferred
1293 from column.
1295 Returns
1296 -------
1297 dtype_str : `str`
1298 Numpy dtype string.
1299 """
1300 # Special-case for string and binary columns
1301 md_name = f"lsst::arrow::len::{name}"
1302 strlen = default_length
1303 metadata = schema.metadata if schema.metadata is not None else {}
1304 if (encoded := md_name.encode("UTF-8")) in metadata:
1305 # String/bytes length from header.
1306 strlen = int(schema.metadata[encoded])
1307 elif numpy_column is not None and len(numpy_column) > 0:
1308 lengths = [len(row) for row in numpy_column if row]
1309 strlen = max(lengths) if lengths else 0
1311 dtype = f"U{strlen}" if _is_string(schema.field(name).type) else f"|S{strlen}"
1313 return dtype
1316def _append_numpy_string_metadata(metadata: dict[bytes, str], name: str, dtype: np.dtype) -> None:
1317 """Append numpy string length keys to arrow metadata.
1319 All column types are handled, but the metadata is only modified for
1320 string and byte columns.
1322 Parameters
1323 ----------
1324 metadata : `dict` [`bytes`, `str`]
1325 Metadata dictionary; modified in place.
1326 name : `str`
1327 Column name.
1328 dtype : `np.dtype`
1329 Numpy dtype.
1330 """
1331 import numpy as np
1333 if dtype.type is np.str_:
1334 metadata[f"lsst::arrow::len::{name}".encode()] = str(dtype.itemsize // 4)
1335 metadata[f"table::len::{name}".encode()] = str(dtype.itemsize // 4)
1336 elif dtype.type is np.bytes_:
1337 metadata[f"lsst::arrow::len::{name}".encode()] = str(dtype.itemsize)
1338 metadata[f"table::len::{name}".encode()] = str(dtype.itemsize)
1341def _append_numpy_multidim_metadata(metadata: dict[bytes, str], name: str, dtype: np.dtype) -> None:
1342 """Append numpy multi-dimensional shapes to arrow metadata.
1344 All column types are handled, but the metadata is only modified for
1345 multi-dimensional columns.
1347 Parameters
1348 ----------
1349 metadata : `dict` [`bytes`, `str`]
1350 Metadata dictionary; modified in place.
1351 name : `str`
1352 Column name.
1353 dtype : `np.dtype`
1354 Numpy dtype.
1355 """
1356 if len(dtype.shape) > 1:
1357 metadata[f"lsst::arrow::shape::{name}".encode()] = str(dtype.shape)
1360def _multidim_shape_from_metadata(metadata: dict[bytes, bytes], list_size: int, name: str) -> tuple[int, ...]:
1361 """Retrieve the shape from the metadata, if available.
1363 Parameters
1364 ----------
1365 metadata : `dict` [`bytes`, `bytes`]
1366 Metadata dictionary.
1367 list_size : `int`
1368 Size of the list datatype.
1369 name : `str`
1370 Column name.
1372 Returns
1373 -------
1374 shape : `tuple` [`int`]
1375 Shape associated with the column.
1377 Raises
1378 ------
1379 RuntimeError
1380 Raised if metadata is found but has incorrect format.
1381 """
1382 md_name = f"lsst::arrow::shape::{name}"
1383 if (encoded := md_name.encode("UTF-8")) in metadata:
1384 groups = re.search(r"\((.*)\)", metadata[encoded].decode("UTF-8"))
1385 if groups is None: 1385 ↛ 1386line 1385 didn't jump to line 1386 because the condition on line 1385 was never true
1386 raise RuntimeError("Illegal value found in metadata.")
1387 shape = tuple(int(x) for x in groups[1].split(",") if x != "")
1388 else:
1389 shape = (list_size,)
1391 return shape
1394def _schema_to_dtype_list(schema: pa.Schema) -> list[tuple[str, tuple[Any] | str]]:
1395 """Convert a pyarrow schema to a numpy dtype.
1397 Parameters
1398 ----------
1399 schema : `pyarrow.Schema`
1400 Input pyarrow schema.
1402 Returns
1403 -------
1404 dtype_list: `list` [`tuple`]
1405 A list with name, type pairs.
1406 """
1407 metadata = schema.metadata if schema.metadata is not None else {}
1409 dtype: list[Any] = []
1410 for name in schema.names:
1411 t = schema.field(name).type
1412 if isinstance(t, pa.FixedSizeListType):
1413 shape = _multidim_shape_from_metadata(metadata, t.list_size, name)
1414 dtype.append((name, (t.value_type.to_pandas_dtype(), shape)))
1415 elif not (_is_string(t) or _is_binary(t)):
1416 dtype.append((name, t.to_pandas_dtype()))
1417 else:
1418 dtype.append((name, _arrow_string_to_numpy_dtype(schema, name)))
1420 return dtype
1423def _numpy_dtype_to_arrow_types(dtype: np.dtype) -> list[Any]:
1424 """Convert a numpy dtype to a list of arrow types.
1426 Parameters
1427 ----------
1428 dtype : `numpy.dtype`
1429 Numpy dtype to convert.
1431 Returns
1432 -------
1433 type_list : `list` [`object`]
1434 Converted list of arrow types.
1435 """
1436 from math import prod
1438 import numpy as np
1440 type_list: list[Any] = []
1441 if dtype.names is None: 1441 ↛ 1442line 1441 didn't jump to line 1442 because the condition on line 1441 was never true
1442 return type_list
1444 for name in dtype.names:
1445 dt = dtype[name]
1446 arrow_type: Any
1447 if len(dt.shape) > 0:
1448 arrow_type = pa.list_(
1449 pa.from_numpy_dtype(cast(tuple[np.dtype, tuple[int, ...]], dt.subdtype)[0].type),
1450 prod(dt.shape),
1451 )
1452 elif dt.type == np.datetime64:
1453 time_unit = "ns" if "ns" in dt.str else "us"
1454 # The pa.timestamp() is the correct datatype to round-trip
1455 # a numpy datetime64[ns] or datetime[us] array.
1456 arrow_type = pa.timestamp(time_unit)
1457 else:
1458 try:
1459 arrow_type = pa.from_numpy_dtype(dt.type)
1460 except pa.ArrowNotImplementedError as e:
1461 msg = f"Could not serialize column {name} (type {str(dt)}) to Parquet."
1462 if dt == np.dtype("O"): 1462 ↛ 1464line 1462 didn't jump to line 1464 because the condition on line 1462 was always true
1463 msg += " This is usually because the column is mixed type or has uneven length rows."
1464 e.add_note(msg)
1465 raise
1466 type_list.append((name, arrow_type))
1468 return type_list
1471def _numpy_dict_to_dtype(numpy_dict: dict[str, np.ndarray]) -> tuple[np.dtype, int]:
1472 """Extract equivalent table dtype from dict of numpy arrays.
1474 Parameters
1475 ----------
1476 numpy_dict : `dict` [`str`, `numpy.ndarray`]
1477 Dict with keys as the column names, values as the arrays.
1479 Returns
1480 -------
1481 dtype : `numpy.dtype`
1482 dtype of equivalent table.
1483 rowcount : `int`
1484 Number of rows in the table.
1486 Raises
1487 ------
1488 ValueError if columns in numpy_dict have unequal numbers of rows.
1489 """
1490 import numpy as np
1492 dtype_list: list[tuple] = []
1493 rowcount = 0
1494 for name, col in numpy_dict.items():
1495 if rowcount == 0:
1496 rowcount = len(col)
1497 if len(col) != rowcount:
1498 raise ValueError(f"Column {name} has a different number of rows.")
1499 if len(col.shape) == 1:
1500 dtype_list.append((name, col.dtype))
1501 else:
1502 dtype_list.append((name, (col.dtype, col.shape[1:])))
1503 dtype = np.dtype(dtype_list)
1505 return (dtype, rowcount)
1508def _numpy_style_arrays_to_arrow_arrays(
1509 dtype: np.dtype,
1510 rowcount: int,
1511 np_style_arrays: dict[str, np.ndarray] | np.ndarray | atable.Table,
1512 schema: pa.Schema,
1513) -> list[pa.Array]:
1514 """Convert numpy-style arrays to arrow arrays.
1516 Parameters
1517 ----------
1518 dtype : `numpy.dtype`
1519 Numpy dtype of input table/arrays.
1520 rowcount : `int`
1521 Number of rows in input table/arrays.
1522 np_style_arrays : `dict` [`str`, `np.ndarray`] or `np.ndarray`
1523 or `astropy.table.Table`
1524 Arrays to convert to arrow.
1525 schema : `pyarrow.Schema`
1526 Schema of arrow table.
1528 Returns
1529 -------
1530 arrow_arrays : `list` [`pyarrow.Array`]
1531 List of converted pyarrow arrays.
1532 """
1533 import numpy as np
1535 arrow_arrays: list[pa.Array] = []
1536 if dtype.names is None: 1536 ↛ 1537line 1536 didn't jump to line 1537 because the condition on line 1536 was never true
1537 return arrow_arrays
1539 for name in dtype.names:
1540 dt = dtype[name]
1541 val: Any
1542 if len(dt.shape) > 0:
1543 if rowcount > 0:
1544 val = np.split(np_style_arrays[name].ravel(), rowcount)
1545 else:
1546 val = []
1547 else:
1548 val = np_style_arrays[name]
1550 try:
1551 arrow_arrays.append(pa.array(val, type=schema.field(name).type))
1552 except pa.ArrowNotImplementedError as err:
1553 # Check if val is big-endian.
1554 if (np.little_endian and val.dtype.byteorder == ">") or ( 1554 ↛ 1563line 1554 didn't jump to line 1563 because the condition on line 1554 was always true
1555 not np.little_endian and val.dtype.byteorder == "="
1556 ):
1557 # We need to convert the array to little-endian.
1558 val2 = val.byteswap()
1559 val2.dtype = val2.dtype.newbyteorder("<")
1560 arrow_arrays.append(pa.array(val2, type=schema.field(name).type))
1561 else:
1562 # This failed for some other reason so raise the exception.
1563 raise err
1565 return arrow_arrays
1568def compute_row_group_size(schema: pa.Schema, target_size: int = TARGET_ROW_GROUP_BYTES) -> int:
1569 """Compute approximate row group size for a given arrow schema.
1571 Given a schema, this routine will compute the number of rows in a row group
1572 that targets the persisted size on disk (or smaller). The exact size on
1573 disk depends on the compression settings and ratios; typical binary data
1574 tables will have around 15-20% compression with the pyarrow default
1575 ``snappy`` compression algorithm.
1577 Parameters
1578 ----------
1579 schema : `pyarrow.Schema`
1580 Arrow table schema.
1581 target_size : `int`, optional
1582 The target size (in bytes).
1584 Returns
1585 -------
1586 row_group_size : `int`
1587 Number of rows per row group to hit the target size.
1588 """
1589 bit_width = 0
1591 metadata = schema.metadata if schema.metadata is not None else {}
1593 for name in schema.names:
1594 if name.startswith("__"):
1595 continue
1597 t = schema.field(name).type
1599 if _is_string(t) or _is_binary(t):
1600 md_name = f"lsst::arrow::len::{name}"
1602 if (encoded := md_name.encode("UTF-8")) in metadata:
1603 # String/bytes length from header.
1604 strlen = int(schema.metadata[encoded])
1605 else:
1606 # We don't know the string width, so guess something.
1607 strlen = 10
1609 # Assuming UTF-8 encoding, and very few wide characters.
1610 t_width = 8 * strlen
1611 elif isinstance(t, pa.FixedSizeListType):
1612 if t.value_type == pa.null(): 1612 ↛ 1613line 1612 didn't jump to line 1613 because the condition on line 1612 was never true
1613 t_width = 0
1614 else:
1615 t_width = t.list_size * t.value_type.bit_width
1616 elif t == pa.null():
1617 t_width = 0
1618 elif isinstance(t, pa.ListType):
1619 if t.value_type == pa.null():
1620 t_width = 0
1621 else:
1622 # This is a variable length list, just choose
1623 # something arbitrary.
1624 t_width = 10 * t.value_type.bit_width
1625 else:
1626 t_width = t.bit_width
1628 bit_width += t_width
1630 # Insist it is at least 1 byte wide to avoid any divide-by-zero errors.
1631 if bit_width < 8:
1632 bit_width = 8
1634 byte_width = bit_width // 8
1636 return target_size // byte_width
1639def _is_string(t: pa.DataType) -> bool:
1640 return pa.types.is_string(t) or pa.types.is_large_string(t) or pa.types.is_string_view(t)
1643def _is_binary(t: pa.DataType) -> bool:
1644 return pa.types.is_binary(t) or pa.types.is_large_binary(t) or pa.types.is_binary_view(t)