Coverage for python/lsst/daf/butler/formatters/parquet.py: 96%
580 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-08-05 01:26 -0700
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 component in ("columns", "schema") or self.file_descriptor.readStorageClass.name in schema_names:
141 # The schema will be translated to column format
142 # depending on the input type.
143 return schema
144 elif component == "rowcount":
145 # Get the rowcount from the metadata if possible, otherwise count.
146 if b"lsst::arrow::rowcount" in schema.metadata:
147 return int(schema.metadata[b"lsst::arrow::rowcount"])
149 with generic_open(path, fs) as handle:
150 temp_table = pq.read_table(
151 handle,
152 columns=[schema.names[0]],
153 use_threads=False,
154 use_pandas_metadata=False,
155 )
157 return len(temp_table[schema.names[0]])
159 par_columns = None
160 strip_astropy_meta_yaml = True
161 if self.file_descriptor.parameters:
162 par_columns = self.file_descriptor.parameters.pop("columns", None)
163 if par_columns:
164 has_pandas_multi_index = False
165 if schema.metadata and b"pandas" in schema.metadata:
166 md = json.loads(schema.metadata[b"pandas"])
167 if len(md["column_indexes"]) > 1:
168 has_pandas_multi_index = True
170 if not has_pandas_multi_index:
171 # Ensure uniqueness, keeping order.
172 par_columns_in = list(dict.fromkeys(ensure_iterable(par_columns)))
173 file_columns = [name for name in schema.names if not name.startswith("__")]
175 # Do case-sensitive glob-style matching, again ensuring
176 # uniqueness and ordering.
177 par_columns = {}
178 for par_column in par_columns_in:
179 found = False
180 for file_column in file_columns:
181 if fnmatchcase(file_column, par_column):
182 found = True
183 par_columns[file_column] = True
184 if not found:
185 raise ValueError(
186 f"Column {par_column} specified in parameters not available in parquet file."
187 )
188 par_columns = list(par_columns.keys())
189 else:
190 par_columns = _standardize_multi_index_columns(
191 arrow_schema_to_pandas_index(schema),
192 par_columns,
193 )
195 strip_astropy_meta_yaml = self.file_descriptor.parameters.pop(
196 "strip_astropy_meta_yaml",
197 True,
198 )
200 if len(self.file_descriptor.parameters): 200 ↛ 201line 200 didn't jump to line 201 because the condition on line 200 was never true
201 raise ValueError(
202 f"Unsupported parameters {self.file_descriptor.parameters} in ArrowTable read."
203 )
205 metadata = schema.metadata if schema.metadata is not None else {}
206 with generic_open(path, fs) as handle:
207 arrow_table = pq.read_table(
208 handle,
209 columns=par_columns,
210 use_threads=False,
211 use_pandas_metadata=(b"pandas" in metadata),
212 )
214 if strip_astropy_meta_yaml:
215 metadata = arrow_table.schema.metadata
216 # Only strip if (a) we have metadata; (b) it contains
217 # ``table_meta_yaml``; (c) it contains ``lsst::arrow::rowcount``
218 # to avoid stripping data from pure astropy tables (not written
219 # by the butler).
220 if metadata and metadata.pop(b"table_meta_yaml", None) and b"lsst::arrow::rowcount" in metadata:
221 arrow_table = arrow_table.replace_schema_metadata(metadata)
223 return arrow_table
225 def add_provenance(self, in_memory_dataset: Any, provenance: DatasetProvenance | None = None) -> Any:
226 return _add_arrow_provenance(in_memory_dataset, self.dataset_ref, provenance)
228 def write_local_file(self, in_memory_dataset: Any, uri: ResourcePath) -> None:
229 """Serialize the in memory dataset to a local parquet file.
231 Parameters
232 ----------
233 in_memory_dataset : `typing.Any`
234 The Python object to serialize.
235 uri : `lsst.resources.ResourcePath`
236 The location to write the local file.
237 """
238 if isinstance(in_memory_dataset, pa.Schema):
239 pq.write_metadata(in_memory_dataset, uri.ospath)
240 return
242 type_string = _checkArrowCompatibleType(in_memory_dataset)
244 if type_string is None:
245 raise ValueError(
246 f"Unsupported type {get_full_type_name(in_memory_dataset)} of "
247 "inMemoryDataset for ParquetFormatter."
248 )
250 if type_string == "arrow":
251 arrow_table = in_memory_dataset
252 elif type_string == "astropy":
253 arrow_table = astropy_to_arrow(in_memory_dataset)
254 elif type_string == "numpy":
255 arrow_table = numpy_to_arrow(in_memory_dataset)
256 elif type_string == "numpydict":
257 arrow_table = numpy_dict_to_arrow(in_memory_dataset)
258 else:
259 arrow_table = pandas_to_arrow(in_memory_dataset)
261 row_group_size = compute_row_group_size(arrow_table.schema)
263 pq.write_table(arrow_table, uri.ospath, row_group_size=row_group_size)
266def arrow_to_pandas(arrow_table: pa.Table) -> pd.DataFrame:
267 """Convert a pyarrow table to a pandas DataFrame.
269 Parameters
270 ----------
271 arrow_table : `pyarrow.Table`
272 Input arrow table to convert. If the table has ``pandas`` metadata
273 in the schema it will be used in the construction of the
274 ``DataFrame``.
276 Returns
277 -------
278 dataframe : `pandas.DataFrame`
279 Converted pandas dataframe.
280 """
281 dataframe = arrow_table.to_pandas(use_threads=False, integer_object_nulls=True)
283 metadata = arrow_table.schema.metadata if arrow_table.schema.metadata is not None else {}
284 if (key := ASTROPY_PANDAS_INDEX_KEY.encode()) in metadata:
285 pandas_index = metadata[key].decode("UTF8")
286 if pandas_index in arrow_table.schema.names:
287 dataframe.set_index(pandas_index, inplace=True)
288 else:
289 log.warning(
290 "Index column ``%s`` not available for arrow table conversion to DataFrame",
291 pandas_index,
292 )
294 return dataframe
297def arrow_to_astropy(arrow_table: pa.Table) -> atable.Table:
298 """Convert a pyarrow table to an `astropy.table.Table`.
300 Parameters
301 ----------
302 arrow_table : `pyarrow.Table`
303 Input arrow table to convert. If the table has astropy unit
304 metadata in the schema it will be used in the construction
305 of the ``astropy.table.Table``.
307 Returns
308 -------
309 table : `astropy.table.Table`
310 Converted astropy table.
311 """
312 from astropy.table import Table
314 astropy_table = Table(arrow_to_numpy_dict(arrow_table))
316 _apply_astropy_metadata(astropy_table, arrow_table.schema)
318 if (key := ASTROPY_PANDAS_INDEX_KEY) in astropy_table.meta:
319 if astropy_table.meta[key] not in astropy_table.columns:
320 astropy_table.meta.pop(key)
322 return astropy_table
325def arrow_to_numpy(arrow_table: pa.Table) -> np.ndarray | np.ma.MaskedArray:
326 """Convert a pyarrow table to a structured numpy array.
328 Parameters
329 ----------
330 arrow_table : `pyarrow.Table`
331 Input arrow table.
333 Returns
334 -------
335 array : `numpy.ndarray` or `numpy.ma.MaskedArray` (N,)
336 Numpy array table with N rows and the same column names
337 as the input arrow table. Will be masked records if any values
338 in the table are null.
339 """
340 import numpy as np
342 numpy_dict = arrow_to_numpy_dict(arrow_table)
344 has_mask = False
345 dtype: list[tuple] = []
346 for name, col in numpy_dict.items():
347 if len(shape := numpy_dict[name].shape) <= 1:
348 dtype.append((name, col.dtype))
349 else:
350 dtype.append((name, (col.dtype, shape[1:])))
352 if not has_mask and isinstance(col, np.ma.MaskedArray):
353 has_mask = True
355 array: Any
357 if has_mask:
358 import numpy.ma.mrecords as mrecords
360 array = mrecords.fromarrays(list(numpy_dict.values()), dtype=dtype)
361 else:
362 array = np.rec.fromarrays(numpy_dict.values(), dtype=dtype)
363 return array
366def arrow_to_numpy_dict(arrow_table: pa.Table) -> dict[str, np.ndarray]:
367 """Convert a pyarrow table to a dict of numpy arrays.
369 Parameters
370 ----------
371 arrow_table : `pyarrow.Table`
372 Input arrow table.
374 Returns
375 -------
376 numpy_dict : `dict` [`str`, `numpy.ndarray`]
377 Dict with keys as the column names, values as the arrays.
378 """
379 import numpy as np
381 schema = arrow_table.schema
382 metadata = schema.metadata if schema.metadata is not None else {}
384 numpy_dict = {}
386 for name in schema.names:
387 t = schema.field(name).type
389 if arrow_table[name].null_count == 0:
390 # Regular non-masked column
391 col = arrow_table[name].to_numpy()
392 else:
393 # For a masked column, we need to ask arrow to fill the null
394 # values with an appropriately typed value before conversion.
395 # Then we apply the mask to get a masked array of the correct type.
396 null_value: Any
397 match t:
398 case t if t in (pa.float64(), pa.float32(), pa.float16()):
399 null_value = np.nan
400 case t if t in (pa.int64(), pa.int32(), pa.int16(), pa.int8()):
401 null_value = -1
402 case t if t in (pa.bool_(),):
403 null_value = True
404 case t if _is_string(t) or _is_binary(t):
405 null_value = ""
406 case _:
407 # This is the fallback for unsigned ints in particular.
408 null_value = 0
410 col = np.ma.MaskedArray(
411 data=arrow_table[name].fill_null(null_value).to_numpy(),
412 mask=arrow_table[name].is_null().to_numpy(),
413 fill_value=null_value,
414 )
416 if _is_string(t) or _is_binary(t):
417 col = col.astype(_arrow_string_to_numpy_dtype(schema, name, col))
418 elif isinstance(t, pa.FixedSizeListType):
419 if len(col) > 0:
420 col = np.stack(col)
421 else:
422 # this is an empty column, and needs to be coerced to type.
423 col = col.astype(t.value_type.to_pandas_dtype())
425 shape = _multidim_shape_from_metadata(metadata, t.list_size, name)
426 col = col.reshape((len(arrow_table), *shape))
428 numpy_dict[name] = col
430 return numpy_dict
433def _numpy_dict_to_numpy(numpy_dict: dict[str, np.ndarray]) -> np.ndarray:
434 """Convert a dict of numpy arrays to a structured numpy array.
436 Parameters
437 ----------
438 numpy_dict : `dict` [`str`, `numpy.ndarray`]
439 Dict with keys as the column names, values as the arrays.
441 Returns
442 -------
443 array : `numpy.ndarray` (N,)
444 Numpy array table with N rows and columns names from the dict keys.
445 """
446 return arrow_to_numpy(numpy_dict_to_arrow(numpy_dict))
449def _numpy_to_numpy_dict(np_array: np.ndarray) -> dict[str, np.ndarray]:
450 """Convert a structured numpy array to a dict of numpy arrays.
452 Parameters
453 ----------
454 np_array : `numpy.ndarray`
455 Input numpy array with multiple fields.
457 Returns
458 -------
459 numpy_dict : `dict` [`str`, `numpy.ndarray`]
460 Dict with keys as the column names, values as the arrays.
461 """
462 return arrow_to_numpy_dict(numpy_to_arrow(np_array))
465def numpy_to_arrow(np_array: np.ndarray) -> pa.Table:
466 """Convert a numpy array table to an arrow table.
468 Parameters
469 ----------
470 np_array : `numpy.ndarray`
471 Input numpy array with multiple fields.
473 Returns
474 -------
475 arrow_table : `pyarrow.Table`
476 Converted arrow table.
477 """
478 type_list = _numpy_dtype_to_arrow_types(np_array.dtype)
480 md = {}
481 md[b"lsst::arrow::rowcount"] = str(len(np_array))
483 names = np_array.dtype.names
484 if names is None: 484 ↛ 485line 484 didn't jump to line 485 because the condition on line 484 was never true
485 names = ()
487 for name in names:
488 _append_numpy_string_metadata(md, name, np_array.dtype[name])
489 _append_numpy_multidim_metadata(md, name, np_array.dtype[name])
491 schema = pa.schema(type_list, metadata=md)
493 arrays = _numpy_style_arrays_to_arrow_arrays(
494 np_array.dtype,
495 len(np_array),
496 np_array,
497 schema,
498 )
500 arrow_table = pa.Table.from_arrays(arrays, schema=schema)
502 return arrow_table
505def numpy_dict_to_arrow(numpy_dict: dict[str, np.ndarray]) -> pa.Table:
506 """Convert a dict of numpy arrays to an arrow table.
508 Parameters
509 ----------
510 numpy_dict : `dict` [`str`, `numpy.ndarray`]
511 Dict with keys as the column names, values as the arrays.
513 Returns
514 -------
515 arrow_table : `pyarrow.Table`
516 Converted arrow table.
518 Raises
519 ------
520 ValueError
521 Raised if columns in ``numpy_dict`` have unequal numbers of
522 rows.
523 """
524 dtype, rowcount = _numpy_dict_to_dtype(numpy_dict)
525 type_list = _numpy_dtype_to_arrow_types(dtype)
527 md = {}
528 md[b"lsst::arrow::rowcount"] = str(rowcount)
530 if dtype.names is not None: 530 ↛ 535line 530 didn't jump to line 535 because the condition on line 530 was always true
531 for name in dtype.names:
532 _append_numpy_string_metadata(md, name, dtype[name])
533 _append_numpy_multidim_metadata(md, name, dtype[name])
535 schema = pa.schema(type_list, metadata=md)
537 arrays = _numpy_style_arrays_to_arrow_arrays(
538 dtype,
539 rowcount,
540 numpy_dict,
541 schema,
542 )
544 arrow_table = pa.Table.from_arrays(arrays, schema=schema)
546 return arrow_table
549def astropy_to_arrow(astropy_table: atable.Table) -> pa.Table:
550 """Convert an astropy table to an arrow table.
552 Parameters
553 ----------
554 astropy_table : `astropy.table.Table`
555 Input astropy table.
557 Returns
558 -------
559 arrow_table : `pyarrow.Table`
560 Converted arrow table.
561 """
562 from astropy.table import meta
564 type_list = _numpy_dtype_to_arrow_types(astropy_table.dtype)
566 md = {}
567 md[b"lsst::arrow::rowcount"] = str(len(astropy_table))
569 if (key := ASTROPY_PANDAS_INDEX_KEY) in astropy_table.meta:
570 md[key.encode()] = astropy_table.meta[key]
572 for name in astropy_table.dtype.names:
573 _append_numpy_string_metadata(md, name, astropy_table.dtype[name])
574 _append_numpy_multidim_metadata(md, name, astropy_table.dtype[name])
576 meta_yaml = meta.get_yaml_from_table(astropy_table)
577 meta_yaml_str = "\n".join(meta_yaml)
578 md[b"table_meta_yaml"] = meta_yaml_str
580 # Convert type list to fields with metadata.
581 fields = []
582 for name, pa_type in type_list:
583 field_metadata = {}
584 if description := astropy_table[name].description:
585 field_metadata["description"] = description
586 if unit := astropy_table[name].unit:
587 field_metadata["unit"] = str(unit)
588 fields.append(
589 pa.field(
590 name,
591 pa_type,
592 metadata=field_metadata,
593 )
594 )
596 schema = pa.schema(fields, metadata=md)
598 arrays = _numpy_style_arrays_to_arrow_arrays(
599 astropy_table.dtype,
600 len(astropy_table),
601 astropy_table,
602 schema,
603 )
605 arrow_table = pa.Table.from_arrays(arrays, schema=schema)
607 return arrow_table
610def astropy_to_pandas(astropy_table: atable.Table, index: str | None = None) -> pd.DataFrame:
611 """Convert an astropy table to a pandas dataframe via arrow.
613 By going via arrow we avoid pandas masked column bugs (e.g.
614 https://github.com/pandas-dev/pandas/issues/58173)
616 Parameters
617 ----------
618 astropy_table : `astropy.table.Table`
619 Input astropy table.
620 index : `str`, optional
621 Name of column to set as index.
623 Returns
624 -------
625 dataframe : `pandas.DataFrame`
626 Output pandas dataframe.
627 """
628 index_requested = False
629 if (key := ASTROPY_PANDAS_INDEX_KEY) in astropy_table.meta:
630 _index = astropy_table.meta[key]
631 if _index not in astropy_table.columns:
632 log.warning(
633 "Index column ``%s`` not available for astropy table conversion to DataFrame",
634 _index,
635 )
636 _index = None
637 else:
638 index_requested = True
639 _index = index
641 dataframe = arrow_to_pandas(astropy_to_arrow(astropy_table))
643 # Set the index if we have a valid index name, and either the
644 # index was requested in the call to the function or the dataframe
645 # was not previously indexed with the call to arrow_to_pandas.
646 if isinstance(_index, str) and (index_requested or dataframe.index.name is None):
647 dataframe.set_index(_index, inplace=True)
648 elif _index and index_requested: 648 ↛ 649line 648 didn't jump to line 649 because the condition on line 648 was never true
649 raise RuntimeError("index must be a string or None.")
651 return dataframe
654def add_pandas_index_to_astropy(astropy_table: atable.Table, index: str) -> None:
655 """Add special metadata to an astropy table to indicate a pandas index.
657 Parameters
658 ----------
659 astropy_table : `astropy.table.Table`
660 Input astropy table.
661 index : `str`
662 Name of column for pandas to set as index, if read as DataFrame.
663 """
664 if index not in astropy_table.columns:
665 raise ValueError("Column ``%s`` not in astropy table columns to use as pandas index.", index)
666 astropy_table.meta[ASTROPY_PANDAS_INDEX_KEY] = index
669def _astropy_to_numpy_dict(astropy_table: atable.Table) -> dict[str, np.ndarray]:
670 """Convert an astropy table to an arrow table.
672 Parameters
673 ----------
674 astropy_table : `astropy.table.Table`
675 Input astropy table.
677 Returns
678 -------
679 numpy_dict : `dict` [`str`, `numpy.ndarray`]
680 Dict with keys as the column names, values as the arrays.
681 """
682 return arrow_to_numpy_dict(astropy_to_arrow(astropy_table))
685def pandas_to_arrow(dataframe: pd.DataFrame, default_length: int = 10) -> pa.Table:
686 """Convert a pandas dataframe to an arrow table.
688 Parameters
689 ----------
690 dataframe : `pandas.DataFrame`
691 Input pandas dataframe.
692 default_length : `int`, optional
693 Default string length when not in metadata or can be inferred
694 from column.
696 Returns
697 -------
698 arrow_table : `pyarrow.Table`
699 Converted arrow table.
700 """
701 import pandas as pd
703 old_index = None
705 if isinstance(dataframe.index, pd.RangeIndex) and dataframe.index.name is not None: 705 ↛ 708line 705 didn't jump to line 708 because the condition on line 705 was never true
706 # Turn the RangeIndex into a regular index, or it won't serialize
707 # interoperably via arrow.
708 old_index = dataframe.index
710 dataframe.index = pd.Index(dataframe.index.to_numpy(), name=dataframe.index.name)
712 try:
713 arrow_table = pa.Table.from_pandas(dataframe)
714 except pa.ArrowInvalid as e:
715 msg = "; ".join(e.args)
716 msg += "; This is usually because the column is mixed type or has uneven length rows."
717 e.add_note(msg)
718 raise
719 finally:
720 if old_index is not None: 720 ↛ 721line 720 didn't jump to line 721 because the condition on line 720 was never true
721 dataframe.index = old_index
723 # Update the metadata
724 md = arrow_table.schema.metadata
726 md[b"lsst::arrow::rowcount"] = str(arrow_table.num_rows)
728 # We loop through the arrow table columns because the datatypes have
729 # been checked and converted from pandas objects.
730 for name in arrow_table.column_names:
731 if not name.startswith("__") and arrow_table[name].type == pa.string():
732 if len(arrow_table[name]) > 0: 732 ↛ 735line 732 didn't jump to line 735 because the condition on line 732 was always true
733 strlen = max(len(row.as_py()) for row in arrow_table[name] if row.is_valid)
734 else:
735 strlen = default_length
736 md[f"lsst::arrow::len::{name}".encode()] = str(strlen)
738 arrow_table = arrow_table.replace_schema_metadata(md)
740 return arrow_table
743def pandas_to_astropy(dataframe: pd.DataFrame) -> atable.Table:
744 """Convert a pandas dataframe to an astropy table, preserving indexes.
746 Parameters
747 ----------
748 dataframe : `pandas.DataFrame`
749 Input pandas dataframe.
751 Returns
752 -------
753 astropy_table : `astropy.table.Table`
754 Converted astropy table.
755 """
756 import pandas as pd
758 if isinstance(dataframe.columns, pd.MultiIndex):
759 raise ValueError("Cannot convert a multi-index dataframe to an astropy table.")
761 return arrow_to_astropy(pandas_to_arrow(dataframe))
764def pandas_to_numpy(dataframe: pd.DataFrame) -> np.ndarray | np.ma.MaskedArray:
765 """Convert a pandas dataframe to a numpy recarray.
767 Parameters
768 ----------
769 dataframe : `pandas.DataFrame`
770 Input pandas dataframe.
772 Returns
773 -------
774 array : `numpy.ndarray` or `numpy.ma.MaskedArray` (N,)
775 Numpy array table with N rows and the same column names
776 as the input dataframe. Will be masked records if any values
777 in the table are null.
778 """
779 # This conversion ensures strings are handled properly.
780 return arrow_to_numpy(pandas_to_arrow(dataframe))
783def _pandas_to_numpy_dict(dataframe: pd.DataFrame) -> dict[str, np.ndarray]:
784 """Convert a pandas dataframe to an dict of numpy arrays.
786 Parameters
787 ----------
788 dataframe : `pandas.DataFrame`
789 Input pandas dataframe.
791 Returns
792 -------
793 numpy_dict : `dict` [`str`, `numpy.ndarray`]
794 Dict with keys as the column names, values as the arrays.
795 """
796 return arrow_to_numpy_dict(pandas_to_arrow(dataframe))
799def numpy_to_astropy(np_array: np.ndarray) -> atable.Table:
800 """Convert a numpy table to an astropy table.
802 Parameters
803 ----------
804 np_array : `numpy.ndarray`
805 Input numpy array with multiple fields.
807 Returns
808 -------
809 astropy_table : `astropy.table.Table`
810 Converted astropy table.
811 """
812 from astropy.table import Table
814 return Table(data=np_array, copy=False)
817def arrow_schema_to_pandas_index(schema: pa.Schema) -> pd.Index | pd.MultiIndex:
818 """Convert an arrow schema to a pandas index/multiindex.
820 Parameters
821 ----------
822 schema : `pyarrow.Schema`
823 Input pyarrow schema.
825 Returns
826 -------
827 index : `pandas.Index` or `pandas.MultiIndex`
828 Converted pandas index.
829 """
830 import pandas as pd
832 if b"pandas" in schema.metadata:
833 md = json.loads(schema.metadata[b"pandas"])
834 indexes = md["column_indexes"]
835 len_indexes = len(indexes)
836 else:
837 len_indexes = 0
839 if len_indexes <= 1:
840 return pd.Index(name for name in schema.names if not name.startswith("__"))
841 else:
842 raw_columns = _split_multi_index_column_names(len(indexes), schema.names)
843 return pd.MultiIndex.from_tuples(raw_columns, names=[f["name"] for f in indexes])
846def arrow_schema_to_column_list(schema: pa.Schema) -> list[str]:
847 """Convert an arrow schema to a list of string column names.
849 Parameters
850 ----------
851 schema : `pyarrow.Schema`
852 Input pyarrow schema.
854 Returns
855 -------
856 column_list : `list` [`str`]
857 Converted list of column names.
858 """
859 return list(schema.names)
862class DataFrameSchema:
863 """Wrapper class for a schema for a pandas DataFrame.
865 Parameters
866 ----------
867 dataframe : `pandas.DataFrame`
868 Dataframe to turn into a schema.
869 """
871 def __init__(self, dataframe: pd.DataFrame) -> None:
872 import pandas as pd
874 self._schema = dataframe.loc[[False] * len(dataframe)]
876 if isinstance(self._schema.index, pd.RangeIndex) and self._schema.index.name is not None: 876 ↛ 879line 876 didn't jump to line 879 because the condition on line 876 was never true
877 # Turn the RangeIndex into a regular index or it won't
878 # give us all the columns via arrow.
879 self._schema.index = pd.Index(self._schema.index.to_numpy(), name=self._schema.index.name)
881 @classmethod
882 def from_arrow(cls, schema: pa.Schema) -> DataFrameSchema:
883 """Convert an arrow schema into a `DataFrameSchema`.
885 Parameters
886 ----------
887 schema : `pyarrow.Schema`
888 The pyarrow schema to convert.
890 Returns
891 -------
892 dataframe_schema : `DataFrameSchema`
893 Converted dataframe schema.
894 """
895 empty_table = pa.Table.from_pylist([] * len(schema.names), schema=schema)
897 return cls(empty_table.to_pandas())
899 def to_arrow_schema(self) -> pa.Schema:
900 """Convert to an arrow schema.
902 Returns
903 -------
904 arrow_schema : `pyarrow.Schema`
905 Converted pyarrow schema.
906 """
907 arrow_table = pa.Table.from_pandas(self._schema)
909 return arrow_table.schema
911 def to_arrow_numpy_schema(self) -> ArrowNumpySchema:
912 """Convert to an `ArrowNumpySchema`.
914 Returns
915 -------
916 arrow_numpy_schema : `ArrowNumpySchema`
917 Converted arrow numpy schema.
918 """
919 return ArrowNumpySchema.from_arrow(self.to_arrow_schema())
921 def to_arrow_astropy_schema(self) -> ArrowAstropySchema:
922 """Convert to an ArrowAstropySchema.
924 Returns
925 -------
926 arrow_astropy_schema : `ArrowAstropySchema`
927 Converted arrow astropy schema.
928 """
929 return ArrowAstropySchema.from_arrow(self.to_arrow_schema())
931 @property
932 def schema(self) -> np.dtype:
933 return self._schema
935 def __repr__(self) -> str:
936 return repr(self._schema)
938 def __eq__(self, other: object) -> bool:
939 if not isinstance(other, DataFrameSchema):
940 return NotImplemented
942 return self._schema.equals(other._schema)
945class ArrowAstropySchema:
946 """Wrapper class for a schema for an astropy table.
948 Parameters
949 ----------
950 astropy_table : `astropy.table.Table`
951 Input astropy table.
952 """
954 def __init__(self, astropy_table: atable.Table) -> None:
955 self._schema = astropy_table[:0]
957 @classmethod
958 def from_arrow(cls, schema: pa.Schema) -> ArrowAstropySchema:
959 """Convert an arrow schema into a ArrowAstropySchema.
961 Parameters
962 ----------
963 schema : `pyarrow.Schema`
964 Input pyarrow schema.
966 Returns
967 -------
968 astropy_schema : `ArrowAstropySchema`
969 Converted arrow astropy schema.
970 """
971 import numpy as np
972 from astropy.table import Table
974 dtype = _schema_to_dtype_list(schema)
976 data = np.zeros(0, dtype=dtype)
977 astropy_table = Table(data=data)
979 _apply_astropy_metadata(astropy_table, schema)
981 return cls(astropy_table)
983 def to_arrow_schema(self) -> pa.Schema:
984 """Convert to an arrow schema.
986 Returns
987 -------
988 arrow_schema : `pyarrow.Schema`
989 Converted pyarrow schema.
990 """
991 return astropy_to_arrow(self._schema).schema
993 def to_dataframe_schema(self) -> DataFrameSchema:
994 """Convert to a DataFrameSchema.
996 Returns
997 -------
998 dataframe_schema : `DataFrameSchema`
999 Converted dataframe schema.
1000 """
1001 return DataFrameSchema.from_arrow(astropy_to_arrow(self._schema).schema)
1003 def to_arrow_numpy_schema(self) -> ArrowNumpySchema:
1004 """Convert to an `ArrowNumpySchema`.
1006 Returns
1007 -------
1008 arrow_numpy_schema : `ArrowNumpySchema`
1009 Converted arrow numpy schema.
1010 """
1011 return ArrowNumpySchema.from_arrow(astropy_to_arrow(self._schema).schema)
1013 @property
1014 def schema(self) -> atable.Table:
1015 return self._schema
1017 def __repr__(self) -> str:
1018 return repr(self._schema)
1020 def __eq__(self, other: object) -> bool:
1021 if not isinstance(other, ArrowAstropySchema):
1022 return NotImplemented
1024 # If this comparison passes then the two tables have the
1025 # same column names.
1026 if self._schema.dtype != other._schema.dtype:
1027 return False
1029 for name in self._schema.columns:
1030 if not self._schema[name].unit == other._schema[name].unit:
1031 return False
1032 if not self._schema[name].description == other._schema[name].description:
1033 return False
1034 if not self._schema[name].format == other._schema[name].format:
1035 return False
1037 return True
1040class ArrowNumpySchema:
1041 """Wrapper class for a schema for a numpy ndarray.
1043 Parameters
1044 ----------
1045 numpy_dtype : `numpy.dtype`
1046 Numpy dtype to convert.
1047 """
1049 def __init__(self, numpy_dtype: np.dtype) -> None:
1050 self._dtype = numpy_dtype
1052 @classmethod
1053 def from_arrow(cls, schema: pa.Schema) -> ArrowNumpySchema:
1054 """Convert an arrow schema into an `ArrowNumpySchema`.
1056 Parameters
1057 ----------
1058 schema : `pyarrow.Schema`
1059 Pyarrow schema to convert.
1061 Returns
1062 -------
1063 numpy_schema : `ArrowNumpySchema`
1064 Converted arrow numpy schema.
1065 """
1066 import numpy as np
1068 dtype = _schema_to_dtype_list(schema)
1070 return cls(np.dtype(dtype))
1072 def to_arrow_astropy_schema(self) -> ArrowAstropySchema:
1073 """Convert to an `ArrowAstropySchema`.
1075 Returns
1076 -------
1077 astropy_schema : `ArrowAstropySchema`
1078 Converted arrow astropy schema.
1079 """
1080 import numpy as np
1082 return ArrowAstropySchema.from_arrow(numpy_to_arrow(np.zeros(0, dtype=self._dtype)).schema)
1084 def to_dataframe_schema(self) -> DataFrameSchema:
1085 """Convert to a `DataFrameSchema`.
1087 Returns
1088 -------
1089 dataframe_schema : `DataFrameSchema`
1090 Converted dataframe schema.
1091 """
1092 import numpy as np
1094 return DataFrameSchema.from_arrow(numpy_to_arrow(np.zeros(0, dtype=self._dtype)).schema)
1096 def to_arrow_schema(self) -> pa.Schema:
1097 """Convert to a `pyarrow.Schema`.
1099 Returns
1100 -------
1101 arrow_schema : `pyarrow.Schema`
1102 Converted pyarrow schema.
1103 """
1104 import numpy as np
1106 return numpy_to_arrow(np.zeros(0, dtype=self._dtype)).schema
1108 @property
1109 def schema(self) -> np.dtype:
1110 return self._dtype
1112 def __repr__(self) -> str:
1113 return repr(self._dtype)
1115 def __eq__(self, other: object) -> bool:
1116 if not isinstance(other, ArrowNumpySchema):
1117 return NotImplemented
1119 if not self._dtype == other._dtype:
1120 return False
1122 return True
1125def _split_multi_index_column_names(n: int, names: Iterable[str]) -> list[Sequence[str]]:
1126 """Split a string that represents a multi-index column.
1128 PyArrow maps Pandas' multi-index column names (which are tuples in Python)
1129 to flat strings on disk. This routine exists to reconstruct the original
1130 tuple.
1132 Parameters
1133 ----------
1134 n : `int`
1135 Number of levels in the `pandas.MultiIndex` that is being
1136 reconstructed.
1137 names : `~collections.abc.Iterable` [`str`]
1138 Strings to be split.
1140 Returns
1141 -------
1142 column_names : `list` [`tuple` [`str`]]
1143 A list of multi-index column name tuples.
1144 """
1145 column_names: list[Sequence[str]] = []
1147 pattern = re.compile(r"\({}\)".format(", ".join(["'(.*)'"] * n)))
1148 for name in names:
1149 m = re.search(pattern, name)
1150 if m is not None:
1151 column_names.append(m.groups())
1153 return column_names
1156def _standardize_multi_index_columns(
1157 pd_index: pd.MultiIndex,
1158 columns: Any,
1159 stringify: bool = True,
1160) -> list[str | Sequence[Any]]:
1161 """Transform a dictionary/iterable index from a multi-index column list
1162 into a string directly understandable by PyArrow.
1164 Parameters
1165 ----------
1166 pd_index : `pandas.MultiIndex`
1167 Pandas multi-index.
1168 columns : `list` [`tuple`] or `dict` [`str`, `str` or `list` [`str`]]
1169 Columns to standardize.
1170 stringify : `bool`, optional
1171 Should the column names be stringified?
1173 Returns
1174 -------
1175 names : `list` [`str`]
1176 Stringified representation of a multi-index column name.
1177 """
1178 index_level_names = tuple(pd_index.names)
1180 names: list[str | Sequence[Any]] = []
1182 if isinstance(columns, list):
1183 for requested in columns:
1184 if not isinstance(requested, tuple):
1185 raise ValueError(
1186 "Columns parameter for multi-index data frame must be a dictionary or list of tuples. "
1187 f"Instead got a {get_full_type_name(requested)}."
1188 )
1189 if stringify:
1190 names.append(str(requested))
1191 else:
1192 names.append(requested)
1193 else:
1194 if not isinstance(columns, collections.abc.Mapping): 1194 ↛ 1195line 1194 didn't jump to line 1195 because the condition on line 1194 was never true
1195 raise ValueError(
1196 "Columns parameter for multi-index data frame must be a dictionary or list of tuples. "
1197 f"Instead got a {get_full_type_name(columns)}."
1198 )
1199 if not set(index_level_names).issuperset(columns.keys()): 1199 ↛ 1200line 1199 didn't jump to line 1200 because the condition on line 1199 was never true
1200 raise ValueError(
1201 f"Cannot use dict with keys {set(columns.keys())} to select columns from {index_level_names}."
1202 )
1203 factors = [
1204 ensure_iterable(columns.get(level, pd_index.levels[i]))
1205 for i, level in enumerate(index_level_names)
1206 ]
1207 for requested in itertools.product(*factors):
1208 for i, value in enumerate(requested):
1209 if value not in pd_index.levels[i]: 1209 ↛ 1210line 1209 didn't jump to line 1210 because the condition on line 1209 was never true
1210 raise ValueError(f"Unrecognized value {value!r} for index {index_level_names[i]!r}.")
1211 if stringify:
1212 names.append(str(requested))
1213 else:
1214 names.append(requested)
1216 return names
1219def _apply_astropy_metadata(astropy_table: atable.Table, arrow_schema: pa.Schema) -> None:
1220 """Apply any astropy metadata from the schema metadata.
1222 Parameters
1223 ----------
1224 astropy_table : `astropy.table.Table`
1225 Table to apply metadata.
1226 arrow_schema : `pyarrow.Schema`
1227 Arrow schema with metadata.
1228 """
1229 from astropy.table import meta
1231 metadata = arrow_schema.metadata if arrow_schema.metadata is not None else {}
1233 # Check if we have a special astropy metadata header yaml.
1234 meta_yaml = metadata.get(b"table_meta_yaml", None)
1235 if meta_yaml:
1236 meta_yaml = meta_yaml.decode("UTF8").split("\n")
1237 meta_hdr = meta.get_header_from_yaml(meta_yaml)
1239 # Set description, format, unit, meta from the column
1240 # metadata that was serialized with the table.
1241 header_cols = {x["name"]: x for x in meta_hdr["datatype"]}
1242 for col in astropy_table.columns.values():
1243 for attr in ("description", "format", "unit", "meta"):
1244 if attr in header_cols[col.name]:
1245 setattr(col, attr, header_cols[col.name][attr])
1247 if "meta" in meta_hdr:
1248 astropy_table.meta.update(meta_hdr["meta"])
1249 else:
1250 # If we don't have astropy header data, we may have arrow field
1251 # metadata.
1252 for name in arrow_schema.names:
1253 field_metadata = arrow_schema.field(name).metadata
1254 if field_metadata is None:
1255 continue
1256 if (
1257 b"description" in field_metadata
1258 and (description := field_metadata[b"description"].decode("UTF-8")) != ""
1259 ):
1260 astropy_table[name].description = description
1261 if b"unit" in field_metadata and (unit := field_metadata[b"unit"].decode("UTF-8")) != "":
1262 astropy_table[name].unit = unit
1264 # Ensure that the special ASTROPY_PANDAS_INDEX_KEY is propagated to
1265 # the table metadata.
1266 if index_key := metadata.get(ASTROPY_PANDAS_INDEX_KEY.encode(), None):
1267 astropy_table.meta[ASTROPY_PANDAS_INDEX_KEY] = index_key.decode("UTF-8")
1270def _arrow_string_to_numpy_dtype(
1271 schema: pa.Schema, name: str, numpy_column: np.ndarray | None = None, default_length: int = 10
1272) -> str:
1273 """Get the numpy dtype string associated with an arrow column.
1275 Parameters
1276 ----------
1277 schema : `pyarrow.Schema`
1278 Arrow table schema.
1279 name : `str`
1280 Column name.
1281 numpy_column : `numpy.ndarray`, optional
1282 Column to determine numpy string dtype.
1283 default_length : `int`, optional
1284 Default string length when not in metadata or can be inferred
1285 from column.
1287 Returns
1288 -------
1289 dtype_str : `str`
1290 Numpy dtype string.
1291 """
1292 # Special-case for string and binary columns
1293 md_name = f"lsst::arrow::len::{name}"
1294 strlen = default_length
1295 metadata = schema.metadata if schema.metadata is not None else {}
1296 if (encoded := md_name.encode("UTF-8")) in metadata:
1297 # String/bytes length from header.
1298 strlen = int(schema.metadata[encoded])
1299 elif numpy_column is not None and len(numpy_column) > 0:
1300 lengths = [len(row) for row in numpy_column if row]
1301 strlen = max(lengths) if lengths else 0
1303 dtype = f"U{strlen}" if _is_string(schema.field(name).type) else f"|S{strlen}"
1305 return dtype
1308def _append_numpy_string_metadata(metadata: dict[bytes, str], name: str, dtype: np.dtype) -> None:
1309 """Append numpy string length keys to arrow metadata.
1311 All column types are handled, but the metadata is only modified for
1312 string and byte columns.
1314 Parameters
1315 ----------
1316 metadata : `dict` [`bytes`, `str`]
1317 Metadata dictionary; modified in place.
1318 name : `str`
1319 Column name.
1320 dtype : `np.dtype`
1321 Numpy dtype.
1322 """
1323 import numpy as np
1325 if dtype.type is np.str_:
1326 metadata[f"lsst::arrow::len::{name}".encode()] = str(dtype.itemsize // 4)
1327 metadata[f"table::len::{name}".encode()] = str(dtype.itemsize // 4)
1328 elif dtype.type is np.bytes_:
1329 metadata[f"lsst::arrow::len::{name}".encode()] = str(dtype.itemsize)
1330 metadata[f"table::len::{name}".encode()] = str(dtype.itemsize)
1333def _append_numpy_multidim_metadata(metadata: dict[bytes, str], name: str, dtype: np.dtype) -> None:
1334 """Append numpy multi-dimensional shapes to arrow metadata.
1336 All column types are handled, but the metadata is only modified for
1337 multi-dimensional columns.
1339 Parameters
1340 ----------
1341 metadata : `dict` [`bytes`, `str`]
1342 Metadata dictionary; modified in place.
1343 name : `str`
1344 Column name.
1345 dtype : `np.dtype`
1346 Numpy dtype.
1347 """
1348 if len(dtype.shape) > 1:
1349 metadata[f"lsst::arrow::shape::{name}".encode()] = str(dtype.shape)
1352def _multidim_shape_from_metadata(metadata: dict[bytes, bytes], list_size: int, name: str) -> tuple[int, ...]:
1353 """Retrieve the shape from the metadata, if available.
1355 Parameters
1356 ----------
1357 metadata : `dict` [`bytes`, `bytes`]
1358 Metadata dictionary.
1359 list_size : `int`
1360 Size of the list datatype.
1361 name : `str`
1362 Column name.
1364 Returns
1365 -------
1366 shape : `tuple` [`int`]
1367 Shape associated with the column.
1369 Raises
1370 ------
1371 RuntimeError
1372 Raised if metadata is found but has incorrect format.
1373 """
1374 md_name = f"lsst::arrow::shape::{name}"
1375 if (encoded := md_name.encode("UTF-8")) in metadata:
1376 groups = re.search(r"\((.*)\)", metadata[encoded].decode("UTF-8"))
1377 if groups is None: 1377 ↛ 1378line 1377 didn't jump to line 1378 because the condition on line 1377 was never true
1378 raise RuntimeError("Illegal value found in metadata.")
1379 shape = tuple(int(x) for x in groups[1].split(",") if x != "")
1380 else:
1381 shape = (list_size,)
1383 return shape
1386def _schema_to_dtype_list(schema: pa.Schema) -> list[tuple[str, tuple[Any] | str]]:
1387 """Convert a pyarrow schema to a numpy dtype.
1389 Parameters
1390 ----------
1391 schema : `pyarrow.Schema`
1392 Input pyarrow schema.
1394 Returns
1395 -------
1396 dtype_list: `list` [`tuple`]
1397 A list with name, type pairs.
1398 """
1399 metadata = schema.metadata if schema.metadata is not None else {}
1401 dtype: list[Any] = []
1402 for name in schema.names:
1403 t = schema.field(name).type
1404 if isinstance(t, pa.FixedSizeListType):
1405 shape = _multidim_shape_from_metadata(metadata, t.list_size, name)
1406 dtype.append((name, (t.value_type.to_pandas_dtype(), shape)))
1407 elif not (_is_string(t) or _is_binary(t)):
1408 dtype.append((name, t.to_pandas_dtype()))
1409 else:
1410 dtype.append((name, _arrow_string_to_numpy_dtype(schema, name)))
1412 return dtype
1415def _numpy_dtype_to_arrow_types(dtype: np.dtype) -> list[Any]:
1416 """Convert a numpy dtype to a list of arrow types.
1418 Parameters
1419 ----------
1420 dtype : `numpy.dtype`
1421 Numpy dtype to convert.
1423 Returns
1424 -------
1425 type_list : `list` [`object`]
1426 Converted list of arrow types.
1427 """
1428 from math import prod
1430 import numpy as np
1432 type_list: list[Any] = []
1433 if dtype.names is None: 1433 ↛ 1434line 1433 didn't jump to line 1434 because the condition on line 1433 was never true
1434 return type_list
1436 for name in dtype.names:
1437 dt = dtype[name]
1438 arrow_type: Any
1439 if len(dt.shape) > 0:
1440 arrow_type = pa.list_(
1441 pa.from_numpy_dtype(cast(tuple[np.dtype, tuple[int, ...]], dt.subdtype)[0].type),
1442 prod(dt.shape),
1443 )
1444 elif dt.type == np.datetime64:
1445 time_unit = "ns" if "ns" in dt.str else "us"
1446 # The pa.timestamp() is the correct datatype to round-trip
1447 # a numpy datetime64[ns] or datetime[us] array.
1448 arrow_type = pa.timestamp(time_unit)
1449 else:
1450 try:
1451 arrow_type = pa.from_numpy_dtype(dt.type)
1452 except pa.ArrowNotImplementedError as e:
1453 msg = f"Could not serialize column {name} (type {str(dt)}) to Parquet."
1454 if dt == np.dtype("O"): 1454 ↛ 1456line 1454 didn't jump to line 1456 because the condition on line 1454 was always true
1455 msg += " This is usually because the column is mixed type or has uneven length rows."
1456 e.add_note(msg)
1457 raise
1458 type_list.append((name, arrow_type))
1460 return type_list
1463def _numpy_dict_to_dtype(numpy_dict: dict[str, np.ndarray]) -> tuple[np.dtype, int]:
1464 """Extract equivalent table dtype from dict of numpy arrays.
1466 Parameters
1467 ----------
1468 numpy_dict : `dict` [`str`, `numpy.ndarray`]
1469 Dict with keys as the column names, values as the arrays.
1471 Returns
1472 -------
1473 dtype : `numpy.dtype`
1474 dtype of equivalent table.
1475 rowcount : `int`
1476 Number of rows in the table.
1478 Raises
1479 ------
1480 ValueError if columns in numpy_dict have unequal numbers of rows.
1481 """
1482 import numpy as np
1484 dtype_list: list[tuple] = []
1485 rowcount = 0
1486 for name, col in numpy_dict.items():
1487 if rowcount == 0:
1488 rowcount = len(col)
1489 if len(col) != rowcount:
1490 raise ValueError(f"Column {name} has a different number of rows.")
1491 if len(col.shape) == 1:
1492 dtype_list.append((name, col.dtype))
1493 else:
1494 dtype_list.append((name, (col.dtype, col.shape[1:])))
1495 dtype = np.dtype(dtype_list)
1497 return (dtype, rowcount)
1500def _numpy_style_arrays_to_arrow_arrays(
1501 dtype: np.dtype,
1502 rowcount: int,
1503 np_style_arrays: dict[str, np.ndarray] | np.ndarray | atable.Table,
1504 schema: pa.Schema,
1505) -> list[pa.Array]:
1506 """Convert numpy-style arrays to arrow arrays.
1508 Parameters
1509 ----------
1510 dtype : `numpy.dtype`
1511 Numpy dtype of input table/arrays.
1512 rowcount : `int`
1513 Number of rows in input table/arrays.
1514 np_style_arrays : `dict` [`str`, `np.ndarray`] or `np.ndarray`
1515 or `astropy.table.Table`
1516 Arrays to convert to arrow.
1517 schema : `pyarrow.Schema`
1518 Schema of arrow table.
1520 Returns
1521 -------
1522 arrow_arrays : `list` [`pyarrow.Array`]
1523 List of converted pyarrow arrays.
1524 """
1525 import numpy as np
1527 arrow_arrays: list[pa.Array] = []
1528 if dtype.names is None: 1528 ↛ 1529line 1528 didn't jump to line 1529 because the condition on line 1528 was never true
1529 return arrow_arrays
1531 for name in dtype.names:
1532 dt = dtype[name]
1533 val: Any
1534 if len(dt.shape) > 0:
1535 if rowcount > 0:
1536 val = np.split(np_style_arrays[name].ravel(), rowcount)
1537 else:
1538 val = []
1539 else:
1540 val = np_style_arrays[name]
1542 try:
1543 arrow_arrays.append(pa.array(val, type=schema.field(name).type))
1544 except pa.ArrowNotImplementedError as err:
1545 # Check if val is big-endian.
1546 if (np.little_endian and val.dtype.byteorder == ">") or ( 1546 ↛ 1555line 1546 didn't jump to line 1555 because the condition on line 1546 was always true
1547 not np.little_endian and val.dtype.byteorder == "="
1548 ):
1549 # We need to convert the array to little-endian.
1550 val2 = val.byteswap()
1551 val2.dtype = val2.dtype.newbyteorder("<")
1552 arrow_arrays.append(pa.array(val2, type=schema.field(name).type))
1553 else:
1554 # This failed for some other reason so raise the exception.
1555 raise err
1557 return arrow_arrays
1560def compute_row_group_size(schema: pa.Schema, target_size: int = TARGET_ROW_GROUP_BYTES) -> int:
1561 """Compute approximate row group size for a given arrow schema.
1563 Given a schema, this routine will compute the number of rows in a row group
1564 that targets the persisted size on disk (or smaller). The exact size on
1565 disk depends on the compression settings and ratios; typical binary data
1566 tables will have around 15-20% compression with the pyarrow default
1567 ``snappy`` compression algorithm.
1569 Parameters
1570 ----------
1571 schema : `pyarrow.Schema`
1572 Arrow table schema.
1573 target_size : `int`, optional
1574 The target size (in bytes).
1576 Returns
1577 -------
1578 row_group_size : `int`
1579 Number of rows per row group to hit the target size.
1580 """
1581 bit_width = 0
1583 metadata = schema.metadata if schema.metadata is not None else {}
1585 for name in schema.names:
1586 if name.startswith("__"):
1587 continue
1589 t = schema.field(name).type
1591 if _is_string(t) or _is_binary(t):
1592 md_name = f"lsst::arrow::len::{name}"
1594 if (encoded := md_name.encode("UTF-8")) in metadata:
1595 # String/bytes length from header.
1596 strlen = int(schema.metadata[encoded])
1597 else:
1598 # We don't know the string width, so guess something.
1599 strlen = 10
1601 # Assuming UTF-8 encoding, and very few wide characters.
1602 t_width = 8 * strlen
1603 elif isinstance(t, pa.FixedSizeListType):
1604 if t.value_type == pa.null(): 1604 ↛ 1605line 1604 didn't jump to line 1605 because the condition on line 1604 was never true
1605 t_width = 0
1606 else:
1607 t_width = t.list_size * t.value_type.bit_width
1608 elif t == pa.null():
1609 t_width = 0
1610 elif isinstance(t, pa.ListType):
1611 if t.value_type == pa.null():
1612 t_width = 0
1613 else:
1614 # This is a variable length list, just choose
1615 # something arbitrary.
1616 t_width = 10 * t.value_type.bit_width
1617 else:
1618 t_width = t.bit_width
1620 bit_width += t_width
1622 # Insist it is at least 1 byte wide to avoid any divide-by-zero errors.
1623 if bit_width < 8:
1624 bit_width = 8
1626 byte_width = bit_width // 8
1628 return target_size // byte_width
1631def _is_string(t: pa.DataType) -> bool:
1632 return pa.types.is_string(t) or pa.types.is_large_string(t) or pa.types.is_string_view(t)
1635def _is_binary(t: pa.DataType) -> bool:
1636 return pa.types.is_binary(t) or pa.types.is_large_binary(t) or pa.types.is_binary_view(t)