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

1# This file is part of daf_butler.

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.

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

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

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

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

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

14#

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

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

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

18# (at your option) any later version.

19#

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

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

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

23# GNU General Public License for more details.

24#

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

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

28from __future__ import annotations

30__all__ = (

31 "ArrowAstropySchema",

32 "ArrowNumpySchema",

33 "DataFrameSchema",

34 "ParquetFormatter",

35 "add_pandas_index_to_astropy",

36 "arrow_schema_to_pandas_index",

37 "arrow_to_astropy",

38 "arrow_to_numpy",

39 "arrow_to_numpy_dict",

40 "arrow_to_pandas",

41 "astropy_to_arrow",

42 "astropy_to_pandas",

43 "compute_row_group_size",

44 "numpy_dict_to_arrow",

45 "numpy_to_arrow",

46 "numpy_to_astropy",

47 "pandas_to_arrow",

48 "pandas_to_astropy",

49 "pandas_to_numpy",

50)

52import collections.abc

53import contextlib

54import itertools

55import json

56import logging

57import re

58from collections.abc import Generator, Iterable, Sequence

59from fnmatch import fnmatchcase

60from typing import IO, TYPE_CHECKING, Any, cast

62import pyarrow as pa

63import pyarrow.parquet as pq

65from lsst.daf.butler import DatasetProvenance, FormatterV2

66from lsst.daf.butler.delegates.arrowtable import _add_arrow_provenance, _checkArrowCompatibleType

67from lsst.resources import ResourcePath

68from lsst.utils.introspection import get_full_type_name

69from lsst.utils.iteration import ensure_iterable

71log = logging.getLogger(__name__)

73if TYPE_CHECKING:

74 import astropy.table as atable

75 import numpy as np

76 import pandas as pd

78 try:

79 import fsspec

80 from fsspec.spec import AbstractFileSystem

81 except ImportError:

82 fsspec = None

83 AbstractFileSystem = type

85TARGET_ROW_GROUP_BYTES = 1_000_000_000

86ASTROPY_PANDAS_INDEX_KEY = "lsst::arrow::astropy_pandas_index"

89@contextlib.contextmanager

90def generic_open(path: str, fs: AbstractFileSystem | None) -> Generator[IO]:

91 if fs is None:

92 with open(path, "rb") as fh:

93 yield fh

94 else:

95 with fs.open(path) as fh:

96 yield fh

99class ParquetFormatter(FormatterV2):

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

101 Parquet files.

102 """

103

104 default_extension = ".parq"

105 can_read_from_uri = True

106 can_read_from_local_file = True

107

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

109 # Docstring inherited.

110 return _checkArrowCompatibleType(in_memory_dataset) is not None

111

112 def read_from_uri(self, uri: ResourcePath, component: str | None = None, expected_size: int = -1) -> Any:

113 # Docstring inherited from Formatter.read.

114 try:

115 fs, path = uri.to_fsspec()

116 except ImportError:

117 log.debug("fsspec not available; falling back to local file access.")

118 # This signals to the formatter to use the read_from_local_file

119 # code path.

120 return NotImplemented

121

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

123

124 def read_from_local_file(self, path: str, component: str | None = None, expected_size: int = -1) -> Any:

125 # Docstring inherited from Formatter.read.

126 return self._read_parquet(path=path, component=component, expected_size=expected_size)

127

128 def _read_parquet(

129 self,

130 path: str,

131 fs: AbstractFileSystem | None = None,

132 component: str | None = None,

133 expected_size: int = -1,

134 ) -> Any:

135 with generic_open(path, fs) as handle:

136 schema = pq.read_schema(handle)

137

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

139

140 if component in ("columns", "schema") or self.file_descriptor.readStorageClass.name in schema_names:

141 # The schema will be translated to column format

142 # depending on the input type.

143 return schema

144 elif component == "rowcount":

145 # Get the rowcount from the metadata if possible, otherwise count.

146 if b"lsst::arrow::rowcount" in schema.metadata:

147 return int(schema.metadata[b"lsst::arrow::rowcount"])

148

149 with generic_open(path, fs) as handle:

150 temp_table = pq.read_table(

151 handle,

152 columns=[schema.names[0]],

153 use_threads=False,

154 use_pandas_metadata=False,

155 )

156

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

158

159 par_columns = None

160 strip_astropy_meta_yaml = True

161 if self.file_descriptor.parameters:

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

163 if par_columns:

164 has_pandas_multi_index = False

165 index_columns = []

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

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

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

169 has_pandas_multi_index = True

170

171 index_columns = _get_pandas_index_columns(md)

172

173 if not has_pandas_multi_index:

174 # Ensure uniqueness, keeping order.

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

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

177

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

179 # uniqueness and ordering.

180 par_columns = {}

181 for par_column in par_columns_in:

182 found = False

183 for file_column in file_columns:

184 if fnmatchcase(file_column, par_column):

185 found = True

186 par_columns[file_column] = True

187 if not found:

188 if par_column not in index_columns: 188 ↛ 181line 188 didn't jump to line 181 because the condition on line 188 was always true

189 # Note that the pandas index column will always

190 # be loaded.

191 raise ValueError(

192 f"Column {par_column} specified in parameters not "

193 "available in parquet file."

194 )

195 par_columns = list(par_columns.keys())

196 else:

197 par_columns = _standardize_multi_index_columns(

198 arrow_schema_to_pandas_index(schema),

199 par_columns,

200 )

201

202 strip_astropy_meta_yaml = self.file_descriptor.parameters.pop(

203 "strip_astropy_meta_yaml",

204 True,

205 )

206

207 if len(self.file_descriptor.parameters): 207 ↛ 208line 207 didn't jump to line 208 because the condition on line 207 was never true

208 raise ValueError(

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

210 )

211

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

213 with generic_open(path, fs) as handle:

214 arrow_table = pq.read_table(

215 handle,

216 columns=par_columns,

217 use_threads=False,

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

219 )

220

221 if strip_astropy_meta_yaml:

222 metadata = arrow_table.schema.metadata

223 # Only strip if (a) we have metadata; (b) it contains

224 # ``table_meta_yaml``; (c) it contains ``lsst::arrow::rowcount``

225 # to avoid stripping data from pure astropy tables (not written

226 # by the butler).

227 if metadata and metadata.pop(b"table_meta_yaml", None) and b"lsst::arrow::rowcount" in metadata:

228 arrow_table = arrow_table.replace_schema_metadata(metadata)

229

230 return arrow_table

231

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

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

234

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

236 """Serialize the in memory dataset to a local parquet file.

237

238 Parameters

239 ----------

240 in_memory_dataset : `typing.Any`

241 The Python object to serialize.

242 uri : `lsst.resources.ResourcePath`

243 The location to write the local file.

244 """

245 if isinstance(in_memory_dataset, pa.Schema):

246 pq.write_metadata(in_memory_dataset, uri.ospath)

247 return

248

249 type_string = _checkArrowCompatibleType(in_memory_dataset)

250

251 if type_string is None:

252 raise ValueError(

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

254 "inMemoryDataset for ParquetFormatter."

255 )

256

257 if type_string == "arrow":

258 arrow_table = in_memory_dataset

259 elif type_string == "astropy":

260 arrow_table = astropy_to_arrow(in_memory_dataset)

261 elif type_string == "numpy":

262 arrow_table = numpy_to_arrow(in_memory_dataset)

263 elif type_string == "numpydict":

264 arrow_table = numpy_dict_to_arrow(in_memory_dataset)

265 else:

266 arrow_table = pandas_to_arrow(in_memory_dataset)

267

268 row_group_size = compute_row_group_size(arrow_table.schema)

269

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

271

272

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

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

275

276 Parameters

277 ----------

278 arrow_table : `pyarrow.Table`

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

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

281 ``DataFrame``.

282

283 Returns

284 -------

285 dataframe : `pandas.DataFrame`

286 Converted pandas dataframe.

287 """

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

289

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

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

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

293 if pandas_index in arrow_table.schema.names:

294 dataframe.set_index(pandas_index, inplace=True)

295 else:

296 log.warning(

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

298 pandas_index,

299 )

300

301 return dataframe

302

303

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

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

306

307 Parameters

308 ----------

309 arrow_table : `pyarrow.Table`

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

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

312 of the ``astropy.table.Table``.

313

314 Returns

315 -------

316 table : `astropy.table.Table`

317 Converted astropy table.

318 """

319 from astropy.table import Table

320

321 astropy_table = Table(arrow_to_numpy_dict(arrow_table))

322

323 _apply_astropy_metadata(astropy_table, arrow_table.schema)

324

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

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

327 astropy_table.meta.pop(key)

328

329 return astropy_table

330

331

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

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

334

335 Parameters

336 ----------

337 arrow_table : `pyarrow.Table`

338 Input arrow table.

339

340 Returns

341 -------

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

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

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

345 in the table are null.

346 """

347 import numpy as np

348

349 numpy_dict = arrow_to_numpy_dict(arrow_table)

350

351 has_mask = False

352 dtype: list[tuple] = []

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

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

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

356 else:

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

358

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

360 has_mask = True

361

362 array: Any

363

364 if has_mask:

365 import numpy.ma.mrecords as mrecords

366

367 array = mrecords.fromarrays(list(numpy_dict.values()), dtype=dtype)

368 else:

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

370 return array

371

372

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

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

375

376 Parameters

377 ----------

378 arrow_table : `pyarrow.Table`

379 Input arrow table.

380

381 Returns

382 -------

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

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

385 """

386 import numpy as np

387

388 schema = arrow_table.schema

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

390

391 numpy_dict = {}

392

393 for name in schema.names:

394 t = schema.field(name).type

395

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

397 # Regular non-masked column

398 col = arrow_table[name].to_numpy()

399 else:

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

401 # values with an appropriately typed value before conversion.

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

403 null_value: Any

404 match t:

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

406 null_value = np.nan

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

408 null_value = -1

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

410 null_value = True

411 case t if _is_string(t) or _is_binary(t):

412 null_value = ""

413 case _:

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

415 null_value = 0

416

417 col = np.ma.MaskedArray(

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

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

420 fill_value=null_value,

421 )

422

423 if _is_string(t) or _is_binary(t):

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

425 elif isinstance(t, pa.FixedSizeListType):

426 if len(col) > 0:

427 col = np.stack(col)

428 else:

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

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

431

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

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

434

435 numpy_dict[name] = col

436

437 return numpy_dict

438

439

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

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

442

443 Parameters

444 ----------

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

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

447

448 Returns

449 -------

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

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

452 """

453 return arrow_to_numpy(numpy_dict_to_arrow(numpy_dict))

454

455

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

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

458

459 Parameters

460 ----------

461 np_array : `numpy.ndarray`

462 Input numpy array with multiple fields.

463

464 Returns

465 -------

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

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

468 """

469 return arrow_to_numpy_dict(numpy_to_arrow(np_array))

470

471

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

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

474

475 Parameters

476 ----------

477 np_array : `numpy.ndarray`

478 Input numpy array with multiple fields.

479

480 Returns

481 -------

482 arrow_table : `pyarrow.Table`

483 Converted arrow table.

484 """

485 type_list = _numpy_dtype_to_arrow_types(np_array.dtype)

486

487 md = {}

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

489

490 names = np_array.dtype.names

491 if names is None: 491 ↛ 492line 491 didn't jump to line 492 because the condition on line 491 was never true

492 names = ()

493

494 for name in names:

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

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

497

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

499

500 arrays = _numpy_style_arrays_to_arrow_arrays(

501 np_array.dtype,

502 len(np_array),

503 np_array,

504 schema,

505 )

506

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

508

509 return arrow_table

510

511

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

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

514

515 Parameters

516 ----------

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

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

519

520 Returns

521 -------

522 arrow_table : `pyarrow.Table`

523 Converted arrow table.

524

525 Raises

526 ------

527 ValueError

528 Raised if columns in ``numpy_dict`` have unequal numbers of

529 rows.

530 """

531 dtype, rowcount = _numpy_dict_to_dtype(numpy_dict)

532 type_list = _numpy_dtype_to_arrow_types(dtype)

533

534 md = {}

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

536

537 if dtype.names is not None: 537 ↛ 542line 537 didn't jump to line 542 because the condition on line 537 was always true

538 for name in dtype.names:

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

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

541

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

543

544 arrays = _numpy_style_arrays_to_arrow_arrays(

545 dtype,

546 rowcount,

547 numpy_dict,

548 schema,

549 )

550

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

552

553 return arrow_table

554

555

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

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

558

559 Parameters

560 ----------

561 astropy_table : `astropy.table.Table`

562 Input astropy table.

563

564 Returns

565 -------

566 arrow_table : `pyarrow.Table`

567 Converted arrow table.

568 """

569 from astropy.table import meta

570

571 type_list = _numpy_dtype_to_arrow_types(astropy_table.dtype)

572

573 md = {}

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

575

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

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

578

579 for name in astropy_table.dtype.names:

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

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

582

583 meta_yaml = meta.get_yaml_from_table(astropy_table)

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

585 md[b"table_meta_yaml"] = meta_yaml_str

586

587 # Convert type list to fields with metadata.

588 fields = []

589 for name, pa_type in type_list:

590 field_metadata = {}

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

592 field_metadata["description"] = description

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

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

595 fields.append(

596 pa.field(

597 name,

598 pa_type,

599 metadata=field_metadata,

600 )

601 )

602

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

604

605 arrays = _numpy_style_arrays_to_arrow_arrays(

606 astropy_table.dtype,

607 len(astropy_table),

608 astropy_table,

609 schema,

610 )

611

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

613

614 return arrow_table

615

616

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

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

619

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

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

622

623 Parameters

624 ----------

625 astropy_table : `astropy.table.Table`

626 Input astropy table.

627 index : `str`, optional

628 Name of column to set as index.

629

630 Returns

631 -------

632 dataframe : `pandas.DataFrame`

633 Output pandas dataframe.

634 """

635 index_requested = False

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

637 _index = astropy_table.meta[key]

638 if _index not in astropy_table.columns:

639 log.warning(

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

641 _index,

642 )

643 _index = None

644 else:

645 index_requested = True

646 _index = index

647

648 dataframe = arrow_to_pandas(astropy_to_arrow(astropy_table))

649

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

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

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

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

654 dataframe.set_index(_index, inplace=True)

655 elif _index and index_requested: 655 ↛ 656line 655 didn't jump to line 656 because the condition on line 655 was never true

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

657

658 return dataframe

659

660

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

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

663

664 Parameters

665 ----------

666 astropy_table : `astropy.table.Table`

667 Input astropy table.

668 index : `str`

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

670 """

671 if index not in astropy_table.columns:

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

673 astropy_table.meta[ASTROPY_PANDAS_INDEX_KEY] = index

674

675

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

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

678

679 Parameters

680 ----------

681 astropy_table : `astropy.table.Table`

682 Input astropy table.

683

684 Returns

685 -------

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

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

688 """

689 return arrow_to_numpy_dict(astropy_to_arrow(astropy_table))

690

691

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

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

694

695 Parameters

696 ----------

697 dataframe : `pandas.DataFrame`

698 Input pandas dataframe.

699 default_length : `int`, optional

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

701 from column.

702

703 Returns

704 -------

705 arrow_table : `pyarrow.Table`

706 Converted arrow table.

707 """

708 try:

709 arrow_table = pa.Table.from_pandas(dataframe, preserve_index=True)

710 except pa.ArrowInvalid as e:

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

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

713 e.add_note(msg)

714 raise

715

716 # Update the metadata

717 md = arrow_table.schema.metadata

718

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

720

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

722 # been checked and converted from pandas objects.

723 for name in arrow_table.column_names:

724 if not name.startswith("__") and _is_string(arrow_table[name].type):

725 col = arrow_table[name]

726 if len(col) > 0 and (col.length() > col.null_count): 726 ↛ 729line 726 didn't jump to line 729 because the condition on line 726 was always true

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

728 else:

729 strlen = default_length

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

731

732 arrow_table = arrow_table.replace_schema_metadata(md)

733

734 return arrow_table

735

736

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

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

739

740 Parameters

741 ----------

742 dataframe : `pandas.DataFrame`

743 Input pandas dataframe.

744

745 Returns

746 -------

747 astropy_table : `astropy.table.Table`

748 Converted astropy table.

749 """

750 import pandas as pd

751

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

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

754

755 return arrow_to_astropy(pandas_to_arrow(dataframe))

756

757

758def pandas_to_numpy(dataframe: pd.DataFrame) -> np.ndarray | np.ma.MaskedArray:

759 """Convert a pandas dataframe to a numpy recarray.

760

761 Parameters

762 ----------

763 dataframe : `pandas.DataFrame`

764 Input pandas dataframe.

765

766 Returns

767 -------

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

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

770 as the input dataframe. Will be masked records if any values

771 in the table are null.

772 """

773 # This conversion ensures strings are handled properly.

774 return arrow_to_numpy(pandas_to_arrow(dataframe))

775

776

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

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

779

780 Parameters

781 ----------

782 dataframe : `pandas.DataFrame`

783 Input pandas dataframe.

784

785 Returns

786 -------

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

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

789 """

790 return arrow_to_numpy_dict(pandas_to_arrow(dataframe))

791

792

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

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

795

796 Parameters

797 ----------

798 np_array : `numpy.ndarray`

799 Input numpy array with multiple fields.

800

801 Returns

802 -------

803 astropy_table : `astropy.table.Table`

804 Converted astropy table.

805 """

806 from astropy.table import Table

807

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

809

810

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

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

813

814 Parameters

815 ----------

816 schema : `pyarrow.Schema`

817 Input pyarrow schema.

818

819 Returns

820 -------

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

822 Converted pandas index.

823 """

824 import pandas as pd

825

826 index_columns = []

827 if b"pandas" in schema.metadata:

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

829 indexes = md["column_indexes"]

830 len_indexes = len(indexes)

831

832 if len_indexes > 0: 832 ↛ 837line 832 didn't jump to line 837 because the condition on line 832 was always true

833 index_columns = _get_pandas_index_columns(md)

834 else:

835 len_indexes = 0

836

837 if len_indexes <= 1:

838 columns = {name for name in schema.names if not name.startswith("__")}

839 columns = columns.union(set(index_columns))

840 return pd.Index(columns)

841 else:

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

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

844

845

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

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

848

849 Parameters

850 ----------

851 schema : `pyarrow.Schema`

852 Input pyarrow schema.

853

854 Returns

855 -------

856 column_list : `list` [`str`]

857 Converted list of column names.

858 """

859 return list(schema.names)

860

861

862class DataFrameSchema:

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

864

865 Parameters

866 ----------

867 dataframe : `pandas.DataFrame`

868 Dataframe to turn into a schema.

869 """

870

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

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

873

874 @classmethod

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

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

877

878 Parameters

879 ----------

880 schema : `pyarrow.Schema`

881 The pyarrow schema to convert.

882

883 Returns

884 -------

885 dataframe_schema : `DataFrameSchema`

886 Converted dataframe schema.

887 """

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

889

890 return cls(empty_table.to_pandas())

891

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

893 """Convert to an arrow schema.

894

895 Returns

896 -------

897 arrow_schema : `pyarrow.Schema`

898 Converted pyarrow schema.

899 """

900 arrow_table = pa.Table.from_pandas(self._schema, preserve_index=True)

901

902 return arrow_table.schema

903

904 def to_arrow_numpy_schema(self) -> ArrowNumpySchema:

905 """Convert to an `ArrowNumpySchema`.

906

907 Returns

908 -------

909 arrow_numpy_schema : `ArrowNumpySchema`

910 Converted arrow numpy schema.

911 """

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

913

914 def to_arrow_astropy_schema(self) -> ArrowAstropySchema:

915 """Convert to an ArrowAstropySchema.

916

917 Returns

918 -------

919 arrow_astropy_schema : `ArrowAstropySchema`

920 Converted arrow astropy schema.

921 """

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

923

924 @property

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

926 return self._schema

927

928 def __repr__(self) -> str:

929 return repr(self._schema)

930

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

932 if not isinstance(other, DataFrameSchema):

933 return NotImplemented

934

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

936

937

938class ArrowAstropySchema:

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

940

941 Parameters

942 ----------

943 astropy_table : `astropy.table.Table`

944 Input astropy table.

945 """

946

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

948 self._schema = astropy_table[:0]

949

950 @classmethod

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

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

953

954 Parameters

955 ----------

956 schema : `pyarrow.Schema`

957 Input pyarrow schema.

958

959 Returns

960 -------

961 astropy_schema : `ArrowAstropySchema`

962 Converted arrow astropy schema.

963 """

964 import numpy as np

965 from astropy.table import Table

966

967 dtype = _schema_to_dtype_list(schema)

968

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

970 astropy_table = Table(data=data)

971

972 _apply_astropy_metadata(astropy_table, schema)

973

974 return cls(astropy_table)

975

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

977 """Convert to an arrow schema.

978

979 Returns

980 -------

981 arrow_schema : `pyarrow.Schema`

982 Converted pyarrow schema.

983 """

984 return astropy_to_arrow(self._schema).schema

985

986 def to_dataframe_schema(self) -> DataFrameSchema:

987 """Convert to a DataFrameSchema.

988

989 Returns

990 -------

991 dataframe_schema : `DataFrameSchema`

992 Converted dataframe schema.

993 """

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

995

996 def to_arrow_numpy_schema(self) -> ArrowNumpySchema:

997 """Convert to an `ArrowNumpySchema`.

998

999 Returns

1000 -------

1001 arrow_numpy_schema : `ArrowNumpySchema`

1002 Converted arrow numpy schema.

1003 """

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

1005

1006 @property

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

1008 return self._schema

1009

1010 def __repr__(self) -> str:

1011 return repr(self._schema)

1012

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

1014 if not isinstance(other, ArrowAstropySchema):

1015 return NotImplemented

1016

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

1018 # same column names.

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

1020 return False

1021

1022 for name in self._schema.columns:

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

1024 return False

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

1026 return False

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

1028 return False

1029

1030 return True

1031

1032

1033class ArrowNumpySchema:

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

1035

1036 Parameters

1037 ----------

1038 numpy_dtype : `numpy.dtype`

1039 Numpy dtype to convert.

1040 """

1041

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

1043 self._dtype = numpy_dtype

1044

1045 @classmethod

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

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

1048

1049 Parameters

1050 ----------

1051 schema : `pyarrow.Schema`

1052 Pyarrow schema to convert.

1053

1054 Returns

1055 -------

1056 numpy_schema : `ArrowNumpySchema`

1057 Converted arrow numpy schema.

1058 """

1059 import numpy as np

1060

1061 dtype = _schema_to_dtype_list(schema)

1062

1063 return cls(np.dtype(dtype))

1064

1065 def to_arrow_astropy_schema(self) -> ArrowAstropySchema:

1066 """Convert to an `ArrowAstropySchema`.

1067

1068 Returns

1069 -------

1070 astropy_schema : `ArrowAstropySchema`

1071 Converted arrow astropy schema.

1072 """

1073 import numpy as np

1074

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

1076

1077 def to_dataframe_schema(self) -> DataFrameSchema:

1078 """Convert to a `DataFrameSchema`.

1079

1080 Returns

1081 -------

1082 dataframe_schema : `DataFrameSchema`

1083 Converted dataframe schema.

1084 """

1085 import numpy as np

1086

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

1088

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

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

1091

1092 Returns

1093 -------

1094 arrow_schema : `pyarrow.Schema`

1095 Converted pyarrow schema.

1096 """

1097 import numpy as np

1098

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

1100

1101 @property

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

1103 return self._dtype

1104

1105 def __repr__(self) -> str:

1106 return repr(self._dtype)

1107

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

1109 if not isinstance(other, ArrowNumpySchema):

1110 return NotImplemented

1111

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

1113 return False

1114

1115 return True

1116

1117

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

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

1120

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

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

1123 tuple.

1124

1125 Parameters

1126 ----------

1127 n : `int`

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

1129 reconstructed.

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

1131 Strings to be split.

1132

1133 Returns

1134 -------

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

1136 A list of multi-index column name tuples.

1137 """

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

1139

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

1141 for name in names:

1142 m = re.search(pattern, name)

1143 if m is not None:

1144 column_names.append(m.groups())

1145

1146 return column_names

1147

1148

1149def _standardize_multi_index_columns(

1150 pd_index: pd.MultiIndex,

1151 columns: Any,

1152 stringify: bool = True,

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

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

1155 into a string directly understandable by PyArrow.

1156

1157 Parameters

1158 ----------

1159 pd_index : `pandas.MultiIndex`

1160 Pandas multi-index.

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

1162 Columns to standardize.

1163 stringify : `bool`, optional

1164 Should the column names be stringified?

1165

1166 Returns

1167 -------

1168 names : `list` [`str`]

1169 Stringified representation of a multi-index column name.

1170 """

1171 index_level_names = tuple(pd_index.names)

1172

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

1174

1175 if isinstance(columns, list):

1176 for requested in columns:

1177 if not isinstance(requested, tuple):

1178 raise ValueError(

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

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

1181 )

1182 if stringify:

1183 names.append(str(requested))

1184 else:

1185 names.append(requested)

1186 else:

1187 if not isinstance(columns, collections.abc.Mapping): 1187 ↛ 1188line 1187 didn't jump to line 1188 because the condition on line 1187 was never true

1188 raise ValueError(

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

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

1191 )

1192 if not set(index_level_names).issuperset(columns.keys()): 1192 ↛ 1193line 1192 didn't jump to line 1193 because the condition on line 1192 was never true

1193 raise ValueError(

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

1195 )

1196 factors = [

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

1198 for i, level in enumerate(index_level_names)

1199 ]

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

1201 for i, value in enumerate(requested):

1202 if value not in pd_index.levels[i]: 1202 ↛ 1203line 1202 didn't jump to line 1203 because the condition on line 1202 was never true

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

1204 if stringify:

1205 names.append(str(requested))

1206 else:

1207 names.append(requested)

1208

1209 return names

1210

1211

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

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

1214

1215 Parameters

1216 ----------

1217 astropy_table : `astropy.table.Table`

1218 Table to apply metadata.

1219 arrow_schema : `pyarrow.Schema`

1220 Arrow schema with metadata.

1221 """

1222 from astropy.table import meta

1223

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

1225

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

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

1228 if meta_yaml:

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

1230 meta_hdr = meta.get_header_from_yaml(meta_yaml)

1231

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

1233 # metadata that was serialized with the table.

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

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

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

1237 if attr in header_cols[col.name]:

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

1239

1240 if "meta" in meta_hdr:

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

1242 else:

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

1244 # metadata.

1245 for name in arrow_schema.names:

1246 field_metadata = arrow_schema.field(name).metadata

1247 if field_metadata is None:

1248 continue

1249 if (

1250 b"description" in field_metadata

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

1252 ):

1253 astropy_table[name].description = description

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

1255 astropy_table[name].unit = unit

1256

1257 # Ensure that the special ASTROPY_PANDAS_INDEX_KEY is propagated to

1258 # the table metadata.

1259 if index_key := metadata.get(ASTROPY_PANDAS_INDEX_KEY.encode(), None):

1260 astropy_table.meta[ASTROPY_PANDAS_INDEX_KEY] = index_key.decode("UTF-8")

1261

1262

1263def _arrow_string_to_numpy_dtype(

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

1265) -> str:

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

1267

1268 Parameters

1269 ----------

1270 schema : `pyarrow.Schema`

1271 Arrow table schema.

1272 name : `str`

1273 Column name.

1274 numpy_column : `numpy.ndarray`, optional

1275 Column to determine numpy string dtype.

1276 default_length : `int`, optional

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

1278 from column.

1279

1280 Returns

1281 -------

1282 dtype_str : `str`

1283 Numpy dtype string.

1284 """

1285 # Special-case for string and binary columns

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

1287 strlen = default_length

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

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

1290 # String/bytes length from header.

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

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

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

1294 strlen = max(lengths) if lengths else 0

1295

1296 dtype = f"U{strlen}" if _is_string(schema.field(name).type) else f"|S{strlen}"

1297

1298 return dtype

1299

1300

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

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

1303

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

1305 string and byte columns.

1306

1307 Parameters

1308 ----------

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

1310 Metadata dictionary; modified in place.

1311 name : `str`

1312 Column name.

1313 dtype : `np.dtype`

1314 Numpy dtype.

1315 """

1316 import numpy as np

1317

1318 if dtype.type is np.str_:

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

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

1321 elif dtype.type is np.bytes_:

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

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

1324

1325

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

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

1328

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

1330 multi-dimensional columns.

1331

1332 Parameters

1333 ----------

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

1335 Metadata dictionary; modified in place.

1336 name : `str`

1337 Column name.

1338 dtype : `np.dtype`

1339 Numpy dtype.

1340 """

1341 if len(dtype.shape) > 1:

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

1343

1344

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

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

1347

1348 Parameters

1349 ----------

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

1351 Metadata dictionary.

1352 list_size : `int`

1353 Size of the list datatype.

1354 name : `str`

1355 Column name.

1356

1357 Returns

1358 -------

1359 shape : `tuple` [`int`]

1360 Shape associated with the column.

1361

1362 Raises

1363 ------

1364 RuntimeError

1365 Raised if metadata is found but has incorrect format.

1366 """

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

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

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

1370 if groups is None: 1370 ↛ 1371line 1370 didn't jump to line 1371 because the condition on line 1370 was never true

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

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

1373 else:

1374 shape = (list_size,)

1375

1376 return shape

1377

1378

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

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

1381

1382 Parameters

1383 ----------

1384 schema : `pyarrow.Schema`

1385 Input pyarrow schema.

1386

1387 Returns

1388 -------

1389 dtype_list: `list` [`tuple`]

1390 A list with name, type pairs.

1391 """

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

1393

1394 dtype: list[Any] = []

1395 for name in schema.names:

1396 t = schema.field(name).type

1397 if isinstance(t, pa.FixedSizeListType):

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

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

1400 elif not (_is_string(t) or _is_binary(t)):

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

1402 else:

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

1404

1405 return dtype

1406

1407

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

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

1410

1411 Parameters

1412 ----------

1413 dtype : `numpy.dtype`

1414 Numpy dtype to convert.

1415

1416 Returns

1417 -------

1418 type_list : `list` [`object`]

1419 Converted list of arrow types.

1420 """

1421 from math import prod

1422

1423 import numpy as np

1424

1425 type_list: list[Any] = []

1426 if dtype.names is None: 1426 ↛ 1427line 1426 didn't jump to line 1427 because the condition on line 1426 was never true

1427 return type_list

1428

1429 for name in dtype.names:

1430 dt = dtype[name]

1431 arrow_type: Any

1432 if len(dt.shape) > 0:

1433 arrow_type = pa.list_(

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

1435 prod(dt.shape),

1436 )

1437 elif dt.type == np.datetime64:

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

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

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

1441 arrow_type = pa.timestamp(time_unit)

1442 else:

1443 try:

1444 arrow_type = pa.from_numpy_dtype(dt.type)

1445 except pa.ArrowNotImplementedError as e:

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

1447 if dt == np.dtype("O"): 1447 ↛ 1449line 1447 didn't jump to line 1449 because the condition on line 1447 was always true

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

1449 e.add_note(msg)

1450 raise

1451 type_list.append((name, arrow_type))

1452

1453 return type_list

1454

1455

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

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

1458

1459 Parameters

1460 ----------

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

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

1463

1464 Returns

1465 -------

1466 dtype : `numpy.dtype`

1467 dtype of equivalent table.

1468 rowcount : `int`

1469 Number of rows in the table.

1470

1471 Raises

1472 ------

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

1474 """

1475 import numpy as np

1476

1477 dtype_list: list[tuple] = []

1478 rowcount = 0

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

1480 if rowcount == 0:

1481 rowcount = len(col)

1482 if len(col) != rowcount:

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

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

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

1486 else:

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

1488 dtype = np.dtype(dtype_list)

1489

1490 return (dtype, rowcount)

1491

1492

1493def _numpy_style_arrays_to_arrow_arrays(

1494 dtype: np.dtype,

1495 rowcount: int,

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

1497 schema: pa.Schema,

1498) -> list[pa.Array]:

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

1500

1501 Parameters

1502 ----------

1503 dtype : `numpy.dtype`

1504 Numpy dtype of input table/arrays.

1505 rowcount : `int`

1506 Number of rows in input table/arrays.

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

1508 or `astropy.table.Table`

1509 Arrays to convert to arrow.

1510 schema : `pyarrow.Schema`

1511 Schema of arrow table.

1512

1513 Returns

1514 -------

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

1516 List of converted pyarrow arrays.

1517 """

1518 import numpy as np

1519

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

1521 if dtype.names is None: 1521 ↛ 1522line 1521 didn't jump to line 1522 because the condition on line 1521 was never true

1522 return arrow_arrays

1523

1524 for name in dtype.names:

1525 dt = dtype[name]

1526 val: Any

1527 if len(dt.shape) > 0:

1528 if rowcount > 0:

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

1530 else:

1531 val = []

1532 else:

1533 val = np_style_arrays[name]

1534

1535 try:

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

1537 except pa.ArrowNotImplementedError as err:

1538 # Check if val is big-endian.

1539 if (np.little_endian and val.dtype.byteorder == ">") or ( 1539 ↛ 1548line 1539 didn't jump to line 1548 because the condition on line 1539 was always true

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

1541 ):

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

1543 val2 = val.byteswap()

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

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

1546 else:

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

1548 raise err

1549

1550 return arrow_arrays

1551

1552

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

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

1555

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

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

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

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

1560 ``snappy`` compression algorithm.

1561

1562 Parameters

1563 ----------

1564 schema : `pyarrow.Schema`

1565 Arrow table schema.

1566 target_size : `int`, optional

1567 The target size (in bytes).

1568

1569 Returns

1570 -------

1571 row_group_size : `int`

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

1573 """

1574 bit_width = 0

1575

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

1577

1578 for name in schema.names:

1579 t = schema.field(name).type

1580

1581 if _is_string(t) or _is_binary(t):

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

1583

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

1585 # String/bytes length from header.

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

1587 else:

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

1589 strlen = 10

1590

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

1592 t_width = 8 * strlen

1593 elif isinstance(t, pa.FixedSizeListType):

1594 if t.value_type == pa.null(): 1594 ↛ 1595line 1594 didn't jump to line 1595 because the condition on line 1594 was never true

1595 t_width = 0

1596 else:

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

1598 elif t == pa.null():

1599 t_width = 0

1600 elif isinstance(t, pa.ListType):

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

1602 t_width = 0

1603 else:

1604 # This is a variable length list, just choose

1605 # something arbitrary.

1606 t_width = 10 * t.value_type.bit_width

1607 else:

1608 t_width = t.bit_width

1609

1610 bit_width += t_width

1611

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

1613 if bit_width < 8:

1614 bit_width = 8

1615

1616 byte_width = bit_width // 8

1617

1618 return target_size // byte_width

1619

1620

1621def _is_string(t: pa.DataType) -> bool:

1622 return pa.types.is_string(t) or pa.types.is_large_string(t) or pa.types.is_string_view(t)

1623

1624

1625def _is_binary(t: pa.DataType) -> bool:

1626 return (

1627 pa.types.is_binary(t)

1628 or pa.types.is_large_binary(t)

1629 or pa.types.is_binary_view(t)

1630 or pa.types.is_fixed_size_binary(t)

1631 )

1632

1633

1634def _get_pandas_index_columns(md: dict) -> list[str]:

1635 if isinstance(md["index_columns"][0], dict): 1635 ↛ 1637line 1635 didn't jump to line 1637 because the condition on line 1635 was never true

1636 # For parquet files written with pandas3 default parquet settings.

1637 index_columns = [col["name"] for col in md["index_columns"]]

1638 else:

1639 # For parquet files written with pandas2 default parquet settings

1640 # or this parquet writer.

1641 index_columns = md["index_columns"]

1642

1643 return index_columns