Coverage for python / lsst / daf / butler / formatters / parquet.py: 12%

553 statements  

« prev     ^ index     » next       coverage.py v7.13.1, created at 2026-01-22 01:49 -0800

1# This file is part of daf_butler. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (http://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This software is dual licensed under the GNU General Public License and also 

10# under a 3-clause BSD license. Recipients may choose which of these licenses 

11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, 

12# respectively. If you choose the GPL option then the following text applies 

13# (but note that there is still no warranty even if you opt for BSD instead): 

14# 

15# This program is free software: you can redistribute it and/or modify 

16# it under the terms of the GNU General Public License as published by 

17# the Free Software Foundation, either version 3 of the License, or 

18# (at your option) any later version. 

19# 

20# This program is distributed in the hope that it will be useful, 

21# but WITHOUT ANY WARRANTY; without even the implied warranty of 

22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

23# GNU General Public License for more details. 

24# 

25# You should have received a copy of the GNU General Public License 

26# along with this program. If not, see <http://www.gnu.org/licenses/>. 

27 

28from __future__ import annotations 

29 

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) 

50 

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 

60 

61import pyarrow as pa 

62import pyarrow.parquet as pq 

63 

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 

69 

70log = logging.getLogger(__name__) 

71 

72if TYPE_CHECKING: 

73 import astropy.table as atable 

74 import numpy as np 

75 import pandas as pd 

76 

77 try: 

78 import fsspec 

79 from fsspec.spec import AbstractFileSystem 

80 except ImportError: 

81 fsspec = None 

82 AbstractFileSystem = type 

83 

84TARGET_ROW_GROUP_BYTES = 1_000_000_000 

85ASTROPY_PANDAS_INDEX_KEY = "lsst::arrow::astropy_pandas_index" 

86 

87 

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 

96 

97 

98class ParquetFormatter(FormatterV2): 

99 """Interface for reading and writing Arrow Table objects to and from 

100 Parquet files. 

101 """ 

102 

103 default_extension = ".parq" 

104 can_read_from_uri = True 

105 can_read_from_local_file = True 

106 

107 def can_accept(self, in_memory_dataset: Any) -> bool: 

108 # Docstring inherited. 

109 return _checkArrowCompatibleType(in_memory_dataset) is not None 

110 

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 

120 

121 return self._read_parquet(path=path, fs=fs, component=component, expected_size=expected_size) 

122 

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) 

126 

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) 

136 

137 schema_names = ["ArrowSchema", "DataFrameSchema", "ArrowAstropySchema", "ArrowNumpySchema"] 

138 

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"]) 

147 

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 ) 

155 

156 return len(temp_table[schema.names[0]]) 

157 

158 par_columns = None 

159 if self.file_descriptor.parameters: 

160 par_columns = self.file_descriptor.parameters.pop("columns", None) 

161 if par_columns: 

162 has_pandas_multi_index = False 

163 if schema.metadata and b"pandas" in schema.metadata: 

164 md = json.loads(schema.metadata[b"pandas"]) 

165 if len(md["column_indexes"]) > 1: 

166 has_pandas_multi_index = True 

167 

168 if not has_pandas_multi_index: 

169 # Ensure uniqueness, keeping order. 

170 par_columns_in = list(dict.fromkeys(ensure_iterable(par_columns))) 

171 file_columns = [name for name in schema.names if not name.startswith("__")] 

172 

173 # Do case-sensitive glob-style matching, again ensuring 

174 # uniqueness and ordering. 

175 par_columns = {} 

176 for par_column in par_columns_in: 

177 found = False 

178 for file_column in file_columns: 

179 if fnmatchcase(file_column, par_column): 

180 found = True 

181 par_columns[file_column] = True 

182 if not found: 

183 raise ValueError( 

184 f"Column {par_column} specified in parameters not available in parquet file." 

185 ) 

186 par_columns = list(par_columns.keys()) 

187 else: 

188 par_columns = _standardize_multi_index_columns( 

189 arrow_schema_to_pandas_index(schema), 

190 par_columns, 

191 ) 

192 

193 if len(self.file_descriptor.parameters): 

194 raise ValueError( 

195 f"Unsupported parameters {self.file_descriptor.parameters} in ArrowTable read." 

196 ) 

197 

198 metadata = schema.metadata if schema.metadata is not None else {} 

199 with generic_open(path, fs) as handle: 

200 arrow_table = pq.read_table( 

201 handle, 

202 columns=par_columns, 

203 use_threads=False, 

204 use_pandas_metadata=(b"pandas" in metadata), 

205 ) 

206 

207 return arrow_table 

208 

209 def add_provenance(self, in_memory_dataset: Any, provenance: DatasetProvenance | None = None) -> Any: 

210 return _add_arrow_provenance(in_memory_dataset, self.dataset_ref, provenance) 

211 

212 def write_local_file(self, in_memory_dataset: Any, uri: ResourcePath) -> None: 

213 if isinstance(in_memory_dataset, pa.Schema): 

214 pq.write_metadata(in_memory_dataset, uri.ospath) 

215 return 

216 

217 type_string = _checkArrowCompatibleType(in_memory_dataset) 

218 

219 if type_string is None: 

220 raise ValueError( 

221 f"Unsupported type {get_full_type_name(in_memory_dataset)} of " 

222 "inMemoryDataset for ParquetFormatter." 

223 ) 

224 

225 if type_string == "arrow": 

226 arrow_table = in_memory_dataset 

227 elif type_string == "astropy": 

228 arrow_table = astropy_to_arrow(in_memory_dataset) 

229 elif type_string == "numpy": 

230 arrow_table = numpy_to_arrow(in_memory_dataset) 

231 elif type_string == "numpydict": 

232 arrow_table = numpy_dict_to_arrow(in_memory_dataset) 

233 else: 

234 arrow_table = pandas_to_arrow(in_memory_dataset) 

235 

236 row_group_size = compute_row_group_size(arrow_table.schema) 

237 

238 pq.write_table(arrow_table, uri.ospath, row_group_size=row_group_size) 

239 

240 

241def arrow_to_pandas(arrow_table: pa.Table) -> pd.DataFrame: 

242 """Convert a pyarrow table to a pandas DataFrame. 

243 

244 Parameters 

245 ---------- 

246 arrow_table : `pyarrow.Table` 

247 Input arrow table to convert. If the table has ``pandas`` metadata 

248 in the schema it will be used in the construction of the 

249 ``DataFrame``. 

250 

251 Returns 

252 ------- 

253 dataframe : `pandas.DataFrame` 

254 Converted pandas dataframe. 

255 """ 

256 dataframe = arrow_table.to_pandas(use_threads=False, integer_object_nulls=True) 

257 

258 metadata = arrow_table.schema.metadata if arrow_table.schema.metadata is not None else {} 

259 if (key := ASTROPY_PANDAS_INDEX_KEY.encode()) in metadata: 

260 pandas_index = metadata[key].decode("UTF8") 

261 if pandas_index in arrow_table.schema.names: 

262 dataframe.set_index(pandas_index, inplace=True) 

263 else: 

264 log.warning( 

265 "Index column ``%s`` not available for arrow table conversion to DataFrame", 

266 pandas_index, 

267 ) 

268 

269 return dataframe 

270 

271 

272def arrow_to_astropy(arrow_table: pa.Table) -> atable.Table: 

273 """Convert a pyarrow table to an `astropy.Table`. 

274 

275 Parameters 

276 ---------- 

277 arrow_table : `pyarrow.Table` 

278 Input arrow table to convert. If the table has astropy unit 

279 metadata in the schema it will be used in the construction 

280 of the ``astropy.Table``. 

281 

282 Returns 

283 ------- 

284 table : `astropy.Table` 

285 Converted astropy table. 

286 """ 

287 from astropy.table import Table 

288 

289 astropy_table = Table(arrow_to_numpy_dict(arrow_table)) 

290 

291 _apply_astropy_metadata(astropy_table, arrow_table.schema) 

292 

293 if (key := ASTROPY_PANDAS_INDEX_KEY) in astropy_table.meta: 

294 if astropy_table.meta[key] not in astropy_table.columns: 

295 astropy_table.meta.pop(key) 

296 

297 return astropy_table 

298 

299 

300def arrow_to_numpy(arrow_table: pa.Table) -> np.ndarray | np.ma.MaskedArray: 

301 """Convert a pyarrow table to a structured numpy array. 

302 

303 Parameters 

304 ---------- 

305 arrow_table : `pyarrow.Table` 

306 Input arrow table. 

307 

308 Returns 

309 ------- 

310 array : `numpy.ndarray` or `numpy.ma.MaskedArray` (N,) 

311 Numpy array table with N rows and the same column names 

312 as the input arrow table. Will be masked records if any values 

313 in the table are null. 

314 """ 

315 import numpy as np 

316 

317 numpy_dict = arrow_to_numpy_dict(arrow_table) 

318 

319 has_mask = False 

320 dtype: list[tuple] = [] 

321 for name, col in numpy_dict.items(): 

322 if len(shape := numpy_dict[name].shape) <= 1: 

323 dtype.append((name, col.dtype)) 

324 else: 

325 dtype.append((name, (col.dtype, shape[1:]))) 

326 

327 if not has_mask and isinstance(col, np.ma.MaskedArray): 

328 has_mask = True 

329 

330 if has_mask: 

331 array = np.ma.mrecords.fromarrays(numpy_dict.values(), dtype=dtype) 

332 else: 

333 array = np.rec.fromarrays(numpy_dict.values(), dtype=dtype) 

334 return array 

335 

336 

337def arrow_to_numpy_dict(arrow_table: pa.Table) -> dict[str, np.ndarray]: 

338 """Convert a pyarrow table to a dict of numpy arrays. 

339 

340 Parameters 

341 ---------- 

342 arrow_table : `pyarrow.Table` 

343 Input arrow table. 

344 

345 Returns 

346 ------- 

347 numpy_dict : `dict` [`str`, `numpy.ndarray`] 

348 Dict with keys as the column names, values as the arrays. 

349 """ 

350 import numpy as np 

351 

352 schema = arrow_table.schema 

353 metadata = schema.metadata if schema.metadata is not None else {} 

354 

355 numpy_dict = {} 

356 

357 for name in schema.names: 

358 t = schema.field(name).type 

359 

360 if arrow_table[name].null_count == 0: 

361 # Regular non-masked column 

362 col = arrow_table[name].to_numpy() 

363 else: 

364 # For a masked column, we need to ask arrow to fill the null 

365 # values with an appropriately typed value before conversion. 

366 # Then we apply the mask to get a masked array of the correct type. 

367 null_value: Any 

368 match t: 

369 case t if t in (pa.float64(), pa.float32(), pa.float16()): 

370 null_value = np.nan 

371 case t if t in (pa.int64(), pa.int32(), pa.int16(), pa.int8()): 

372 null_value = -1 

373 case t if t in (pa.bool_(),): 

374 null_value = True 

375 case t if t in (pa.string(), pa.binary()): 

376 null_value = "" 

377 case _: 

378 # This is the fallback for unsigned ints in particular. 

379 null_value = 0 

380 

381 col = np.ma.masked_array( 

382 data=arrow_table[name].fill_null(null_value).to_numpy(), 

383 mask=arrow_table[name].is_null().to_numpy(), 

384 fill_value=null_value, 

385 ) 

386 

387 if t in (pa.string(), pa.binary()): 

388 col = col.astype(_arrow_string_to_numpy_dtype(schema, name, col)) 

389 elif isinstance(t, pa.FixedSizeListType): 

390 if len(col) > 0: 

391 col = np.stack(col) 

392 else: 

393 # this is an empty column, and needs to be coerced to type. 

394 col = col.astype(t.value_type.to_pandas_dtype()) 

395 

396 shape = _multidim_shape_from_metadata(metadata, t.list_size, name) 

397 col = col.reshape((len(arrow_table), *shape)) 

398 

399 numpy_dict[name] = col 

400 

401 return numpy_dict 

402 

403 

404def _numpy_dict_to_numpy(numpy_dict: dict[str, np.ndarray]) -> np.ndarray: 

405 """Convert a dict of numpy arrays to a structured numpy array. 

406 

407 Parameters 

408 ---------- 

409 numpy_dict : `dict` [`str`, `numpy.ndarray`] 

410 Dict with keys as the column names, values as the arrays. 

411 

412 Returns 

413 ------- 

414 array : `numpy.ndarray` (N,) 

415 Numpy array table with N rows and columns names from the dict keys. 

416 """ 

417 return arrow_to_numpy(numpy_dict_to_arrow(numpy_dict)) 

418 

419 

420def _numpy_to_numpy_dict(np_array: np.ndarray) -> dict[str, np.ndarray]: 

421 """Convert a structured numpy array to a dict of numpy arrays. 

422 

423 Parameters 

424 ---------- 

425 np_array : `numpy.ndarray` 

426 Input numpy array with multiple fields. 

427 

428 Returns 

429 ------- 

430 numpy_dict : `dict` [`str`, `numpy.ndarray`] 

431 Dict with keys as the column names, values as the arrays. 

432 """ 

433 return arrow_to_numpy_dict(numpy_to_arrow(np_array)) 

434 

435 

436def numpy_to_arrow(np_array: np.ndarray) -> pa.Table: 

437 """Convert a numpy array table to an arrow table. 

438 

439 Parameters 

440 ---------- 

441 np_array : `numpy.ndarray` 

442 Input numpy array with multiple fields. 

443 

444 Returns 

445 ------- 

446 arrow_table : `pyarrow.Table` 

447 Converted arrow table. 

448 """ 

449 type_list = _numpy_dtype_to_arrow_types(np_array.dtype) 

450 

451 md = {} 

452 md[b"lsst::arrow::rowcount"] = str(len(np_array)) 

453 

454 names = np_array.dtype.names 

455 if names is None: 

456 names = () 

457 

458 for name in names: 

459 _append_numpy_string_metadata(md, name, np_array.dtype[name]) 

460 _append_numpy_multidim_metadata(md, name, np_array.dtype[name]) 

461 

462 schema = pa.schema(type_list, metadata=md) 

463 

464 arrays = _numpy_style_arrays_to_arrow_arrays( 

465 np_array.dtype, 

466 len(np_array), 

467 np_array, 

468 schema, 

469 ) 

470 

471 arrow_table = pa.Table.from_arrays(arrays, schema=schema) 

472 

473 return arrow_table 

474 

475 

476def numpy_dict_to_arrow(numpy_dict: dict[str, np.ndarray]) -> pa.Table: 

477 """Convert a dict of numpy arrays to an arrow table. 

478 

479 Parameters 

480 ---------- 

481 numpy_dict : `dict` [`str`, `numpy.ndarray`] 

482 Dict with keys as the column names, values as the arrays. 

483 

484 Returns 

485 ------- 

486 arrow_table : `pyarrow.Table` 

487 Converted arrow table. 

488 

489 Raises 

490 ------ 

491 ValueError if columns in numpy_dict have unequal numbers of rows. 

492 """ 

493 dtype, rowcount = _numpy_dict_to_dtype(numpy_dict) 

494 type_list = _numpy_dtype_to_arrow_types(dtype) 

495 

496 md = {} 

497 md[b"lsst::arrow::rowcount"] = str(rowcount) 

498 

499 if dtype.names is not None: 

500 for name in dtype.names: 

501 _append_numpy_string_metadata(md, name, dtype[name]) 

502 _append_numpy_multidim_metadata(md, name, dtype[name]) 

503 

504 schema = pa.schema(type_list, metadata=md) 

505 

506 arrays = _numpy_style_arrays_to_arrow_arrays( 

507 dtype, 

508 rowcount, 

509 numpy_dict, 

510 schema, 

511 ) 

512 

513 arrow_table = pa.Table.from_arrays(arrays, schema=schema) 

514 

515 return arrow_table 

516 

517 

518def astropy_to_arrow(astropy_table: atable.Table) -> pa.Table: 

519 """Convert an astropy table to an arrow table. 

520 

521 Parameters 

522 ---------- 

523 astropy_table : `astropy.Table` 

524 Input astropy table. 

525 

526 Returns 

527 ------- 

528 arrow_table : `pyarrow.Table` 

529 Converted arrow table. 

530 """ 

531 from astropy.table import meta 

532 

533 type_list = _numpy_dtype_to_arrow_types(astropy_table.dtype) 

534 

535 md = {} 

536 md[b"lsst::arrow::rowcount"] = str(len(astropy_table)) 

537 

538 if (key := ASTROPY_PANDAS_INDEX_KEY) in astropy_table.meta: 

539 md[key.encode()] = astropy_table.meta[key] 

540 

541 for name in astropy_table.dtype.names: 

542 _append_numpy_string_metadata(md, name, astropy_table.dtype[name]) 

543 _append_numpy_multidim_metadata(md, name, astropy_table.dtype[name]) 

544 

545 meta_yaml = meta.get_yaml_from_table(astropy_table) 

546 meta_yaml_str = "\n".join(meta_yaml) 

547 md[b"table_meta_yaml"] = meta_yaml_str 

548 

549 # Convert type list to fields with metadata. 

550 fields = [] 

551 for name, pa_type in type_list: 

552 field_metadata = {} 

553 if description := astropy_table[name].description: 

554 field_metadata["description"] = description 

555 if unit := astropy_table[name].unit: 

556 field_metadata["unit"] = str(unit) 

557 fields.append( 

558 pa.field( 

559 name, 

560 pa_type, 

561 metadata=field_metadata, 

562 ) 

563 ) 

564 

565 schema = pa.schema(fields, metadata=md) 

566 

567 arrays = _numpy_style_arrays_to_arrow_arrays( 

568 astropy_table.dtype, 

569 len(astropy_table), 

570 astropy_table, 

571 schema, 

572 ) 

573 

574 arrow_table = pa.Table.from_arrays(arrays, schema=schema) 

575 

576 return arrow_table 

577 

578 

579def astropy_to_pandas(astropy_table: atable.Table, index: str | None = None) -> pd.DataFrame: 

580 """Convert an astropy table to a pandas dataframe via arrow. 

581 

582 By going via arrow we avoid pandas masked column bugs (e.g. 

583 https://github.com/pandas-dev/pandas/issues/58173) 

584 

585 Parameters 

586 ---------- 

587 astropy_table : `astropy.Table` 

588 Input astropy table. 

589 index : `str`, optional 

590 Name of column to set as index. 

591 

592 Returns 

593 ------- 

594 dataframe : `pandas.DataFrame` 

595 Output pandas dataframe. 

596 """ 

597 index_requested = False 

598 if (key := ASTROPY_PANDAS_INDEX_KEY) in astropy_table.meta: 

599 _index = astropy_table.meta[key] 

600 if _index not in astropy_table.columns: 

601 log.warning( 

602 "Index column ``%s`` not available for astropy table conversion to DataFrame", 

603 _index, 

604 ) 

605 _index = None 

606 else: 

607 index_requested = True 

608 _index = index 

609 

610 dataframe = arrow_to_pandas(astropy_to_arrow(astropy_table)) 

611 

612 # Set the index if we have a valid index name, and either the 

613 # index was requested in the call to the function or the dataframe 

614 # was not previously indexed with the call to arrow_to_pandas. 

615 if isinstance(_index, str) and (index_requested or dataframe.index.name is None): 

616 dataframe.set_index(_index, inplace=True) 

617 elif _index and index_requested: 

618 raise RuntimeError("index must be a string or None.") 

619 

620 return dataframe 

621 

622 

623def add_pandas_index_to_astropy(astropy_table: atable.Table, index: str) -> None: 

624 """Add special metadata to an astropy table to indicate a pandas index. 

625 

626 Parameters 

627 ---------- 

628 astropy_table : `astropy.table.Table` 

629 Input astropy table. 

630 index : `str` 

631 Name of column for pandas to set as index, if read as DataFrame. 

632 """ 

633 if index not in astropy_table.columns: 

634 raise ValueError("Column ``%s`` not in astropy table columns to use as pandas index.", index) 

635 astropy_table.meta[ASTROPY_PANDAS_INDEX_KEY] = index 

636 

637 

638def _astropy_to_numpy_dict(astropy_table: atable.Table) -> dict[str, np.ndarray]: 

639 """Convert an astropy table to an arrow table. 

640 

641 Parameters 

642 ---------- 

643 astropy_table : `astropy.Table` 

644 Input astropy table. 

645 

646 Returns 

647 ------- 

648 numpy_dict : `dict` [`str`, `numpy.ndarray`] 

649 Dict with keys as the column names, values as the arrays. 

650 """ 

651 return arrow_to_numpy_dict(astropy_to_arrow(astropy_table)) 

652 

653 

654def pandas_to_arrow(dataframe: pd.DataFrame, default_length: int = 10) -> pa.Table: 

655 """Convert a pandas dataframe to an arrow table. 

656 

657 Parameters 

658 ---------- 

659 dataframe : `pandas.DataFrame` 

660 Input pandas dataframe. 

661 default_length : `int`, optional 

662 Default string length when not in metadata or can be inferred 

663 from column. 

664 

665 Returns 

666 ------- 

667 arrow_table : `pyarrow.Table` 

668 Converted arrow table. 

669 """ 

670 try: 

671 arrow_table = pa.Table.from_pandas(dataframe) 

672 except pa.ArrowInvalid as e: 

673 msg = "; ".join(e.args) 

674 msg += "; This is usually because the column is mixed type or has uneven length rows." 

675 e.add_note(msg) 

676 raise 

677 

678 # Update the metadata 

679 md = arrow_table.schema.metadata 

680 

681 md[b"lsst::arrow::rowcount"] = str(arrow_table.num_rows) 

682 

683 # We loop through the arrow table columns because the datatypes have 

684 # been checked and converted from pandas objects. 

685 for name in arrow_table.column_names: 

686 if not name.startswith("__") and arrow_table[name].type == pa.string(): 

687 if len(arrow_table[name]) > 0: 

688 strlen = max(len(row.as_py()) for row in arrow_table[name] if row.is_valid) 

689 else: 

690 strlen = default_length 

691 md[f"lsst::arrow::len::{name}".encode()] = str(strlen) 

692 

693 arrow_table = arrow_table.replace_schema_metadata(md) 

694 

695 return arrow_table 

696 

697 

698def pandas_to_astropy(dataframe: pd.DataFrame) -> atable.Table: 

699 """Convert a pandas dataframe to an astropy table, preserving indexes. 

700 

701 Parameters 

702 ---------- 

703 dataframe : `pandas.DataFrame` 

704 Input pandas dataframe. 

705 

706 Returns 

707 ------- 

708 astropy_table : `astropy.table.Table` 

709 Converted astropy table. 

710 """ 

711 import pandas as pd 

712 

713 if isinstance(dataframe.columns, pd.MultiIndex): 

714 raise ValueError("Cannot convert a multi-index dataframe to an astropy table.") 

715 

716 return arrow_to_astropy(pandas_to_arrow(dataframe)) 

717 

718 

719def _pandas_to_numpy_dict(dataframe: pd.DataFrame) -> dict[str, np.ndarray]: 

720 """Convert a pandas dataframe to an dict of numpy arrays. 

721 

722 Parameters 

723 ---------- 

724 dataframe : `pandas.DataFrame` 

725 Input pandas dataframe. 

726 

727 Returns 

728 ------- 

729 numpy_dict : `dict` [`str`, `numpy.ndarray`] 

730 Dict with keys as the column names, values as the arrays. 

731 """ 

732 return arrow_to_numpy_dict(pandas_to_arrow(dataframe)) 

733 

734 

735def numpy_to_astropy(np_array: np.ndarray) -> atable.Table: 

736 """Convert a numpy table to an astropy table. 

737 

738 Parameters 

739 ---------- 

740 np_array : `numpy.ndarray` 

741 Input numpy array with multiple fields. 

742 

743 Returns 

744 ------- 

745 astropy_table : `astropy.table.Table` 

746 Converted astropy table. 

747 """ 

748 from astropy.table import Table 

749 

750 return Table(data=np_array, copy=False) 

751 

752 

753def arrow_schema_to_pandas_index(schema: pa.Schema) -> pd.Index | pd.MultiIndex: 

754 """Convert an arrow schema to a pandas index/multiindex. 

755 

756 Parameters 

757 ---------- 

758 schema : `pyarrow.Schema` 

759 Input pyarrow schema. 

760 

761 Returns 

762 ------- 

763 index : `pandas.Index` or `pandas.MultiIndex` 

764 Converted pandas index. 

765 """ 

766 import pandas as pd 

767 

768 if b"pandas" in schema.metadata: 

769 md = json.loads(schema.metadata[b"pandas"]) 

770 indexes = md["column_indexes"] 

771 len_indexes = len(indexes) 

772 else: 

773 len_indexes = 0 

774 

775 if len_indexes <= 1: 

776 return pd.Index(name for name in schema.names if not name.startswith("__")) 

777 else: 

778 raw_columns = _split_multi_index_column_names(len(indexes), schema.names) 

779 return pd.MultiIndex.from_tuples(raw_columns, names=[f["name"] for f in indexes]) 

780 

781 

782def arrow_schema_to_column_list(schema: pa.Schema) -> list[str]: 

783 """Convert an arrow schema to a list of string column names. 

784 

785 Parameters 

786 ---------- 

787 schema : `pyarrow.Schema` 

788 Input pyarrow schema. 

789 

790 Returns 

791 ------- 

792 column_list : `list` [`str`] 

793 Converted list of column names. 

794 """ 

795 return list(schema.names) 

796 

797 

798class DataFrameSchema: 

799 """Wrapper class for a schema for a pandas DataFrame. 

800 

801 Parameters 

802 ---------- 

803 dataframe : `pandas.DataFrame` 

804 Dataframe to turn into a schema. 

805 """ 

806 

807 def __init__(self, dataframe: pd.DataFrame) -> None: 

808 self._schema = dataframe.loc[[False] * len(dataframe)] 

809 

810 @classmethod 

811 def from_arrow(cls, schema: pa.Schema) -> DataFrameSchema: 

812 """Convert an arrow schema into a `DataFrameSchema`. 

813 

814 Parameters 

815 ---------- 

816 schema : `pyarrow.Schema` 

817 The pyarrow schema to convert. 

818 

819 Returns 

820 ------- 

821 dataframe_schema : `DataFrameSchema` 

822 Converted dataframe schema. 

823 """ 

824 empty_table = pa.Table.from_pylist([] * len(schema.names), schema=schema) 

825 

826 return cls(empty_table.to_pandas()) 

827 

828 def to_arrow_schema(self) -> pa.Schema: 

829 """Convert to an arrow schema. 

830 

831 Returns 

832 ------- 

833 arrow_schema : `pyarrow.Schema` 

834 Converted pyarrow schema. 

835 """ 

836 arrow_table = pa.Table.from_pandas(self._schema) 

837 

838 return arrow_table.schema 

839 

840 def to_arrow_numpy_schema(self) -> ArrowNumpySchema: 

841 """Convert to an `ArrowNumpySchema`. 

842 

843 Returns 

844 ------- 

845 arrow_numpy_schema : `ArrowNumpySchema` 

846 Converted arrow numpy schema. 

847 """ 

848 return ArrowNumpySchema.from_arrow(self.to_arrow_schema()) 

849 

850 def to_arrow_astropy_schema(self) -> ArrowAstropySchema: 

851 """Convert to an ArrowAstropySchema. 

852 

853 Returns 

854 ------- 

855 arrow_astropy_schema : `ArrowAstropySchema` 

856 Converted arrow astropy schema. 

857 """ 

858 return ArrowAstropySchema.from_arrow(self.to_arrow_schema()) 

859 

860 @property 

861 def schema(self) -> np.dtype: 

862 return self._schema 

863 

864 def __repr__(self) -> str: 

865 return repr(self._schema) 

866 

867 def __eq__(self, other: object) -> bool: 

868 if not isinstance(other, DataFrameSchema): 

869 return NotImplemented 

870 

871 return self._schema.equals(other._schema) 

872 

873 

874class ArrowAstropySchema: 

875 """Wrapper class for a schema for an astropy table. 

876 

877 Parameters 

878 ---------- 

879 astropy_table : `astropy.table.Table` 

880 Input astropy table. 

881 """ 

882 

883 def __init__(self, astropy_table: atable.Table) -> None: 

884 self._schema = astropy_table[:0] 

885 

886 @classmethod 

887 def from_arrow(cls, schema: pa.Schema) -> ArrowAstropySchema: 

888 """Convert an arrow schema into a ArrowAstropySchema. 

889 

890 Parameters 

891 ---------- 

892 schema : `pyarrow.Schema` 

893 Input pyarrow schema. 

894 

895 Returns 

896 ------- 

897 astropy_schema : `ArrowAstropySchema` 

898 Converted arrow astropy schema. 

899 """ 

900 import numpy as np 

901 from astropy.table import Table 

902 

903 dtype = _schema_to_dtype_list(schema) 

904 

905 data = np.zeros(0, dtype=dtype) 

906 astropy_table = Table(data=data) 

907 

908 _apply_astropy_metadata(astropy_table, schema) 

909 

910 return cls(astropy_table) 

911 

912 def to_arrow_schema(self) -> pa.Schema: 

913 """Convert to an arrow schema. 

914 

915 Returns 

916 ------- 

917 arrow_schema : `pyarrow.Schema` 

918 Converted pyarrow schema. 

919 """ 

920 return astropy_to_arrow(self._schema).schema 

921 

922 def to_dataframe_schema(self) -> DataFrameSchema: 

923 """Convert to a DataFrameSchema. 

924 

925 Returns 

926 ------- 

927 dataframe_schema : `DataFrameSchema` 

928 Converted dataframe schema. 

929 """ 

930 return DataFrameSchema.from_arrow(astropy_to_arrow(self._schema).schema) 

931 

932 def to_arrow_numpy_schema(self) -> ArrowNumpySchema: 

933 """Convert to an `ArrowNumpySchema`. 

934 

935 Returns 

936 ------- 

937 arrow_numpy_schema : `ArrowNumpySchema` 

938 Converted arrow numpy schema. 

939 """ 

940 return ArrowNumpySchema.from_arrow(astropy_to_arrow(self._schema).schema) 

941 

942 @property 

943 def schema(self) -> atable.Table: 

944 return self._schema 

945 

946 def __repr__(self) -> str: 

947 return repr(self._schema) 

948 

949 def __eq__(self, other: object) -> bool: 

950 if not isinstance(other, ArrowAstropySchema): 

951 return NotImplemented 

952 

953 # If this comparison passes then the two tables have the 

954 # same column names. 

955 if self._schema.dtype != other._schema.dtype: 

956 return False 

957 

958 for name in self._schema.columns: 

959 if not self._schema[name].unit == other._schema[name].unit: 

960 return False 

961 if not self._schema[name].description == other._schema[name].description: 

962 return False 

963 if not self._schema[name].format == other._schema[name].format: 

964 return False 

965 

966 return True 

967 

968 

969class ArrowNumpySchema: 

970 """Wrapper class for a schema for a numpy ndarray. 

971 

972 Parameters 

973 ---------- 

974 numpy_dtype : `numpy.dtype` 

975 Numpy dtype to convert. 

976 """ 

977 

978 def __init__(self, numpy_dtype: np.dtype) -> None: 

979 self._dtype = numpy_dtype 

980 

981 @classmethod 

982 def from_arrow(cls, schema: pa.Schema) -> ArrowNumpySchema: 

983 """Convert an arrow schema into an `ArrowNumpySchema`. 

984 

985 Parameters 

986 ---------- 

987 schema : `pyarrow.Schema` 

988 Pyarrow schema to convert. 

989 

990 Returns 

991 ------- 

992 numpy_schema : `ArrowNumpySchema` 

993 Converted arrow numpy schema. 

994 """ 

995 import numpy as np 

996 

997 dtype = _schema_to_dtype_list(schema) 

998 

999 return cls(np.dtype(dtype)) 

1000 

1001 def to_arrow_astropy_schema(self) -> ArrowAstropySchema: 

1002 """Convert to an `ArrowAstropySchema`. 

1003 

1004 Returns 

1005 ------- 

1006 astropy_schema : `ArrowAstropySchema` 

1007 Converted arrow astropy schema. 

1008 """ 

1009 import numpy as np 

1010 

1011 return ArrowAstropySchema.from_arrow(numpy_to_arrow(np.zeros(0, dtype=self._dtype)).schema) 

1012 

1013 def to_dataframe_schema(self) -> DataFrameSchema: 

1014 """Convert to a `DataFrameSchema`. 

1015 

1016 Returns 

1017 ------- 

1018 dataframe_schema : `DataFrameSchema` 

1019 Converted dataframe schema. 

1020 """ 

1021 import numpy as np 

1022 

1023 return DataFrameSchema.from_arrow(numpy_to_arrow(np.zeros(0, dtype=self._dtype)).schema) 

1024 

1025 def to_arrow_schema(self) -> pa.Schema: 

1026 """Convert to a `pyarrow.Schema`. 

1027 

1028 Returns 

1029 ------- 

1030 arrow_schema : `pyarrow.Schema` 

1031 Converted pyarrow schema. 

1032 """ 

1033 import numpy as np 

1034 

1035 return numpy_to_arrow(np.zeros(0, dtype=self._dtype)).schema 

1036 

1037 @property 

1038 def schema(self) -> np.dtype: 

1039 return self._dtype 

1040 

1041 def __repr__(self) -> str: 

1042 return repr(self._dtype) 

1043 

1044 def __eq__(self, other: object) -> bool: 

1045 if not isinstance(other, ArrowNumpySchema): 

1046 return NotImplemented 

1047 

1048 if not self._dtype == other._dtype: 

1049 return False 

1050 

1051 return True 

1052 

1053 

1054def _split_multi_index_column_names(n: int, names: Iterable[str]) -> list[Sequence[str]]: 

1055 """Split a string that represents a multi-index column. 

1056 

1057 PyArrow maps Pandas' multi-index column names (which are tuples in Python) 

1058 to flat strings on disk. This routine exists to reconstruct the original 

1059 tuple. 

1060 

1061 Parameters 

1062 ---------- 

1063 n : `int` 

1064 Number of levels in the `pandas.MultiIndex` that is being 

1065 reconstructed. 

1066 names : `~collections.abc.Iterable` [`str`] 

1067 Strings to be split. 

1068 

1069 Returns 

1070 ------- 

1071 column_names : `list` [`tuple` [`str`]] 

1072 A list of multi-index column name tuples. 

1073 """ 

1074 column_names: list[Sequence[str]] = [] 

1075 

1076 pattern = re.compile(r"\({}\)".format(", ".join(["'(.*)'"] * n))) 

1077 for name in names: 

1078 m = re.search(pattern, name) 

1079 if m is not None: 

1080 column_names.append(m.groups()) 

1081 

1082 return column_names 

1083 

1084 

1085def _standardize_multi_index_columns( 

1086 pd_index: pd.MultiIndex, 

1087 columns: Any, 

1088 stringify: bool = True, 

1089) -> list[str | Sequence[Any]]: 

1090 """Transform a dictionary/iterable index from a multi-index column list 

1091 into a string directly understandable by PyArrow. 

1092 

1093 Parameters 

1094 ---------- 

1095 pd_index : `pandas.MultiIndex` 

1096 Pandas multi-index. 

1097 columns : `list` [`tuple`] or `dict` [`str`, `str` or `list` [`str`]] 

1098 Columns to standardize. 

1099 stringify : `bool`, optional 

1100 Should the column names be stringified? 

1101 

1102 Returns 

1103 ------- 

1104 names : `list` [`str`] 

1105 Stringified representation of a multi-index column name. 

1106 """ 

1107 index_level_names = tuple(pd_index.names) 

1108 

1109 names: list[str | Sequence[Any]] = [] 

1110 

1111 if isinstance(columns, list): 

1112 for requested in columns: 

1113 if not isinstance(requested, tuple): 

1114 raise ValueError( 

1115 "Columns parameter for multi-index data frame must be a dictionary or list of tuples. " 

1116 f"Instead got a {get_full_type_name(requested)}." 

1117 ) 

1118 if stringify: 

1119 names.append(str(requested)) 

1120 else: 

1121 names.append(requested) 

1122 else: 

1123 if not isinstance(columns, collections.abc.Mapping): 

1124 raise ValueError( 

1125 "Columns parameter for multi-index data frame must be a dictionary or list of tuples. " 

1126 f"Instead got a {get_full_type_name(columns)}." 

1127 ) 

1128 if not set(index_level_names).issuperset(columns.keys()): 

1129 raise ValueError( 

1130 f"Cannot use dict with keys {set(columns.keys())} to select columns from {index_level_names}." 

1131 ) 

1132 factors = [ 

1133 ensure_iterable(columns.get(level, pd_index.levels[i])) 

1134 for i, level in enumerate(index_level_names) 

1135 ] 

1136 for requested in itertools.product(*factors): 

1137 for i, value in enumerate(requested): 

1138 if value not in pd_index.levels[i]: 

1139 raise ValueError(f"Unrecognized value {value!r} for index {index_level_names[i]!r}.") 

1140 if stringify: 

1141 names.append(str(requested)) 

1142 else: 

1143 names.append(requested) 

1144 

1145 return names 

1146 

1147 

1148def _apply_astropy_metadata(astropy_table: atable.Table, arrow_schema: pa.Schema) -> None: 

1149 """Apply any astropy metadata from the schema metadata. 

1150 

1151 Parameters 

1152 ---------- 

1153 astropy_table : `astropy.table.Table` 

1154 Table to apply metadata. 

1155 arrow_schema : `pyarrow.Schema` 

1156 Arrow schema with metadata. 

1157 """ 

1158 from astropy.table import meta 

1159 

1160 metadata = arrow_schema.metadata if arrow_schema.metadata is not None else {} 

1161 

1162 # Check if we have a special astropy metadata header yaml. 

1163 meta_yaml = metadata.get(b"table_meta_yaml", None) 

1164 if meta_yaml: 

1165 meta_yaml = meta_yaml.decode("UTF8").split("\n") 

1166 meta_hdr = meta.get_header_from_yaml(meta_yaml) 

1167 

1168 # Set description, format, unit, meta from the column 

1169 # metadata that was serialized with the table. 

1170 header_cols = {x["name"]: x for x in meta_hdr["datatype"]} 

1171 for col in astropy_table.columns.values(): 

1172 for attr in ("description", "format", "unit", "meta"): 

1173 if attr in header_cols[col.name]: 

1174 setattr(col, attr, header_cols[col.name][attr]) 

1175 

1176 if "meta" in meta_hdr: 

1177 astropy_table.meta.update(meta_hdr["meta"]) 

1178 else: 

1179 # If we don't have astropy header data, we may have arrow field 

1180 # metadata. 

1181 for name in arrow_schema.names: 

1182 field_metadata = arrow_schema.field(name).metadata 

1183 if field_metadata is None: 

1184 continue 

1185 if ( 

1186 b"description" in field_metadata 

1187 and (description := field_metadata[b"description"].decode("UTF-8")) != "" 

1188 ): 

1189 astropy_table[name].description = description 

1190 if b"unit" in field_metadata and (unit := field_metadata[b"unit"].decode("UTF-8")) != "": 

1191 astropy_table[name].unit = unit 

1192 

1193 

1194def _arrow_string_to_numpy_dtype( 

1195 schema: pa.Schema, name: str, numpy_column: np.ndarray | None = None, default_length: int = 10 

1196) -> str: 

1197 """Get the numpy dtype string associated with an arrow column. 

1198 

1199 Parameters 

1200 ---------- 

1201 schema : `pyarrow.Schema` 

1202 Arrow table schema. 

1203 name : `str` 

1204 Column name. 

1205 numpy_column : `numpy.ndarray`, optional 

1206 Column to determine numpy string dtype. 

1207 default_length : `int`, optional 

1208 Default string length when not in metadata or can be inferred 

1209 from column. 

1210 

1211 Returns 

1212 ------- 

1213 dtype_str : `str` 

1214 Numpy dtype string. 

1215 """ 

1216 # Special-case for string and binary columns 

1217 md_name = f"lsst::arrow::len::{name}" 

1218 strlen = default_length 

1219 metadata = schema.metadata if schema.metadata is not None else {} 

1220 if (encoded := md_name.encode("UTF-8")) in metadata: 

1221 # String/bytes length from header. 

1222 strlen = int(schema.metadata[encoded]) 

1223 elif numpy_column is not None and len(numpy_column) > 0: 

1224 lengths = [len(row) for row in numpy_column if row] 

1225 strlen = max(lengths) if lengths else 0 

1226 

1227 dtype = f"U{strlen}" if schema.field(name).type == pa.string() else f"|S{strlen}" 

1228 

1229 return dtype 

1230 

1231 

1232def _append_numpy_string_metadata(metadata: dict[bytes, str], name: str, dtype: np.dtype) -> None: 

1233 """Append numpy string length keys to arrow metadata. 

1234 

1235 All column types are handled, but the metadata is only modified for 

1236 string and byte columns. 

1237 

1238 Parameters 

1239 ---------- 

1240 metadata : `dict` [`bytes`, `str`] 

1241 Metadata dictionary; modified in place. 

1242 name : `str` 

1243 Column name. 

1244 dtype : `np.dtype` 

1245 Numpy dtype. 

1246 """ 

1247 import numpy as np 

1248 

1249 if dtype.type is np.str_: 

1250 metadata[f"lsst::arrow::len::{name}".encode()] = str(dtype.itemsize // 4) 

1251 metadata[f"table::len::{name}".encode()] = str(dtype.itemsize // 4) 

1252 elif dtype.type is np.bytes_: 

1253 metadata[f"lsst::arrow::len::{name}".encode()] = str(dtype.itemsize) 

1254 metadata[f"table::len::{name}".encode()] = str(dtype.itemsize) 

1255 

1256 

1257def _append_numpy_multidim_metadata(metadata: dict[bytes, str], name: str, dtype: np.dtype) -> None: 

1258 """Append numpy multi-dimensional shapes to arrow metadata. 

1259 

1260 All column types are handled, but the metadata is only modified for 

1261 multi-dimensional columns. 

1262 

1263 Parameters 

1264 ---------- 

1265 metadata : `dict` [`bytes`, `str`] 

1266 Metadata dictionary; modified in place. 

1267 name : `str` 

1268 Column name. 

1269 dtype : `np.dtype` 

1270 Numpy dtype. 

1271 """ 

1272 if len(dtype.shape) > 1: 

1273 metadata[f"lsst::arrow::shape::{name}".encode()] = str(dtype.shape) 

1274 

1275 

1276def _multidim_shape_from_metadata(metadata: dict[bytes, bytes], list_size: int, name: str) -> tuple[int, ...]: 

1277 """Retrieve the shape from the metadata, if available. 

1278 

1279 Parameters 

1280 ---------- 

1281 metadata : `dict` [`bytes`, `bytes`] 

1282 Metadata dictionary. 

1283 list_size : `int` 

1284 Size of the list datatype. 

1285 name : `str` 

1286 Column name. 

1287 

1288 Returns 

1289 ------- 

1290 shape : `tuple` [`int`] 

1291 Shape associated with the column. 

1292 

1293 Raises 

1294 ------ 

1295 RuntimeError 

1296 Raised if metadata is found but has incorrect format. 

1297 """ 

1298 md_name = f"lsst::arrow::shape::{name}" 

1299 if (encoded := md_name.encode("UTF-8")) in metadata: 

1300 groups = re.search(r"\((.*)\)", metadata[encoded].decode("UTF-8")) 

1301 if groups is None: 

1302 raise RuntimeError("Illegal value found in metadata.") 

1303 shape = tuple(int(x) for x in groups[1].split(",") if x != "") 

1304 else: 

1305 shape = (list_size,) 

1306 

1307 return shape 

1308 

1309 

1310def _schema_to_dtype_list(schema: pa.Schema) -> list[tuple[str, tuple[Any] | str]]: 

1311 """Convert a pyarrow schema to a numpy dtype. 

1312 

1313 Parameters 

1314 ---------- 

1315 schema : `pyarrow.Schema` 

1316 Input pyarrow schema. 

1317 

1318 Returns 

1319 ------- 

1320 dtype_list: `list` [`tuple`] 

1321 A list with name, type pairs. 

1322 """ 

1323 metadata = schema.metadata if schema.metadata is not None else {} 

1324 

1325 dtype: list[Any] = [] 

1326 for name in schema.names: 

1327 t = schema.field(name).type 

1328 if isinstance(t, pa.FixedSizeListType): 

1329 shape = _multidim_shape_from_metadata(metadata, t.list_size, name) 

1330 dtype.append((name, (t.value_type.to_pandas_dtype(), shape))) 

1331 elif t not in (pa.string(), pa.binary()): 

1332 dtype.append((name, t.to_pandas_dtype())) 

1333 else: 

1334 dtype.append((name, _arrow_string_to_numpy_dtype(schema, name))) 

1335 

1336 return dtype 

1337 

1338 

1339def _numpy_dtype_to_arrow_types(dtype: np.dtype) -> list[Any]: 

1340 """Convert a numpy dtype to a list of arrow types. 

1341 

1342 Parameters 

1343 ---------- 

1344 dtype : `numpy.dtype` 

1345 Numpy dtype to convert. 

1346 

1347 Returns 

1348 ------- 

1349 type_list : `list` [`object`] 

1350 Converted list of arrow types. 

1351 """ 

1352 from math import prod 

1353 

1354 import numpy as np 

1355 

1356 type_list: list[Any] = [] 

1357 if dtype.names is None: 

1358 return type_list 

1359 

1360 for name in dtype.names: 

1361 dt = dtype[name] 

1362 arrow_type: Any 

1363 if len(dt.shape) > 0: 

1364 arrow_type = pa.list_( 

1365 pa.from_numpy_dtype(cast(tuple[np.dtype, tuple[int, ...]], dt.subdtype)[0].type), 

1366 prod(dt.shape), 

1367 ) 

1368 elif dt.type == np.datetime64: 

1369 time_unit = "ns" if "ns" in dt.str else "us" 

1370 # The pa.timestamp() is the correct datatype to round-trip 

1371 # a numpy datetime64[ns] or datetime[us] array. 

1372 arrow_type = pa.timestamp(time_unit) 

1373 else: 

1374 try: 

1375 arrow_type = pa.from_numpy_dtype(dt.type) 

1376 except pa.ArrowNotImplementedError as e: 

1377 msg = f"Could not serialize column {name} (type {str(dt)}) to Parquet." 

1378 if dt == np.dtype("O"): 

1379 msg += " This is usually because the column is mixed type or has uneven length rows." 

1380 e.add_note(msg) 

1381 raise 

1382 type_list.append((name, arrow_type)) 

1383 

1384 return type_list 

1385 

1386 

1387def _numpy_dict_to_dtype(numpy_dict: dict[str, np.ndarray]) -> tuple[np.dtype, int]: 

1388 """Extract equivalent table dtype from dict of numpy arrays. 

1389 

1390 Parameters 

1391 ---------- 

1392 numpy_dict : `dict` [`str`, `numpy.ndarray`] 

1393 Dict with keys as the column names, values as the arrays. 

1394 

1395 Returns 

1396 ------- 

1397 dtype : `numpy.dtype` 

1398 dtype of equivalent table. 

1399 rowcount : `int` 

1400 Number of rows in the table. 

1401 

1402 Raises 

1403 ------ 

1404 ValueError if columns in numpy_dict have unequal numbers of rows. 

1405 """ 

1406 import numpy as np 

1407 

1408 dtype_list: list[tuple] = [] 

1409 rowcount = 0 

1410 for name, col in numpy_dict.items(): 

1411 if rowcount == 0: 

1412 rowcount = len(col) 

1413 if len(col) != rowcount: 

1414 raise ValueError(f"Column {name} has a different number of rows.") 

1415 if len(col.shape) == 1: 

1416 dtype_list.append((name, col.dtype)) 

1417 else: 

1418 dtype_list.append((name, (col.dtype, col.shape[1:]))) 

1419 dtype = np.dtype(dtype_list) 

1420 

1421 return (dtype, rowcount) 

1422 

1423 

1424def _numpy_style_arrays_to_arrow_arrays( 

1425 dtype: np.dtype, 

1426 rowcount: int, 

1427 np_style_arrays: dict[str, np.ndarray] | np.ndarray | atable.Table, 

1428 schema: pa.Schema, 

1429) -> list[pa.Array]: 

1430 """Convert numpy-style arrays to arrow arrays. 

1431 

1432 Parameters 

1433 ---------- 

1434 dtype : `numpy.dtype` 

1435 Numpy dtype of input table/arrays. 

1436 rowcount : `int` 

1437 Number of rows in input table/arrays. 

1438 np_style_arrays : `dict` [`str`, `np.ndarray`] or `np.ndarray` 

1439 or `astropy.table.Table` 

1440 Arrays to convert to arrow. 

1441 schema : `pyarrow.Schema` 

1442 Schema of arrow table. 

1443 

1444 Returns 

1445 ------- 

1446 arrow_arrays : `list` [`pyarrow.Array`] 

1447 List of converted pyarrow arrays. 

1448 """ 

1449 import numpy as np 

1450 

1451 arrow_arrays: list[pa.Array] = [] 

1452 if dtype.names is None: 

1453 return arrow_arrays 

1454 

1455 for name in dtype.names: 

1456 dt = dtype[name] 

1457 val: Any 

1458 if len(dt.shape) > 0: 

1459 if rowcount > 0: 

1460 val = np.split(np_style_arrays[name].ravel(), rowcount) 

1461 else: 

1462 val = [] 

1463 else: 

1464 val = np_style_arrays[name] 

1465 

1466 try: 

1467 arrow_arrays.append(pa.array(val, type=schema.field(name).type)) 

1468 except pa.ArrowNotImplementedError as err: 

1469 # Check if val is big-endian. 

1470 if (np.little_endian and val.dtype.byteorder == ">") or ( 

1471 not np.little_endian and val.dtype.byteorder == "=" 

1472 ): 

1473 # We need to convert the array to little-endian. 

1474 val2 = val.byteswap() 

1475 val2.dtype = val2.dtype.newbyteorder("<") 

1476 arrow_arrays.append(pa.array(val2, type=schema.field(name).type)) 

1477 else: 

1478 # This failed for some other reason so raise the exception. 

1479 raise err 

1480 

1481 return arrow_arrays 

1482 

1483 

1484def compute_row_group_size(schema: pa.Schema, target_size: int = TARGET_ROW_GROUP_BYTES) -> int: 

1485 """Compute approximate row group size for a given arrow schema. 

1486 

1487 Given a schema, this routine will compute the number of rows in a row group 

1488 that targets the persisted size on disk (or smaller). The exact size on 

1489 disk depends on the compression settings and ratios; typical binary data 

1490 tables will have around 15-20% compression with the pyarrow default 

1491 ``snappy`` compression algorithm. 

1492 

1493 Parameters 

1494 ---------- 

1495 schema : `pyarrow.Schema` 

1496 Arrow table schema. 

1497 target_size : `int`, optional 

1498 The target size (in bytes). 

1499 

1500 Returns 

1501 ------- 

1502 row_group_size : `int` 

1503 Number of rows per row group to hit the target size. 

1504 """ 

1505 bit_width = 0 

1506 

1507 metadata = schema.metadata if schema.metadata is not None else {} 

1508 

1509 for name in schema.names: 

1510 t = schema.field(name).type 

1511 

1512 if t in (pa.string(), pa.binary()): 

1513 md_name = f"lsst::arrow::len::{name}" 

1514 

1515 if (encoded := md_name.encode("UTF-8")) in metadata: 

1516 # String/bytes length from header. 

1517 strlen = int(schema.metadata[encoded]) 

1518 else: 

1519 # We don't know the string width, so guess something. 

1520 strlen = 10 

1521 

1522 # Assuming UTF-8 encoding, and very few wide characters. 

1523 t_width = 8 * strlen 

1524 elif isinstance(t, pa.FixedSizeListType): 

1525 if t.value_type == pa.null(): 

1526 t_width = 0 

1527 else: 

1528 t_width = t.list_size * t.value_type.bit_width 

1529 elif t == pa.null(): 

1530 t_width = 0 

1531 elif isinstance(t, pa.ListType): 

1532 if t.value_type == pa.null(): 

1533 t_width = 0 

1534 else: 

1535 # This is a variable length list, just choose 

1536 # something arbitrary. 

1537 t_width = 10 * t.value_type.bit_width 

1538 else: 

1539 t_width = t.bit_width 

1540 

1541 bit_width += t_width 

1542 

1543 # Insist it is at least 1 byte wide to avoid any divide-by-zero errors. 

1544 if bit_width < 8: 

1545 bit_width = 8 

1546 

1547 byte_width = bit_width // 8 

1548 

1549 return target_size // byte_width