Coverage for python/lsst/daf/butler/formatters/parquet.py: 96%
562 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-06-29 15:11 -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)
51import collections.abc
52import contextlib
53import itertools
54import json
55import logging
56import re
57from collections.abc import Generator, Iterable, Sequence
58from fnmatch import fnmatchcase
59from typing import IO, TYPE_CHECKING, Any, cast
61import pyarrow as pa
62import pyarrow.parquet as pq
64from lsst.daf.butler import DatasetProvenance, FormatterV2
65from lsst.daf.butler.delegates.arrowtable import _add_arrow_provenance, _checkArrowCompatibleType
66from lsst.resources import ResourcePath
67from lsst.utils.introspection import get_full_type_name
68from lsst.utils.iteration import ensure_iterable
70log = logging.getLogger(__name__)
72if TYPE_CHECKING:
73 import astropy.table as atable
74 import numpy as np
75 import pandas as pd
77 try:
78 import fsspec
79 from fsspec.spec import AbstractFileSystem
80 except ImportError:
81 fsspec = None
82 AbstractFileSystem = type
84TARGET_ROW_GROUP_BYTES = 1_000_000_000
85ASTROPY_PANDAS_INDEX_KEY = "lsst::arrow::astropy_pandas_index"
88@contextlib.contextmanager
89def generic_open(path: str, fs: AbstractFileSystem | None) -> Generator[IO]:
90 if fs is None:
91 with open(path, "rb") as fh:
92 yield fh
93 else:
94 with fs.open(path) as fh:
95 yield fh
98class ParquetFormatter(FormatterV2):
99 """Interface for reading and writing Arrow Table objects to and from
100 Parquet files.
101 """
103 default_extension = ".parq"
104 can_read_from_uri = True
105 can_read_from_local_file = True
107 def can_accept(self, in_memory_dataset: Any) -> bool:
108 # Docstring inherited.
109 return _checkArrowCompatibleType(in_memory_dataset) is not None
111 def read_from_uri(self, uri: ResourcePath, component: str | None = None, expected_size: int = -1) -> Any:
112 # Docstring inherited from Formatter.read.
113 try:
114 fs, path = uri.to_fsspec()
115 except ImportError:
116 log.debug("fsspec not available; falling back to local file access.")
117 # This signals to the formatter to use the read_from_local_file
118 # code path.
119 return NotImplemented
121 return self._read_parquet(path=path, fs=fs, component=component, expected_size=expected_size)
123 def read_from_local_file(self, path: str, component: str | None = None, expected_size: int = -1) -> Any:
124 # Docstring inherited from Formatter.read.
125 return self._read_parquet(path=path, component=component, expected_size=expected_size)
127 def _read_parquet(
128 self,
129 path: str,
130 fs: AbstractFileSystem | None = None,
131 component: str | None = None,
132 expected_size: int = -1,
133 ) -> Any:
134 with generic_open(path, fs) as handle:
135 schema = pq.read_schema(handle)
137 schema_names = ["ArrowSchema", "DataFrameSchema", "ArrowAstropySchema", "ArrowNumpySchema"]
139 if component in ("columns", "schema") or self.file_descriptor.readStorageClass.name in schema_names:
140 # The schema will be translated to column format
141 # depending on the input type.
142 return schema
143 elif component == "rowcount":
144 # Get the rowcount from the metadata if possible, otherwise count.
145 if b"lsst::arrow::rowcount" in schema.metadata:
146 return int(schema.metadata[b"lsst::arrow::rowcount"])
148 with generic_open(path, fs) as handle:
149 temp_table = pq.read_table(
150 handle,
151 columns=[schema.names[0]],
152 use_threads=False,
153 use_pandas_metadata=False,
154 )
156 return len(temp_table[schema.names[0]])
158 par_columns = None
159 strip_astropy_meta_yaml = True
160 if self.file_descriptor.parameters:
161 par_columns = self.file_descriptor.parameters.pop("columns", None)
162 if par_columns:
163 has_pandas_multi_index = False
164 if schema.metadata and b"pandas" in schema.metadata:
165 md = json.loads(schema.metadata[b"pandas"])
166 if len(md["column_indexes"]) > 1:
167 has_pandas_multi_index = True
169 if not has_pandas_multi_index:
170 # Ensure uniqueness, keeping order.
171 par_columns_in = list(dict.fromkeys(ensure_iterable(par_columns)))
172 file_columns = [name for name in schema.names if not name.startswith("__")]
174 # Do case-sensitive glob-style matching, again ensuring
175 # uniqueness and ordering.
176 par_columns = {}
177 for par_column in par_columns_in:
178 found = False
179 for file_column in file_columns:
180 if fnmatchcase(file_column, par_column):
181 found = True
182 par_columns[file_column] = True
183 if not found:
184 raise ValueError(
185 f"Column {par_column} specified in parameters not available in parquet file."
186 )
187 par_columns = list(par_columns.keys())
188 else:
189 par_columns = _standardize_multi_index_columns(
190 arrow_schema_to_pandas_index(schema),
191 par_columns,
192 )
194 strip_astropy_meta_yaml = self.file_descriptor.parameters.pop(
195 "strip_astropy_meta_yaml",
196 True,
197 )
199 if len(self.file_descriptor.parameters): 199 ↛ 200line 199 didn't jump to line 200 because the condition on line 199 was never true
200 raise ValueError(
201 f"Unsupported parameters {self.file_descriptor.parameters} in ArrowTable read."
202 )
204 metadata = schema.metadata if schema.metadata is not None else {}
205 with generic_open(path, fs) as handle:
206 arrow_table = pq.read_table(
207 handle,
208 columns=par_columns,
209 use_threads=False,
210 use_pandas_metadata=(b"pandas" in metadata),
211 )
213 if strip_astropy_meta_yaml:
214 metadata = arrow_table.schema.metadata
215 # Only strip if (a) we have metadata; (b) it contains
216 # ``table_meta_yaml``; (c) it contains ``lsst::arrow::rowcount``
217 # to avoid stripping data from pure astropy tables (not written
218 # by the butler).
219 if metadata and metadata.pop(b"table_meta_yaml", None) and b"lsst::arrow::rowcount" in metadata:
220 arrow_table = arrow_table.replace_schema_metadata(metadata)
222 return arrow_table
224 def add_provenance(self, in_memory_dataset: Any, provenance: DatasetProvenance | None = None) -> Any:
225 return _add_arrow_provenance(in_memory_dataset, self.dataset_ref, provenance)
227 def write_local_file(self, in_memory_dataset: Any, uri: ResourcePath) -> None:
228 """Serialize the in memory dataset to a local parquet file.
230 Parameters
231 ----------
232 in_memory_dataset : `typing.Any`
233 The Python object to serialize.
234 uri : `lsst.resources.ResourcePath`
235 The location to write the local file.
236 """
237 if isinstance(in_memory_dataset, pa.Schema):
238 pq.write_metadata(in_memory_dataset, uri.ospath)
239 return
241 type_string = _checkArrowCompatibleType(in_memory_dataset)
243 if type_string is None:
244 raise ValueError(
245 f"Unsupported type {get_full_type_name(in_memory_dataset)} of "
246 "inMemoryDataset for ParquetFormatter."
247 )
249 if type_string == "arrow":
250 arrow_table = in_memory_dataset
251 elif type_string == "astropy":
252 arrow_table = astropy_to_arrow(in_memory_dataset)
253 elif type_string == "numpy":
254 arrow_table = numpy_to_arrow(in_memory_dataset)
255 elif type_string == "numpydict":
256 arrow_table = numpy_dict_to_arrow(in_memory_dataset)
257 else:
258 arrow_table = pandas_to_arrow(in_memory_dataset)
260 row_group_size = compute_row_group_size(arrow_table.schema)
262 pq.write_table(arrow_table, uri.ospath, row_group_size=row_group_size)
265def arrow_to_pandas(arrow_table: pa.Table) -> pd.DataFrame:
266 """Convert a pyarrow table to a pandas DataFrame.
268 Parameters
269 ----------
270 arrow_table : `pyarrow.Table`
271 Input arrow table to convert. If the table has ``pandas`` metadata
272 in the schema it will be used in the construction of the
273 ``DataFrame``.
275 Returns
276 -------
277 dataframe : `pandas.DataFrame`
278 Converted pandas dataframe.
279 """
280 dataframe = arrow_table.to_pandas(use_threads=False, integer_object_nulls=True)
282 metadata = arrow_table.schema.metadata if arrow_table.schema.metadata is not None else {}
283 if (key := ASTROPY_PANDAS_INDEX_KEY.encode()) in metadata:
284 pandas_index = metadata[key].decode("UTF8")
285 if pandas_index in arrow_table.schema.names:
286 dataframe.set_index(pandas_index, inplace=True)
287 else:
288 log.warning(
289 "Index column ``%s`` not available for arrow table conversion to DataFrame",
290 pandas_index,
291 )
293 return dataframe
296def arrow_to_astropy(arrow_table: pa.Table) -> atable.Table:
297 """Convert a pyarrow table to an `astropy.table.Table`.
299 Parameters
300 ----------
301 arrow_table : `pyarrow.Table`
302 Input arrow table to convert. If the table has astropy unit
303 metadata in the schema it will be used in the construction
304 of the ``astropy.table.Table``.
306 Returns
307 -------
308 table : `astropy.table.Table`
309 Converted astropy table.
310 """
311 from astropy.table import Table
313 astropy_table = Table(arrow_to_numpy_dict(arrow_table))
315 _apply_astropy_metadata(astropy_table, arrow_table.schema)
317 if (key := ASTROPY_PANDAS_INDEX_KEY) in astropy_table.meta:
318 if astropy_table.meta[key] not in astropy_table.columns:
319 astropy_table.meta.pop(key)
321 return astropy_table
324def arrow_to_numpy(arrow_table: pa.Table) -> np.ndarray | np.ma.MaskedArray:
325 """Convert a pyarrow table to a structured numpy array.
327 Parameters
328 ----------
329 arrow_table : `pyarrow.Table`
330 Input arrow table.
332 Returns
333 -------
334 array : `numpy.ndarray` or `numpy.ma.MaskedArray` (N,)
335 Numpy array table with N rows and the same column names
336 as the input arrow table. Will be masked records if any values
337 in the table are null.
338 """
339 import numpy as np
341 numpy_dict = arrow_to_numpy_dict(arrow_table)
343 has_mask = False
344 dtype: list[tuple] = []
345 for name, col in numpy_dict.items():
346 if len(shape := numpy_dict[name].shape) <= 1:
347 dtype.append((name, col.dtype))
348 else:
349 dtype.append((name, (col.dtype, shape[1:])))
351 if not has_mask and isinstance(col, np.ma.MaskedArray):
352 has_mask = True
354 array: Any
356 if has_mask:
357 import numpy.ma.mrecords as mrecords
359 array = mrecords.fromarrays(list(numpy_dict.values()), dtype=dtype)
360 else:
361 array = np.rec.fromarrays(numpy_dict.values(), dtype=dtype)
362 return array
365def arrow_to_numpy_dict(arrow_table: pa.Table) -> dict[str, np.ndarray]:
366 """Convert a pyarrow table to a dict of numpy arrays.
368 Parameters
369 ----------
370 arrow_table : `pyarrow.Table`
371 Input arrow table.
373 Returns
374 -------
375 numpy_dict : `dict` [`str`, `numpy.ndarray`]
376 Dict with keys as the column names, values as the arrays.
377 """
378 import numpy as np
380 schema = arrow_table.schema
381 metadata = schema.metadata if schema.metadata is not None else {}
383 numpy_dict = {}
385 for name in schema.names:
386 t = schema.field(name).type
388 if arrow_table[name].null_count == 0:
389 # Regular non-masked column
390 col = arrow_table[name].to_numpy()
391 else:
392 # For a masked column, we need to ask arrow to fill the null
393 # values with an appropriately typed value before conversion.
394 # Then we apply the mask to get a masked array of the correct type.
395 null_value: Any
396 match t:
397 case t if t in (pa.float64(), pa.float32(), pa.float16()):
398 null_value = np.nan
399 case t if t in (pa.int64(), pa.int32(), pa.int16(), pa.int8()):
400 null_value = -1
401 case t if t in (pa.bool_(),):
402 null_value = True
403 case t if t in (pa.string(), pa.binary()):
404 null_value = ""
405 case _:
406 # This is the fallback for unsigned ints in particular.
407 null_value = 0
409 col = np.ma.MaskedArray(
410 data=arrow_table[name].fill_null(null_value).to_numpy(),
411 mask=arrow_table[name].is_null().to_numpy(),
412 fill_value=null_value,
413 )
415 if t in (pa.string(), pa.binary()):
416 col = col.astype(_arrow_string_to_numpy_dtype(schema, name, col))
417 elif isinstance(t, pa.FixedSizeListType):
418 if len(col) > 0:
419 col = np.stack(col)
420 else:
421 # this is an empty column, and needs to be coerced to type.
422 col = col.astype(t.value_type.to_pandas_dtype())
424 shape = _multidim_shape_from_metadata(metadata, t.list_size, name)
425 col = col.reshape((len(arrow_table), *shape))
427 numpy_dict[name] = col
429 return numpy_dict
432def _numpy_dict_to_numpy(numpy_dict: dict[str, np.ndarray]) -> np.ndarray:
433 """Convert a dict of numpy arrays to a structured numpy array.
435 Parameters
436 ----------
437 numpy_dict : `dict` [`str`, `numpy.ndarray`]
438 Dict with keys as the column names, values as the arrays.
440 Returns
441 -------
442 array : `numpy.ndarray` (N,)
443 Numpy array table with N rows and columns names from the dict keys.
444 """
445 return arrow_to_numpy(numpy_dict_to_arrow(numpy_dict))
448def _numpy_to_numpy_dict(np_array: np.ndarray) -> dict[str, np.ndarray]:
449 """Convert a structured numpy array to a dict of numpy arrays.
451 Parameters
452 ----------
453 np_array : `numpy.ndarray`
454 Input numpy array with multiple fields.
456 Returns
457 -------
458 numpy_dict : `dict` [`str`, `numpy.ndarray`]
459 Dict with keys as the column names, values as the arrays.
460 """
461 return arrow_to_numpy_dict(numpy_to_arrow(np_array))
464def numpy_to_arrow(np_array: np.ndarray) -> pa.Table:
465 """Convert a numpy array table to an arrow table.
467 Parameters
468 ----------
469 np_array : `numpy.ndarray`
470 Input numpy array with multiple fields.
472 Returns
473 -------
474 arrow_table : `pyarrow.Table`
475 Converted arrow table.
476 """
477 type_list = _numpy_dtype_to_arrow_types(np_array.dtype)
479 md = {}
480 md[b"lsst::arrow::rowcount"] = str(len(np_array))
482 names = np_array.dtype.names
483 if names is None: 483 ↛ 484line 483 didn't jump to line 484 because the condition on line 483 was never true
484 names = ()
486 for name in names:
487 _append_numpy_string_metadata(md, name, np_array.dtype[name])
488 _append_numpy_multidim_metadata(md, name, np_array.dtype[name])
490 schema = pa.schema(type_list, metadata=md)
492 arrays = _numpy_style_arrays_to_arrow_arrays(
493 np_array.dtype,
494 len(np_array),
495 np_array,
496 schema,
497 )
499 arrow_table = pa.Table.from_arrays(arrays, schema=schema)
501 return arrow_table
504def numpy_dict_to_arrow(numpy_dict: dict[str, np.ndarray]) -> pa.Table:
505 """Convert a dict of numpy arrays to an arrow table.
507 Parameters
508 ----------
509 numpy_dict : `dict` [`str`, `numpy.ndarray`]
510 Dict with keys as the column names, values as the arrays.
512 Returns
513 -------
514 arrow_table : `pyarrow.Table`
515 Converted arrow table.
517 Raises
518 ------
519 ValueError
520 Raised if columns in ``numpy_dict`` have unequal numbers of
521 rows.
522 """
523 dtype, rowcount = _numpy_dict_to_dtype(numpy_dict)
524 type_list = _numpy_dtype_to_arrow_types(dtype)
526 md = {}
527 md[b"lsst::arrow::rowcount"] = str(rowcount)
529 if dtype.names is not None: 529 ↛ 534line 529 didn't jump to line 534 because the condition on line 529 was always true
530 for name in dtype.names:
531 _append_numpy_string_metadata(md, name, dtype[name])
532 _append_numpy_multidim_metadata(md, name, dtype[name])
534 schema = pa.schema(type_list, metadata=md)
536 arrays = _numpy_style_arrays_to_arrow_arrays(
537 dtype,
538 rowcount,
539 numpy_dict,
540 schema,
541 )
543 arrow_table = pa.Table.from_arrays(arrays, schema=schema)
545 return arrow_table
548def astropy_to_arrow(astropy_table: atable.Table) -> pa.Table:
549 """Convert an astropy table to an arrow table.
551 Parameters
552 ----------
553 astropy_table : `astropy.table.Table`
554 Input astropy table.
556 Returns
557 -------
558 arrow_table : `pyarrow.Table`
559 Converted arrow table.
560 """
561 from astropy.table import meta
563 type_list = _numpy_dtype_to_arrow_types(astropy_table.dtype)
565 md = {}
566 md[b"lsst::arrow::rowcount"] = str(len(astropy_table))
568 if (key := ASTROPY_PANDAS_INDEX_KEY) in astropy_table.meta:
569 md[key.encode()] = astropy_table.meta[key]
571 for name in astropy_table.dtype.names:
572 _append_numpy_string_metadata(md, name, astropy_table.dtype[name])
573 _append_numpy_multidim_metadata(md, name, astropy_table.dtype[name])
575 meta_yaml = meta.get_yaml_from_table(astropy_table)
576 meta_yaml_str = "\n".join(meta_yaml)
577 md[b"table_meta_yaml"] = meta_yaml_str
579 # Convert type list to fields with metadata.
580 fields = []
581 for name, pa_type in type_list:
582 field_metadata = {}
583 if description := astropy_table[name].description:
584 field_metadata["description"] = description
585 if unit := astropy_table[name].unit:
586 field_metadata["unit"] = str(unit)
587 fields.append(
588 pa.field(
589 name,
590 pa_type,
591 metadata=field_metadata,
592 )
593 )
595 schema = pa.schema(fields, metadata=md)
597 arrays = _numpy_style_arrays_to_arrow_arrays(
598 astropy_table.dtype,
599 len(astropy_table),
600 astropy_table,
601 schema,
602 )
604 arrow_table = pa.Table.from_arrays(arrays, schema=schema)
606 return arrow_table
609def astropy_to_pandas(astropy_table: atable.Table, index: str | None = None) -> pd.DataFrame:
610 """Convert an astropy table to a pandas dataframe via arrow.
612 By going via arrow we avoid pandas masked column bugs (e.g.
613 https://github.com/pandas-dev/pandas/issues/58173)
615 Parameters
616 ----------
617 astropy_table : `astropy.table.Table`
618 Input astropy table.
619 index : `str`, optional
620 Name of column to set as index.
622 Returns
623 -------
624 dataframe : `pandas.DataFrame`
625 Output pandas dataframe.
626 """
627 index_requested = False
628 if (key := ASTROPY_PANDAS_INDEX_KEY) in astropy_table.meta:
629 _index = astropy_table.meta[key]
630 if _index not in astropy_table.columns:
631 log.warning(
632 "Index column ``%s`` not available for astropy table conversion to DataFrame",
633 _index,
634 )
635 _index = None
636 else:
637 index_requested = True
638 _index = index
640 dataframe = arrow_to_pandas(astropy_to_arrow(astropy_table))
642 # Set the index if we have a valid index name, and either the
643 # index was requested in the call to the function or the dataframe
644 # was not previously indexed with the call to arrow_to_pandas.
645 if isinstance(_index, str) and (index_requested or dataframe.index.name is None):
646 dataframe.set_index(_index, inplace=True)
647 elif _index and index_requested: 647 ↛ 648line 647 didn't jump to line 648 because the condition on line 647 was never true
648 raise RuntimeError("index must be a string or None.")
650 return dataframe
653def add_pandas_index_to_astropy(astropy_table: atable.Table, index: str) -> None:
654 """Add special metadata to an astropy table to indicate a pandas index.
656 Parameters
657 ----------
658 astropy_table : `astropy.table.Table`
659 Input astropy table.
660 index : `str`
661 Name of column for pandas to set as index, if read as DataFrame.
662 """
663 if index not in astropy_table.columns:
664 raise ValueError("Column ``%s`` not in astropy table columns to use as pandas index.", index)
665 astropy_table.meta[ASTROPY_PANDAS_INDEX_KEY] = index
668def _astropy_to_numpy_dict(astropy_table: atable.Table) -> dict[str, np.ndarray]:
669 """Convert an astropy table to an arrow table.
671 Parameters
672 ----------
673 astropy_table : `astropy.table.Table`
674 Input astropy table.
676 Returns
677 -------
678 numpy_dict : `dict` [`str`, `numpy.ndarray`]
679 Dict with keys as the column names, values as the arrays.
680 """
681 return arrow_to_numpy_dict(astropy_to_arrow(astropy_table))
684def pandas_to_arrow(dataframe: pd.DataFrame, default_length: int = 10) -> pa.Table:
685 """Convert a pandas dataframe to an arrow table.
687 Parameters
688 ----------
689 dataframe : `pandas.DataFrame`
690 Input pandas dataframe.
691 default_length : `int`, optional
692 Default string length when not in metadata or can be inferred
693 from column.
695 Returns
696 -------
697 arrow_table : `pyarrow.Table`
698 Converted arrow table.
699 """
700 try:
701 arrow_table = pa.Table.from_pandas(dataframe)
702 except pa.ArrowInvalid as e:
703 msg = "; ".join(e.args)
704 msg += "; This is usually because the column is mixed type or has uneven length rows."
705 e.add_note(msg)
706 raise
708 # Update the metadata
709 md = arrow_table.schema.metadata
711 md[b"lsst::arrow::rowcount"] = str(arrow_table.num_rows)
713 # We loop through the arrow table columns because the datatypes have
714 # been checked and converted from pandas objects.
715 for name in arrow_table.column_names:
716 if not name.startswith("__") and arrow_table[name].type == pa.string():
717 if len(arrow_table[name]) > 0: 717 ↛ 720line 717 didn't jump to line 720 because the condition on line 717 was always true
718 strlen = max(len(row.as_py()) for row in arrow_table[name] if row.is_valid)
719 else:
720 strlen = default_length
721 md[f"lsst::arrow::len::{name}".encode()] = str(strlen)
723 arrow_table = arrow_table.replace_schema_metadata(md)
725 return arrow_table
728def pandas_to_astropy(dataframe: pd.DataFrame) -> atable.Table:
729 """Convert a pandas dataframe to an astropy table, preserving indexes.
731 Parameters
732 ----------
733 dataframe : `pandas.DataFrame`
734 Input pandas dataframe.
736 Returns
737 -------
738 astropy_table : `astropy.table.Table`
739 Converted astropy table.
740 """
741 import pandas as pd
743 if isinstance(dataframe.columns, pd.MultiIndex):
744 raise ValueError("Cannot convert a multi-index dataframe to an astropy table.")
746 return arrow_to_astropy(pandas_to_arrow(dataframe))
749def _pandas_to_numpy_dict(dataframe: pd.DataFrame) -> dict[str, np.ndarray]:
750 """Convert a pandas dataframe to an dict of numpy arrays.
752 Parameters
753 ----------
754 dataframe : `pandas.DataFrame`
755 Input pandas dataframe.
757 Returns
758 -------
759 numpy_dict : `dict` [`str`, `numpy.ndarray`]
760 Dict with keys as the column names, values as the arrays.
761 """
762 return arrow_to_numpy_dict(pandas_to_arrow(dataframe))
765def numpy_to_astropy(np_array: np.ndarray) -> atable.Table:
766 """Convert a numpy table to an astropy table.
768 Parameters
769 ----------
770 np_array : `numpy.ndarray`
771 Input numpy array with multiple fields.
773 Returns
774 -------
775 astropy_table : `astropy.table.Table`
776 Converted astropy table.
777 """
778 from astropy.table import Table
780 return Table(data=np_array, copy=False)
783def arrow_schema_to_pandas_index(schema: pa.Schema) -> pd.Index | pd.MultiIndex:
784 """Convert an arrow schema to a pandas index/multiindex.
786 Parameters
787 ----------
788 schema : `pyarrow.Schema`
789 Input pyarrow schema.
791 Returns
792 -------
793 index : `pandas.Index` or `pandas.MultiIndex`
794 Converted pandas index.
795 """
796 import pandas as pd
798 if b"pandas" in schema.metadata:
799 md = json.loads(schema.metadata[b"pandas"])
800 indexes = md["column_indexes"]
801 len_indexes = len(indexes)
802 else:
803 len_indexes = 0
805 if len_indexes <= 1:
806 return pd.Index(name for name in schema.names if not name.startswith("__"))
807 else:
808 raw_columns = _split_multi_index_column_names(len(indexes), schema.names)
809 return pd.MultiIndex.from_tuples(raw_columns, names=[f["name"] for f in indexes])
812def arrow_schema_to_column_list(schema: pa.Schema) -> list[str]:
813 """Convert an arrow schema to a list of string column names.
815 Parameters
816 ----------
817 schema : `pyarrow.Schema`
818 Input pyarrow schema.
820 Returns
821 -------
822 column_list : `list` [`str`]
823 Converted list of column names.
824 """
825 return list(schema.names)
828class DataFrameSchema:
829 """Wrapper class for a schema for a pandas DataFrame.
831 Parameters
832 ----------
833 dataframe : `pandas.DataFrame`
834 Dataframe to turn into a schema.
835 """
837 def __init__(self, dataframe: pd.DataFrame) -> None:
838 self._schema = dataframe.loc[[False] * len(dataframe)]
840 @classmethod
841 def from_arrow(cls, schema: pa.Schema) -> DataFrameSchema:
842 """Convert an arrow schema into a `DataFrameSchema`.
844 Parameters
845 ----------
846 schema : `pyarrow.Schema`
847 The pyarrow schema to convert.
849 Returns
850 -------
851 dataframe_schema : `DataFrameSchema`
852 Converted dataframe schema.
853 """
854 empty_table = pa.Table.from_pylist([] * len(schema.names), schema=schema)
856 return cls(empty_table.to_pandas())
858 def to_arrow_schema(self) -> pa.Schema:
859 """Convert to an arrow schema.
861 Returns
862 -------
863 arrow_schema : `pyarrow.Schema`
864 Converted pyarrow schema.
865 """
866 arrow_table = pa.Table.from_pandas(self._schema)
868 return arrow_table.schema
870 def to_arrow_numpy_schema(self) -> ArrowNumpySchema:
871 """Convert to an `ArrowNumpySchema`.
873 Returns
874 -------
875 arrow_numpy_schema : `ArrowNumpySchema`
876 Converted arrow numpy schema.
877 """
878 return ArrowNumpySchema.from_arrow(self.to_arrow_schema())
880 def to_arrow_astropy_schema(self) -> ArrowAstropySchema:
881 """Convert to an ArrowAstropySchema.
883 Returns
884 -------
885 arrow_astropy_schema : `ArrowAstropySchema`
886 Converted arrow astropy schema.
887 """
888 return ArrowAstropySchema.from_arrow(self.to_arrow_schema())
890 @property
891 def schema(self) -> np.dtype:
892 return self._schema
894 def __repr__(self) -> str:
895 return repr(self._schema)
897 def __eq__(self, other: object) -> bool:
898 if not isinstance(other, DataFrameSchema):
899 return NotImplemented
901 return self._schema.equals(other._schema)
904class ArrowAstropySchema:
905 """Wrapper class for a schema for an astropy table.
907 Parameters
908 ----------
909 astropy_table : `astropy.table.Table`
910 Input astropy table.
911 """
913 def __init__(self, astropy_table: atable.Table) -> None:
914 self._schema = astropy_table[:0]
916 @classmethod
917 def from_arrow(cls, schema: pa.Schema) -> ArrowAstropySchema:
918 """Convert an arrow schema into a ArrowAstropySchema.
920 Parameters
921 ----------
922 schema : `pyarrow.Schema`
923 Input pyarrow schema.
925 Returns
926 -------
927 astropy_schema : `ArrowAstropySchema`
928 Converted arrow astropy schema.
929 """
930 import numpy as np
931 from astropy.table import Table
933 dtype = _schema_to_dtype_list(schema)
935 data = np.zeros(0, dtype=dtype)
936 astropy_table = Table(data=data)
938 _apply_astropy_metadata(astropy_table, schema)
940 return cls(astropy_table)
942 def to_arrow_schema(self) -> pa.Schema:
943 """Convert to an arrow schema.
945 Returns
946 -------
947 arrow_schema : `pyarrow.Schema`
948 Converted pyarrow schema.
949 """
950 return astropy_to_arrow(self._schema).schema
952 def to_dataframe_schema(self) -> DataFrameSchema:
953 """Convert to a DataFrameSchema.
955 Returns
956 -------
957 dataframe_schema : `DataFrameSchema`
958 Converted dataframe schema.
959 """
960 return DataFrameSchema.from_arrow(astropy_to_arrow(self._schema).schema)
962 def to_arrow_numpy_schema(self) -> ArrowNumpySchema:
963 """Convert to an `ArrowNumpySchema`.
965 Returns
966 -------
967 arrow_numpy_schema : `ArrowNumpySchema`
968 Converted arrow numpy schema.
969 """
970 return ArrowNumpySchema.from_arrow(astropy_to_arrow(self._schema).schema)
972 @property
973 def schema(self) -> atable.Table:
974 return self._schema
976 def __repr__(self) -> str:
977 return repr(self._schema)
979 def __eq__(self, other: object) -> bool:
980 if not isinstance(other, ArrowAstropySchema):
981 return NotImplemented
983 # If this comparison passes then the two tables have the
984 # same column names.
985 if self._schema.dtype != other._schema.dtype:
986 return False
988 for name in self._schema.columns:
989 if not self._schema[name].unit == other._schema[name].unit:
990 return False
991 if not self._schema[name].description == other._schema[name].description:
992 return False
993 if not self._schema[name].format == other._schema[name].format:
994 return False
996 return True
999class ArrowNumpySchema:
1000 """Wrapper class for a schema for a numpy ndarray.
1002 Parameters
1003 ----------
1004 numpy_dtype : `numpy.dtype`
1005 Numpy dtype to convert.
1006 """
1008 def __init__(self, numpy_dtype: np.dtype) -> None:
1009 self._dtype = numpy_dtype
1011 @classmethod
1012 def from_arrow(cls, schema: pa.Schema) -> ArrowNumpySchema:
1013 """Convert an arrow schema into an `ArrowNumpySchema`.
1015 Parameters
1016 ----------
1017 schema : `pyarrow.Schema`
1018 Pyarrow schema to convert.
1020 Returns
1021 -------
1022 numpy_schema : `ArrowNumpySchema`
1023 Converted arrow numpy schema.
1024 """
1025 import numpy as np
1027 dtype = _schema_to_dtype_list(schema)
1029 return cls(np.dtype(dtype))
1031 def to_arrow_astropy_schema(self) -> ArrowAstropySchema:
1032 """Convert to an `ArrowAstropySchema`.
1034 Returns
1035 -------
1036 astropy_schema : `ArrowAstropySchema`
1037 Converted arrow astropy schema.
1038 """
1039 import numpy as np
1041 return ArrowAstropySchema.from_arrow(numpy_to_arrow(np.zeros(0, dtype=self._dtype)).schema)
1043 def to_dataframe_schema(self) -> DataFrameSchema:
1044 """Convert to a `DataFrameSchema`.
1046 Returns
1047 -------
1048 dataframe_schema : `DataFrameSchema`
1049 Converted dataframe schema.
1050 """
1051 import numpy as np
1053 return DataFrameSchema.from_arrow(numpy_to_arrow(np.zeros(0, dtype=self._dtype)).schema)
1055 def to_arrow_schema(self) -> pa.Schema:
1056 """Convert to a `pyarrow.Schema`.
1058 Returns
1059 -------
1060 arrow_schema : `pyarrow.Schema`
1061 Converted pyarrow schema.
1062 """
1063 import numpy as np
1065 return numpy_to_arrow(np.zeros(0, dtype=self._dtype)).schema
1067 @property
1068 def schema(self) -> np.dtype:
1069 return self._dtype
1071 def __repr__(self) -> str:
1072 return repr(self._dtype)
1074 def __eq__(self, other: object) -> bool:
1075 if not isinstance(other, ArrowNumpySchema):
1076 return NotImplemented
1078 if not self._dtype == other._dtype:
1079 return False
1081 return True
1084def _split_multi_index_column_names(n: int, names: Iterable[str]) -> list[Sequence[str]]:
1085 """Split a string that represents a multi-index column.
1087 PyArrow maps Pandas' multi-index column names (which are tuples in Python)
1088 to flat strings on disk. This routine exists to reconstruct the original
1089 tuple.
1091 Parameters
1092 ----------
1093 n : `int`
1094 Number of levels in the `pandas.MultiIndex` that is being
1095 reconstructed.
1096 names : `~collections.abc.Iterable` [`str`]
1097 Strings to be split.
1099 Returns
1100 -------
1101 column_names : `list` [`tuple` [`str`]]
1102 A list of multi-index column name tuples.
1103 """
1104 column_names: list[Sequence[str]] = []
1106 pattern = re.compile(r"\({}\)".format(", ".join(["'(.*)'"] * n)))
1107 for name in names:
1108 m = re.search(pattern, name)
1109 if m is not None:
1110 column_names.append(m.groups())
1112 return column_names
1115def _standardize_multi_index_columns(
1116 pd_index: pd.MultiIndex,
1117 columns: Any,
1118 stringify: bool = True,
1119) -> list[str | Sequence[Any]]:
1120 """Transform a dictionary/iterable index from a multi-index column list
1121 into a string directly understandable by PyArrow.
1123 Parameters
1124 ----------
1125 pd_index : `pandas.MultiIndex`
1126 Pandas multi-index.
1127 columns : `list` [`tuple`] or `dict` [`str`, `str` or `list` [`str`]]
1128 Columns to standardize.
1129 stringify : `bool`, optional
1130 Should the column names be stringified?
1132 Returns
1133 -------
1134 names : `list` [`str`]
1135 Stringified representation of a multi-index column name.
1136 """
1137 index_level_names = tuple(pd_index.names)
1139 names: list[str | Sequence[Any]] = []
1141 if isinstance(columns, list):
1142 for requested in columns:
1143 if not isinstance(requested, tuple):
1144 raise ValueError(
1145 "Columns parameter for multi-index data frame must be a dictionary or list of tuples. "
1146 f"Instead got a {get_full_type_name(requested)}."
1147 )
1148 if stringify:
1149 names.append(str(requested))
1150 else:
1151 names.append(requested)
1152 else:
1153 if not isinstance(columns, collections.abc.Mapping): 1153 ↛ 1154line 1153 didn't jump to line 1154 because the condition on line 1153 was never true
1154 raise ValueError(
1155 "Columns parameter for multi-index data frame must be a dictionary or list of tuples. "
1156 f"Instead got a {get_full_type_name(columns)}."
1157 )
1158 if not set(index_level_names).issuperset(columns.keys()): 1158 ↛ 1159line 1158 didn't jump to line 1159 because the condition on line 1158 was never true
1159 raise ValueError(
1160 f"Cannot use dict with keys {set(columns.keys())} to select columns from {index_level_names}."
1161 )
1162 factors = [
1163 ensure_iterable(columns.get(level, pd_index.levels[i]))
1164 for i, level in enumerate(index_level_names)
1165 ]
1166 for requested in itertools.product(*factors):
1167 for i, value in enumerate(requested):
1168 if value not in pd_index.levels[i]: 1168 ↛ 1169line 1168 didn't jump to line 1169 because the condition on line 1168 was never true
1169 raise ValueError(f"Unrecognized value {value!r} for index {index_level_names[i]!r}.")
1170 if stringify:
1171 names.append(str(requested))
1172 else:
1173 names.append(requested)
1175 return names
1178def _apply_astropy_metadata(astropy_table: atable.Table, arrow_schema: pa.Schema) -> None:
1179 """Apply any astropy metadata from the schema metadata.
1181 Parameters
1182 ----------
1183 astropy_table : `astropy.table.Table`
1184 Table to apply metadata.
1185 arrow_schema : `pyarrow.Schema`
1186 Arrow schema with metadata.
1187 """
1188 from astropy.table import meta
1190 metadata = arrow_schema.metadata if arrow_schema.metadata is not None else {}
1192 # Check if we have a special astropy metadata header yaml.
1193 meta_yaml = metadata.get(b"table_meta_yaml", None)
1194 if meta_yaml:
1195 meta_yaml = meta_yaml.decode("UTF8").split("\n")
1196 meta_hdr = meta.get_header_from_yaml(meta_yaml)
1198 # Set description, format, unit, meta from the column
1199 # metadata that was serialized with the table.
1200 header_cols = {x["name"]: x for x in meta_hdr["datatype"]}
1201 for col in astropy_table.columns.values():
1202 for attr in ("description", "format", "unit", "meta"):
1203 if attr in header_cols[col.name]:
1204 setattr(col, attr, header_cols[col.name][attr])
1206 if "meta" in meta_hdr:
1207 astropy_table.meta.update(meta_hdr["meta"])
1208 else:
1209 # If we don't have astropy header data, we may have arrow field
1210 # metadata.
1211 for name in arrow_schema.names:
1212 field_metadata = arrow_schema.field(name).metadata
1213 if field_metadata is None:
1214 continue
1215 if (
1216 b"description" in field_metadata
1217 and (description := field_metadata[b"description"].decode("UTF-8")) != ""
1218 ):
1219 astropy_table[name].description = description
1220 if b"unit" in field_metadata and (unit := field_metadata[b"unit"].decode("UTF-8")) != "":
1221 astropy_table[name].unit = unit
1223 # Ensure that the special ASTROPY_PANDAS_INDEX_KEY is propagated to
1224 # the table metadata.
1225 if index_key := metadata.get(ASTROPY_PANDAS_INDEX_KEY.encode(), None):
1226 astropy_table.meta[ASTROPY_PANDAS_INDEX_KEY] = index_key.decode("UTF-8")
1229def _arrow_string_to_numpy_dtype(
1230 schema: pa.Schema, name: str, numpy_column: np.ndarray | None = None, default_length: int = 10
1231) -> str:
1232 """Get the numpy dtype string associated with an arrow column.
1234 Parameters
1235 ----------
1236 schema : `pyarrow.Schema`
1237 Arrow table schema.
1238 name : `str`
1239 Column name.
1240 numpy_column : `numpy.ndarray`, optional
1241 Column to determine numpy string dtype.
1242 default_length : `int`, optional
1243 Default string length when not in metadata or can be inferred
1244 from column.
1246 Returns
1247 -------
1248 dtype_str : `str`
1249 Numpy dtype string.
1250 """
1251 # Special-case for string and binary columns
1252 md_name = f"lsst::arrow::len::{name}"
1253 strlen = default_length
1254 metadata = schema.metadata if schema.metadata is not None else {}
1255 if (encoded := md_name.encode("UTF-8")) in metadata:
1256 # String/bytes length from header.
1257 strlen = int(schema.metadata[encoded])
1258 elif numpy_column is not None and len(numpy_column) > 0:
1259 lengths = [len(row) for row in numpy_column if row]
1260 strlen = max(lengths) if lengths else 0
1262 dtype = f"U{strlen}" if schema.field(name).type == pa.string() else f"|S{strlen}"
1264 return dtype
1267def _append_numpy_string_metadata(metadata: dict[bytes, str], name: str, dtype: np.dtype) -> None:
1268 """Append numpy string length keys to arrow metadata.
1270 All column types are handled, but the metadata is only modified for
1271 string and byte columns.
1273 Parameters
1274 ----------
1275 metadata : `dict` [`bytes`, `str`]
1276 Metadata dictionary; modified in place.
1277 name : `str`
1278 Column name.
1279 dtype : `np.dtype`
1280 Numpy dtype.
1281 """
1282 import numpy as np
1284 if dtype.type is np.str_:
1285 metadata[f"lsst::arrow::len::{name}".encode()] = str(dtype.itemsize // 4)
1286 metadata[f"table::len::{name}".encode()] = str(dtype.itemsize // 4)
1287 elif dtype.type is np.bytes_:
1288 metadata[f"lsst::arrow::len::{name}".encode()] = str(dtype.itemsize)
1289 metadata[f"table::len::{name}".encode()] = str(dtype.itemsize)
1292def _append_numpy_multidim_metadata(metadata: dict[bytes, str], name: str, dtype: np.dtype) -> None:
1293 """Append numpy multi-dimensional shapes to arrow metadata.
1295 All column types are handled, but the metadata is only modified for
1296 multi-dimensional columns.
1298 Parameters
1299 ----------
1300 metadata : `dict` [`bytes`, `str`]
1301 Metadata dictionary; modified in place.
1302 name : `str`
1303 Column name.
1304 dtype : `np.dtype`
1305 Numpy dtype.
1306 """
1307 if len(dtype.shape) > 1:
1308 metadata[f"lsst::arrow::shape::{name}".encode()] = str(dtype.shape)
1311def _multidim_shape_from_metadata(metadata: dict[bytes, bytes], list_size: int, name: str) -> tuple[int, ...]:
1312 """Retrieve the shape from the metadata, if available.
1314 Parameters
1315 ----------
1316 metadata : `dict` [`bytes`, `bytes`]
1317 Metadata dictionary.
1318 list_size : `int`
1319 Size of the list datatype.
1320 name : `str`
1321 Column name.
1323 Returns
1324 -------
1325 shape : `tuple` [`int`]
1326 Shape associated with the column.
1328 Raises
1329 ------
1330 RuntimeError
1331 Raised if metadata is found but has incorrect format.
1332 """
1333 md_name = f"lsst::arrow::shape::{name}"
1334 if (encoded := md_name.encode("UTF-8")) in metadata:
1335 groups = re.search(r"\((.*)\)", metadata[encoded].decode("UTF-8"))
1336 if groups is None: 1336 ↛ 1337line 1336 didn't jump to line 1337 because the condition on line 1336 was never true
1337 raise RuntimeError("Illegal value found in metadata.")
1338 shape = tuple(int(x) for x in groups[1].split(",") if x != "")
1339 else:
1340 shape = (list_size,)
1342 return shape
1345def _schema_to_dtype_list(schema: pa.Schema) -> list[tuple[str, tuple[Any] | str]]:
1346 """Convert a pyarrow schema to a numpy dtype.
1348 Parameters
1349 ----------
1350 schema : `pyarrow.Schema`
1351 Input pyarrow schema.
1353 Returns
1354 -------
1355 dtype_list: `list` [`tuple`]
1356 A list with name, type pairs.
1357 """
1358 metadata = schema.metadata if schema.metadata is not None else {}
1360 dtype: list[Any] = []
1361 for name in schema.names:
1362 t = schema.field(name).type
1363 if isinstance(t, pa.FixedSizeListType):
1364 shape = _multidim_shape_from_metadata(metadata, t.list_size, name)
1365 dtype.append((name, (t.value_type.to_pandas_dtype(), shape)))
1366 elif t not in (pa.string(), pa.binary()):
1367 dtype.append((name, t.to_pandas_dtype()))
1368 else:
1369 dtype.append((name, _arrow_string_to_numpy_dtype(schema, name)))
1371 return dtype
1374def _numpy_dtype_to_arrow_types(dtype: np.dtype) -> list[Any]:
1375 """Convert a numpy dtype to a list of arrow types.
1377 Parameters
1378 ----------
1379 dtype : `numpy.dtype`
1380 Numpy dtype to convert.
1382 Returns
1383 -------
1384 type_list : `list` [`object`]
1385 Converted list of arrow types.
1386 """
1387 from math import prod
1389 import numpy as np
1391 type_list: list[Any] = []
1392 if dtype.names is None: 1392 ↛ 1393line 1392 didn't jump to line 1393 because the condition on line 1392 was never true
1393 return type_list
1395 for name in dtype.names:
1396 dt = dtype[name]
1397 arrow_type: Any
1398 if len(dt.shape) > 0:
1399 arrow_type = pa.list_(
1400 pa.from_numpy_dtype(cast(tuple[np.dtype, tuple[int, ...]], dt.subdtype)[0].type),
1401 prod(dt.shape),
1402 )
1403 elif dt.type == np.datetime64:
1404 time_unit = "ns" if "ns" in dt.str else "us"
1405 # The pa.timestamp() is the correct datatype to round-trip
1406 # a numpy datetime64[ns] or datetime[us] array.
1407 arrow_type = pa.timestamp(time_unit)
1408 else:
1409 try:
1410 arrow_type = pa.from_numpy_dtype(dt.type)
1411 except pa.ArrowNotImplementedError as e:
1412 msg = f"Could not serialize column {name} (type {str(dt)}) to Parquet."
1413 if dt == np.dtype("O"): 1413 ↛ 1415line 1413 didn't jump to line 1415 because the condition on line 1413 was always true
1414 msg += " This is usually because the column is mixed type or has uneven length rows."
1415 e.add_note(msg)
1416 raise
1417 type_list.append((name, arrow_type))
1419 return type_list
1422def _numpy_dict_to_dtype(numpy_dict: dict[str, np.ndarray]) -> tuple[np.dtype, int]:
1423 """Extract equivalent table dtype from dict of numpy arrays.
1425 Parameters
1426 ----------
1427 numpy_dict : `dict` [`str`, `numpy.ndarray`]
1428 Dict with keys as the column names, values as the arrays.
1430 Returns
1431 -------
1432 dtype : `numpy.dtype`
1433 dtype of equivalent table.
1434 rowcount : `int`
1435 Number of rows in the table.
1437 Raises
1438 ------
1439 ValueError if columns in numpy_dict have unequal numbers of rows.
1440 """
1441 import numpy as np
1443 dtype_list: list[tuple] = []
1444 rowcount = 0
1445 for name, col in numpy_dict.items():
1446 if rowcount == 0:
1447 rowcount = len(col)
1448 if len(col) != rowcount:
1449 raise ValueError(f"Column {name} has a different number of rows.")
1450 if len(col.shape) == 1:
1451 dtype_list.append((name, col.dtype))
1452 else:
1453 dtype_list.append((name, (col.dtype, col.shape[1:])))
1454 dtype = np.dtype(dtype_list)
1456 return (dtype, rowcount)
1459def _numpy_style_arrays_to_arrow_arrays(
1460 dtype: np.dtype,
1461 rowcount: int,
1462 np_style_arrays: dict[str, np.ndarray] | np.ndarray | atable.Table,
1463 schema: pa.Schema,
1464) -> list[pa.Array]:
1465 """Convert numpy-style arrays to arrow arrays.
1467 Parameters
1468 ----------
1469 dtype : `numpy.dtype`
1470 Numpy dtype of input table/arrays.
1471 rowcount : `int`
1472 Number of rows in input table/arrays.
1473 np_style_arrays : `dict` [`str`, `np.ndarray`] or `np.ndarray`
1474 or `astropy.table.Table`
1475 Arrays to convert to arrow.
1476 schema : `pyarrow.Schema`
1477 Schema of arrow table.
1479 Returns
1480 -------
1481 arrow_arrays : `list` [`pyarrow.Array`]
1482 List of converted pyarrow arrays.
1483 """
1484 import numpy as np
1486 arrow_arrays: list[pa.Array] = []
1487 if dtype.names is None: 1487 ↛ 1488line 1487 didn't jump to line 1488 because the condition on line 1487 was never true
1488 return arrow_arrays
1490 for name in dtype.names:
1491 dt = dtype[name]
1492 val: Any
1493 if len(dt.shape) > 0:
1494 if rowcount > 0:
1495 val = np.split(np_style_arrays[name].ravel(), rowcount)
1496 else:
1497 val = []
1498 else:
1499 val = np_style_arrays[name]
1501 try:
1502 arrow_arrays.append(pa.array(val, type=schema.field(name).type))
1503 except pa.ArrowNotImplementedError as err:
1504 # Check if val is big-endian.
1505 if (np.little_endian and val.dtype.byteorder == ">") or ( 1505 ↛ 1514line 1505 didn't jump to line 1514 because the condition on line 1505 was always true
1506 not np.little_endian and val.dtype.byteorder == "="
1507 ):
1508 # We need to convert the array to little-endian.
1509 val2 = val.byteswap()
1510 val2.dtype = val2.dtype.newbyteorder("<")
1511 arrow_arrays.append(pa.array(val2, type=schema.field(name).type))
1512 else:
1513 # This failed for some other reason so raise the exception.
1514 raise err
1516 return arrow_arrays
1519def compute_row_group_size(schema: pa.Schema, target_size: int = TARGET_ROW_GROUP_BYTES) -> int:
1520 """Compute approximate row group size for a given arrow schema.
1522 Given a schema, this routine will compute the number of rows in a row group
1523 that targets the persisted size on disk (or smaller). The exact size on
1524 disk depends on the compression settings and ratios; typical binary data
1525 tables will have around 15-20% compression with the pyarrow default
1526 ``snappy`` compression algorithm.
1528 Parameters
1529 ----------
1530 schema : `pyarrow.Schema`
1531 Arrow table schema.
1532 target_size : `int`, optional
1533 The target size (in bytes).
1535 Returns
1536 -------
1537 row_group_size : `int`
1538 Number of rows per row group to hit the target size.
1539 """
1540 bit_width = 0
1542 metadata = schema.metadata if schema.metadata is not None else {}
1544 for name in schema.names:
1545 t = schema.field(name).type
1547 if t in (pa.string(), pa.binary()):
1548 md_name = f"lsst::arrow::len::{name}"
1550 if (encoded := md_name.encode("UTF-8")) in metadata:
1551 # String/bytes length from header.
1552 strlen = int(schema.metadata[encoded])
1553 else:
1554 # We don't know the string width, so guess something.
1555 strlen = 10
1557 # Assuming UTF-8 encoding, and very few wide characters.
1558 t_width = 8 * strlen
1559 elif isinstance(t, pa.FixedSizeListType):
1560 if t.value_type == pa.null(): 1560 ↛ 1561line 1560 didn't jump to line 1561 because the condition on line 1560 was never true
1561 t_width = 0
1562 else:
1563 t_width = t.list_size * t.value_type.bit_width
1564 elif t == pa.null():
1565 t_width = 0
1566 elif isinstance(t, pa.ListType):
1567 if t.value_type == pa.null():
1568 t_width = 0
1569 else:
1570 # This is a variable length list, just choose
1571 # something arbitrary.
1572 t_width = 10 * t.value_type.bit_width
1573 else:
1574 t_width = t.bit_width
1576 bit_width += t_width
1578 # Insist it is at least 1 byte wide to avoid any divide-by-zero errors.
1579 if bit_width < 8:
1580 bit_width = 8
1582 byte_width = bit_width // 8
1584 return target_size // byte_width