Coverage for tests/test_parquet.py: 17%

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/>.

28"""Tests for ParquetFormatter.

30Tests in this module are disabled unless pandas and pyarrow are importable.

31"""

33import datetime

34import os

35import posixpath

36import shutil

37import unittest

38import uuid

40try:

41 import pyarrow as pa

42except ImportError:

43 pa = None

44try:

45 import astropy.table as atable

46 from astropy import units

47except ImportError:

48 atable = None

49try:

50 import numpy as np

51except ImportError:

52 np = None

53try:

54 import pandas as pd

55except ImportError:

56 pd = None

58try:

59 import boto3

60 import botocore

62 from lsst.resources.s3utils import clean_test_environment_for_s3

64 try:

65 from moto import mock_aws # v5

66 except ImportError:

67 from moto import mock_s3 as mock_aws

68except ImportError:

69 boto3 = None

71try:

72 import fsspec

73except ImportError:

74 fsspec = None

76try:

77 import s3fs

78except ImportError:

79 s3fs = None

82from lsst.daf.butler import (

83 Butler,

84 Config,

85 DatasetProvenance,

86 DatasetRef,

87 DatasetType,

88 FileDataset,

89 StorageClassConfig,

90 StorageClassFactory,

91)

92from lsst.resources import ResourcePath

94try:

95 from lsst.daf.butler.delegates.arrowtable import ArrowTableDelegate

96except ImportError:

97 pa = None

99try:

100 from lsst.daf.butler.formatters.parquet import (

101 ASTROPY_PANDAS_INDEX_KEY,

102 ArrowAstropySchema,

103 ArrowNumpySchema,

104 DataFrameSchema,

105 ParquetFormatter,

106 _append_numpy_multidim_metadata,

107 _astropy_to_numpy_dict,

108 _numpy_dict_to_numpy,

109 _numpy_dtype_to_arrow_types,

110 _numpy_style_arrays_to_arrow_arrays,

111 _numpy_to_numpy_dict,

112 add_pandas_index_to_astropy,

113 arrow_to_astropy,

114 arrow_to_numpy,

115 arrow_to_numpy_dict,

116 arrow_to_pandas,

117 astropy_to_arrow,

118 astropy_to_pandas,

119 compute_row_group_size,

120 numpy_dict_to_arrow,

121 numpy_to_arrow,

122 pandas_to_arrow,

123 pandas_to_astropy,

124 )

125except ImportError:

126 pa = None

127 pd = None

128 atable = None

129 np = None

130from lsst.daf.butler.tests.utils import makeTestTempDir, removeTestTempDir

131

132TESTDIR = os.path.abspath(os.path.dirname(__file__))

133

134

135def _makeSimpleNumpyTable(include_multidim=False, include_bigendian=False):

136 """Make a simple numpy table with random data.

137

138 Parameters

139 ----------

140 include_multidim : `bool`

141 Include multi-dimensional columns.

142 include_bigendian : `bool`

143 Include big-endian columns.

144

145 Returns

146 -------

147 numpyTable : `numpy.ndarray`

148 """

149 nrow = 5

150

151 dtype = [

152 ("index", "i4"),

153 ("a", "f8"),

154 ("b", "f8"),

155 ("c", "f8"),

156 ("ddd", "f8"),

157 ("f", "i8"),

158 ("strcol", "U10"),

159 ("bytecol", "S10"),

160 ("dtn", "datetime64[ns]"),

161 ("dtu", "datetime64[us]"),

162 ]

163

164 if include_multidim:

165 dtype.extend(

166 [

167 ("d1", "f4", (5,)),

168 ("d2", "i8", (5, 10)),

169 ("d3", "f8", (5, 10)),

170 ]

171 )

172

173 if include_bigendian:

174 dtype.extend([("a_bigendian", ">f8"), ("f_bigendian", ">i8")])

175

176 data = np.zeros(nrow, dtype=dtype)

177 data["index"][:] = np.arange(nrow)

178 data["a"] = np.random.randn(nrow)

179 data["b"] = np.random.randn(nrow)

180 data["c"] = np.random.randn(nrow)

181 data["ddd"] = np.random.randn(nrow)

182 data["f"] = np.arange(nrow) * 10

183 data["strcol"][:] = "teststring"

184 data["bytecol"][:] = "teststring"

185 data["dtn"] = datetime.datetime.fromisoformat("2024-07-23")

186 data["dtu"] = datetime.datetime.fromisoformat("2024-07-23")

187

188 if include_multidim:

189 data["d1"] = np.random.randn(data["d1"].size).reshape(data["d1"].shape)

190 data["d2"] = np.arange(data["d2"].size).reshape(data["d2"].shape)

191 data["d3"] = np.asfortranarray(np.random.randn(data["d3"].size).reshape(data["d3"].shape))

192

193 if include_bigendian:

194 data["a_bigendian"][:] = data["a"]

195 data["f_bigendian"][:] = data["f"]

196

197 return data

198

199

200def _makeSingleIndexDataFrame(include_masked=False, include_lists=False):

201 """Make a single index data frame for testing.

202

203 Parameters

204 ----------

205 include_masked : `bool`

206 Include masked columns.

207 include_lists : `bool`

208 Include list columns.

209

210 Returns

211 -------

212 dataFrame : `~pandas.DataFrame`

213 The test dataframe.

214 allColumns : `list` [`str`]

215 List of all the columns (including index columns).

216 """

217 data = _makeSimpleNumpyTable()

218 df = pd.DataFrame(data)

219 df = df.set_index("index")

220

221 if include_masked:

222 nrow = len(df)

223

224 df["m1"] = pd.array(np.arange(nrow), dtype=pd.Int64Dtype())

225 df["m2"] = pd.array(np.arange(nrow), dtype=np.float32)

226 df["mstrcol"] = pd.array(np.array(["text"] * nrow))

227 df.loc[1, ["m1", "m2", "mstrcol"]] = None

228 df.loc[0, "m1"] = 1649900760361600113

229

230 if include_lists:

231 nrow = len(df)

232

233 df["l1"] = [[0, 0]] * nrow

234 df["l2"] = [[0.0, 0.0]] * nrow

235 df["l3"] = [[]] * nrow

236

237 allColumns = df.columns.append(pd.Index(df.index.names))

238

239 return df, allColumns

240

241

242def _makeMultiIndexDataFrame():

243 """Make a multi-index data frame for testing.

244

245 Returns

246 -------

247 dataFrame : `~pandas.DataFrame`

248 The test dataframe.

249 """

250 columns = pd.MultiIndex.from_tuples(

251 [

252 ("g", "a"),

253 ("g", "b"),

254 ("g", "c"),

255 ("r", "a"),

256 ("r", "b"),

257 ("r", "c"),

258 ],

259 names=["filter", "column"],

260 )

261 df = pd.DataFrame(np.random.randn(5, 6), index=np.arange(5, dtype=int), columns=columns)

262

263 return df

264

265

266def _makeSimpleAstropyTable(include_multidim=False, include_masked=False, include_bigendian=False):

267 """Make an astropy table for testing.

268

269 Parameters

270 ----------

271 include_multidim : `bool`

272 Include multi-dimensional columns.

273 include_masked : `bool`

274 Include masked columns.

275 include_bigendian : `bool`

276 Include big-endian columns.

277

278 Returns

279 -------

280 astropyTable : `astropy.table.Table`

281 The test table.

282 """

283 data = _makeSimpleNumpyTable(include_multidim=include_multidim, include_bigendian=include_bigendian)

284 # Add a couple of units.

285 table = atable.Table(data)

286 table["a"].unit = units.degree

287 table["a"].description = "Description of column a"

288 table["b"].unit = units.meter

289 table["b"].description = "Description of column b"

290

291 # Add some masked columns.

292 if include_masked:

293 nrow = len(table)

294 mask = np.zeros(nrow, dtype=bool)

295 mask[1] = True

296 # We set the masked columns with the underlying sentinel value

297 # to be able test after serialization.

298

299 # Masked 64-bit integer.

300 arr = np.arange(nrow, dtype="i8")

301 arr[mask] = -1

302 arr[0] = 1649900760361600113

303 table["m_i8"] = np.ma.masked_array(data=arr, mask=mask, fill_value=-1)

304 # Masked 32-bit float.

305 arr = np.arange(nrow, dtype="f4")

306 arr[mask] = np.nan

307 table["m_f4"] = np.ma.masked_array(data=arr, mask=mask, fill_value=np.nan)

308 # Unmasked 32-bit float with NaNs.

309 table["um_f4"] = arr

310 # Masked 64-bit float.

311 arr = np.arange(nrow, dtype="f8")

312 arr[mask] = np.nan

313 table["m_f8"] = np.ma.masked_array(data=arr, mask=mask, fill_value=np.nan)

314 # Unmasked 64-bit float with NaNs.

315 table["um_f8"] = arr

316 # Masked boolean.

317 arr = np.zeros(nrow, dtype=np.bool_)

318 arr[mask] = True

319 table["m_bool"] = np.ma.masked_array(data=arr, mask=mask, fill_value=True)

320 # Masked unsigned 32-bit unsigned int.

321 arr = np.arange(nrow, dtype="u4")

322 arr[mask] = 0

323 table["m_u4"] = np.ma.masked_array(data=arr, mask=mask, fill_value=0)

324 # Masked string.

325 table["m_str"] = np.ma.masked_array(data=np.array(["text"] * nrow), mask=mask, fill_value="")

326 # Masked bytes.

327 table["m_byte"] = np.ma.masked_array(data=np.array([b"bytes"] * nrow), mask=mask, fill_value=b"")

328

329 return table

330

331

332def _makeSimpleArrowTable(include_multidim=False, include_masked=False):

333 """Make an arrow table for testing.

334

335 Parameters

336 ----------

337 include_multidim : `bool`

338 Include multi-dimensional columns.

339 include_masked : `bool`

340 Include masked columns.

341

342 Returns

343 -------

344 arrowTable : `pyarrow.Table`

345 The test table.

346 """

347 data = _makeSimpleAstropyTable(include_multidim=include_multidim, include_masked=include_masked)

348 return astropy_to_arrow(data)

349

350

351@unittest.skipUnless(pd is not None, "Cannot test ParquetFormatterDataFrame without pandas.")

352@unittest.skipUnless(pa is not None, "Cannot test ParquetFormatterDataFrame without pyarrow.")

353class ParquetFormatterDataFrameTestCase(unittest.TestCase):

354 """Tests for ParquetFormatter, DataFrame, using local file datastore."""

355

356 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")

357

358 def setUp(self):

359 """Create a new butler root for each test."""

360 self.root = makeTestTempDir(TESTDIR)

361 config = Config(self.configFile)

362 self.run = "test_run"

363 self.butler = Butler.from_config(

364 Butler.makeRepo(self.root, config=config), writeable=True, run=self.run

365 )

366 # No dimensions in dataset type so we don't have to worry about

367 # inserting dimension data or defining data IDs.

368 self.datasetType = DatasetType(

369 "data", dimensions=(), storageClass="DataFrame", universe=self.butler.dimensions

370 )

371 self.butler.registry.registerDatasetType(self.datasetType)

372

373 def tearDown(self):

374 removeTestTempDir(self.root)

375

376 def testSingleIndexDataFrame(self):

377 df1, allColumns = _makeSingleIndexDataFrame(include_masked=True)

378

379 self.butler.put(df1, self.datasetType, dataId={})

380 # Read the whole DataFrame.

381 df2 = self.butler.get(self.datasetType, dataId={})

382 self.assertTrue(df1.equals(df2))

383 # Read just the column descriptions.

384 columns2 = self.butler.get(self.datasetType.componentTypeName("columns"), dataId={})

385 self.assertTrue(allColumns.equals(columns2))

386 # Read the rowcount.

387 rowcount = self.butler.get(self.datasetType.componentTypeName("rowcount"), dataId={})

388 self.assertEqual(rowcount, len(df1))

389 # Read the schema.

390 schema = self.butler.get(self.datasetType.componentTypeName("schema"), dataId={})

391 self.assertEqual(schema, DataFrameSchema(df1))

392 # Read just some columns a few different ways.

393 df3 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["a", "c"]})

394 self.assertTrue(df1.loc[:, ["a", "c"]].equals(df3))

395 df4 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": "a"})

396 self.assertTrue(df1.loc[:, ["a"]].equals(df4))

397 df5 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["index", "a"]})

398 self.assertTrue(df1.loc[:, ["a"]].equals(df5))

399 df6 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": "ddd"})

400 self.assertTrue(df1.loc[:, ["ddd"]].equals(df6))

401 df7 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["a", "a"]})

402 self.assertTrue(df1.loc[:, ["a"]].equals(df7))

403 df8 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["d*"]})

404 self.assertTrue(df1.loc[:, ["ddd", "dtn", "dtu"]].equals(df8))

405 df9 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["d*", "d*"]})

406 self.assertTrue(df1.loc[:, ["ddd", "dtn", "dtu"]].equals(df9))

407 # Passing an unrecognized column should be a ValueError.

408 with self.assertRaises(ValueError):

409 self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["e"]})

410

411 def testSingleIndexDataFrameWithLists(self):

412 df1, allColumns = _makeSingleIndexDataFrame(include_lists=True)

413

414 self.butler.put(df1, self.datasetType, dataId={})

415 # Read the whole DataFrame.

416 df2 = self.butler.get(self.datasetType, dataId={})

417

418 # We need to check the list columns specially because they go

419 # from lists to arrays.

420 for col in ["l1", "l2", "l3"]:

421 for i in range(len(df1)):

422 self.assertTrue(np.all(df2[col].values[i] == df1[col].values[i]))

423

424 def testMultiIndexDataFrame(self):

425 df1 = _makeMultiIndexDataFrame()

426

427 self.butler.put(df1, self.datasetType, dataId={})

428 # Read the whole DataFrame.

429 df2 = self.butler.get(self.datasetType, dataId={})

430 self.assertTrue(df1.equals(df2))

431 # Read just the column descriptions.

432 columns2 = self.butler.get(self.datasetType.componentTypeName("columns"), dataId={})

433 self.assertTrue(df1.columns.equals(columns2))

434 self.assertEqual(columns2.names, df1.columns.names)

435 # Read the rowcount.

436 rowcount = self.butler.get(self.datasetType.componentTypeName("rowcount"), dataId={})

437 self.assertEqual(rowcount, len(df1))

438 # Read the schema.

439 schema = self.butler.get(self.datasetType.componentTypeName("schema"), dataId={})

440 self.assertEqual(schema, DataFrameSchema(df1))

441 # Read just some columns a few different ways.

442 df3 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": {"filter": "g"}})

443 self.assertTrue(df1.loc[:, ["g"]].equals(df3))

444 df4 = self.butler.get(

445 self.datasetType, dataId={}, parameters={"columns": {"filter": ["r"], "column": "a"}}

446 )

447 self.assertTrue(df1.loc[:, [("r", "a")]].equals(df4))

448 column_list = [("g", "a"), ("r", "c")]

449 df5 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": column_list})

450 self.assertTrue(df1.loc[:, column_list].equals(df5))

451 column_dict = {"filter": "r", "column": ["a", "b"]}

452 df6 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": column_dict})

453 self.assertTrue(df1.loc[:, [("r", "a"), ("r", "b")]].equals(df6))

454 # Passing an unrecognized column should be a ValueError.

455 with self.assertRaises(ValueError):

456 self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["d"]})

457

458 def testSingleIndexDataFrameEmptyString(self):

459 """Test persisting a single index dataframe with empty strings."""

460 df1, _ = _makeSingleIndexDataFrame()

461

462 # Set one of the strings to None

463 df1.at[1, "strcol"] = None

464

465 self.butler.put(df1, self.datasetType, dataId={})

466 # Read the whole DataFrame.

467 df2 = self.butler.get(self.datasetType, dataId={})

468 self.assertTrue(df1.equals(df2))

469

470 def testSingleIndexDataFrameAllEmptyStrings(self):

471 """Test persisting a single index dataframe with an empty string

472 column.

473 """

474 df1, _ = _makeSingleIndexDataFrame()

475

476 # Set all of the strings to None

477 df1.loc[0:, "strcol"] = None

478

479 self.butler.put(df1, self.datasetType, dataId={})

480 # Read the whole DataFrame.

481 df2 = self.butler.get(self.datasetType, dataId={})

482 self.assertTrue(df1.equals(df2))

483

484 def testLegacyDataFrame(self):

485 """Test writing a dataframe to parquet via pandas (without additional

486 metadata) and ensure that we can read it back with all the new

487 functionality.

488 """

489 df1, allColumns = _makeSingleIndexDataFrame()

490

491 fname = os.path.join(self.root, "test_dataframe.parq")

492 df1.to_parquet(fname)

493

494 legacy_type = DatasetType(

495 "legacy_dataframe",

496 dimensions=(),

497 storageClass="DataFrame",

498 universe=self.butler.dimensions,

499 )

500 self.butler.registry.registerDatasetType(legacy_type)

501

502 data_id = {}

503 ref = DatasetRef(legacy_type, data_id, run=self.run)

504 dataset = FileDataset(path=fname, refs=[ref], formatter=ParquetFormatter)

505

506 self.butler.ingest(dataset, transfer="copy")

507

508 self.butler.put(df1, self.datasetType, dataId={})

509

510 df2a = self.butler.get(self.datasetType, dataId={})

511 df2b = self.butler.get("legacy_dataframe", dataId={})

512 self.assertTrue(df2a.equals(df2b))

513

514 df3a = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["a"]})

515 df3b = self.butler.get("legacy_dataframe", dataId={}, parameters={"columns": ["a"]})

516 self.assertTrue(df3a.equals(df3b))

517

518 columns2a = self.butler.get(self.datasetType.componentTypeName("columns"), dataId={})

519 columns2b = self.butler.get("legacy_dataframe.columns", dataId={})

520 self.assertTrue(columns2a.equals(columns2b))

521

522 rowcount2a = self.butler.get(self.datasetType.componentTypeName("rowcount"), dataId={})

523 rowcount2b = self.butler.get("legacy_dataframe.rowcount", dataId={})

524 self.assertEqual(rowcount2a, rowcount2b)

525

526 schema2a = self.butler.get(self.datasetType.componentTypeName("schema"), dataId={})

527 schema2b = self.butler.get("legacy_dataframe.schema", dataId={})

528 self.assertEqual(schema2a, schema2b)

529

530 def testDataFrameSchema(self):

531 tab1 = _makeSimpleArrowTable()

532

533 schema = DataFrameSchema.from_arrow(tab1.schema)

534

535 self.assertIsInstance(schema.schema, pd.DataFrame)

536 self.assertEqual(repr(schema), repr(schema._schema))

537 self.assertNotEqual(schema, "not_a_schema")

538 self.assertEqual(schema, schema)

539

540 tab2 = _makeMultiIndexDataFrame()

541 schema2 = DataFrameSchema(tab2)

542

543 self.assertNotEqual(schema, schema2)

544

545 @unittest.skipUnless(atable is not None, "Cannot test reading as astropy without astropy.")

546 def testWriteSingleIndexDataFrameReadAsAstropyTable(self):

547 df1, allColumns = _makeSingleIndexDataFrame()

548

549 self.butler.put(df1, self.datasetType, dataId={})

550

551 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowAstropy")

552

553 tab2_df = tab2.to_pandas(index="index")

554 self.assertTrue(df1.equals(tab2_df))

555

556 # Check reading the columns.

557 columns = list(tab2.columns.keys())

558 columns2 = self.butler.get(

559 self.datasetType.componentTypeName("columns"), dataId={}, storageClass="ArrowColumnList"

560 )

561 # We check the set because pandas reorders the columns.

562 self.assertEqual(set(columns2), set(columns))

563

564 # Check reading the schema.

565 schema = ArrowAstropySchema(tab2)

566 schema2 = self.butler.get(

567 self.datasetType.componentTypeName("schema"), dataId={}, storageClass="ArrowAstropySchema"

568 )

569

570 # The string types are objectified by pandas, and the order

571 # will be changed because of pandas indexing.

572 self.assertEqual(len(schema2.schema.columns), len(schema.schema.columns))

573 for name in schema.schema.columns:

574 self.assertIn(name, schema2.schema.columns)

575 if schema2.schema[name].dtype != np.dtype("O"):

576 self.assertEqual(schema2.schema[name].dtype, schema.schema[name].dtype)

577

578 @unittest.skipUnless(atable is not None, "Cannot test reading as astropy without astropy.")

579 def testWriteSingleIndexDataFrameWithMaskedColsReadAsAstropyTable(self):

580 # We need to special-case the write-as-pandas read-as-astropy code

581 # with masks because pandas has multiple ways to use masked columns.

582 # (The string column mask handling in particular is frustratingly

583 # inconsistent.)

584 df1, allColumns = _makeSingleIndexDataFrame(include_masked=True)

585

586 self.butler.put(df1, self.datasetType, dataId={})

587

588 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowAstropy")

589 tab2_df = astropy_to_pandas(tab2, index="index")

590

591 self.assertTrue(df1.columns.equals(tab2_df.columns))

592 for name in tab2_df.columns:

593 col1 = df1[name]

594 col2 = tab2_df[name]

595

596 if col1.hasnans:

597 notNull = col1.notnull()

598 self.assertTrue(notNull.equals(col2.notnull()))

599 # Need to check value-by-value because column may

600 # be made of objects, depending on what pandas decides.

601 for index in notNull.values.nonzero()[0]:

602 self.assertEqual(col1[index], col2[index])

603 else:

604 self.assertTrue(col1.equals(col2))

605

606 @unittest.skipUnless(atable is not None, "Cannot test reading as astropy without astropy.")

607 def testWriteMultiIndexDataFrameReadAsAstropyTable(self):

608 df1 = _makeMultiIndexDataFrame()

609

610 self.butler.put(df1, self.datasetType, dataId={})

611

612 _ = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowAstropy")

613

614 # This is an odd duck, it doesn't really round-trip.

615 # This test simply checks that it's readable, but definitely not

616 # recommended.

617

618 @unittest.skipUnless(atable is not None, "Cannot test writing as astropy without astropy.")

619 def testWriteAstropyTableWithMaskedColsReadAsSingleIndexDataFrame(self):

620 tab1 = _makeSimpleAstropyTable(include_masked=True)

621

622 self.butler.put(tab1, self.datasetType, dataId={})

623

624 tab2 = self.butler.get(self.datasetType, dataId={})

625

626 tab1_df = astropy_to_pandas(tab1)

627 self.assertTrue(tab1_df.equals(tab2))

628

629 tab2_astropy = pandas_to_astropy(tab2)

630 for col in tab1.dtype.names:

631 np.testing.assert_array_equal(tab2_astropy[col], tab1[col])

632 if isinstance(tab1[col], atable.column.MaskedColumn):

633 np.testing.assert_array_equal(tab2_astropy[col].mask, tab1[col].mask)

634

635 @unittest.skipUnless(pa is not None, "Cannot test reading as arrow without pyarrow.")

636 def testWriteSingleIndexDataFrameReadAsArrowTable(self):

637 df1, allColumns = _makeSingleIndexDataFrame()

638

639 self.butler.put(df1, self.datasetType, dataId={})

640

641 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowTable")

642

643 tab2_df = arrow_to_pandas(tab2)

644 self.assertTrue(df1.equals(tab2_df))

645

646 # Check reading the columns.

647 columns = list(tab2.schema.names)

648 columns2 = self.butler.get(

649 self.datasetType.componentTypeName("columns"), dataId={}, storageClass="ArrowColumnList"

650 )

651 # We check the set because pandas reorders the columns.

652 self.assertEqual(set(columns), set(columns2))

653

654 # Override the component using a dataset type.

655 columnsType = self.datasetType.makeComponentDatasetType("columns").overrideStorageClass(

656 "ArrowColumnList"

657 )

658 self.assertEqual(columns2, self.butler.get(columnsType))

659

660 # Check getting a component while overriding the storage class via

661 # the dataset type. This overrides the parent storage class and then

662 # selects the component.

663 columnsType = self.datasetType.overrideStorageClass("ArrowAstropy").makeComponentDatasetType(

664 "columns"

665 )

666 self.assertEqual(columns2, self.butler.get(columnsType))

667

668 # Check reading the schema.

669 schema = tab2.schema

670 schema2 = self.butler.get(

671 self.datasetType.componentTypeName("schema"), dataId={}, storageClass="ArrowSchema"

672 )

673

674 # These will not have the same metadata, nor will the string column

675 # information be maintained.

676 self.assertEqual(len(schema.names), len(schema2.names))

677 for name in schema.names:

678 if schema.field(name).type not in (pa.string(), pa.binary()):

679 self.assertEqual(schema.field(name).type, schema2.field(name).type)

680

681 @unittest.skipUnless(pa is not None, "Cannot test reading as arrow without pyarrow.")

682 def testWriteMultiIndexDataFrameReadAsArrowTable(self):

683 df1 = _makeMultiIndexDataFrame()

684

685 self.butler.put(df1, self.datasetType, dataId={})

686

687 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowTable")

688

689 tab2_df = arrow_to_pandas(tab2)

690 self.assertTrue(df1.equals(tab2_df))

691

692 @unittest.skipUnless(np is not None, "Cannot test reading as numpy without numpy.")

693 def testWriteSingleIndexDataFrameReadAsNumpyTable(self):

694 df1, allColumns = _makeSingleIndexDataFrame()

695

696 self.butler.put(df1, self.datasetType, dataId={})

697

698 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowNumpy")

699

700 tab2_df = pd.DataFrame.from_records(tab2, index=["index"])

701 self.assertTrue(df1.equals(tab2_df))

702

703 # Check reading the columns.

704 columns = list(tab2.dtype.names)

705 columns2 = self.butler.get(

706 self.datasetType.componentTypeName("columns"), dataId={}, storageClass="ArrowColumnList"

707 )

708 # We check the set because pandas reorders the columns.

709 self.assertEqual(set(columns2), set(columns))

710

711 # Check reading the schema.

712 schema = ArrowNumpySchema(tab2.dtype)

713 schema2 = self.butler.get(

714 self.datasetType.componentTypeName("schema"), dataId={}, storageClass="ArrowNumpySchema"

715 )

716

717 # The string types will be objectified by pandas, and the order

718 # will be changed because of pandas indexing.

719 self.assertEqual(len(schema.schema.names), len(schema2.schema.names))

720 for name in schema.schema.names:

721 self.assertIn(name, schema2.schema.names)

722 self.assertEqual(schema2.schema[name].type, schema.schema[name].type)

723

724 @unittest.skipUnless(np is not None, "Cannot test reading as numpy without numpy.")

725 def testWriteMultiIndexDataFrameReadAsNumpyTable(self):

726 df1 = _makeMultiIndexDataFrame()

727

728 self.butler.put(df1, self.datasetType, dataId={})

729

730 _ = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowNumpy")

731

732 # This is an odd duck, it doesn't really round-trip.

733 # This test simply checks that it's readable, but definitely not

734 # recommended.

735

736 @unittest.skipUnless(np is not None, "Cannot test reading as numpy dict without numpy.")

737 def testWriteSingleIndexDataFrameReadAsNumpyDict(self):

738 df1, allColumns = _makeSingleIndexDataFrame()

739

740 self.butler.put(df1, self.datasetType, dataId={})

741

742 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowNumpyDict")

743

744 tab2_df = pd.DataFrame.from_records(tab2, index=["index"])

745 # The column order is not maintained.

746 self.assertEqual(set(df1.columns), set(tab2_df.columns))

747 for col in df1.columns:

748 self.assertTrue(np.all(df1[col].values == tab2_df[col].values))

749

750 @unittest.skipUnless(np is not None, "Cannot test reading as numpy dict without numpy.")

751 def testWriteMultiIndexDataFrameReadAsNumpyDict(self):

752 df1 = _makeMultiIndexDataFrame()

753

754 self.butler.put(df1, self.datasetType, dataId={})

755

756 _ = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowNumpyDict")

757

758 # This is an odd duck, it doesn't really round-trip.

759 # This test simply checks that it's readable, but definitely not

760 # recommended.

761

762 def testBadDataFrameColumnParquet(self):

763 df1, allColumns = _makeSingleIndexDataFrame()

764

765 # Make a column with mixed type.

766 bad_col1 = [0.0] * len(df1)

767 bad_col1[1] = 0.0 * units.nJy

768 bad_df = df1.copy()

769 bad_df["bad_col1"] = bad_col1

770

771 # At the moment we cannot check that the correct note is added

772 # to the exception, but that will be possible in the future.

773 with self.assertRaises(RuntimeError):

774 self.butler.put(bad_df, self.datasetType, dataId={})

775

776 @unittest.skipUnless(atable is not None, "Cannot test reading as astropy without astropy.")

777 def testWriteReadAstropyTableLossless(self):

778 tab1 = _makeSimpleAstropyTable(include_multidim=True, include_masked=True)

779

780 put_ref = self.butler.put(tab1, self.datasetType, dataId={})

781

782 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowAstropy")

783

784 # Check that minimal provenance was written by default.

785 expected = {

786 "LSST.BUTLER.ID": str(put_ref.id),

787 "LSST.BUTLER.RUN": "test_run",

788 "LSST.BUTLER.DATASETTYPE": "data",

789 "LSST.BUTLER.N_INPUTS": 0,

790 }

791

792 self.assertEqual(tab2.meta, expected)

793

794 _checkAstropyTableEquality(tab1, tab2)

795

796 @unittest.skipUnless(atable is not None, "Cannot test reading as astropy without astropy.")

797 def testWriteReadAstropyTableProvenance(self):

798 tab1 = _makeSimpleAstropyTable()

799

800 # Create a ref for provenance.

801 astropy_type = DatasetType(

802 "astropy_parquet",

803 dimensions=(),

804 storageClass="ArrowAstropy",

805 universe=self.butler.dimensions,

806 )

807 self.butler.registry.registerDatasetType(astropy_type)

808 input_ref = DatasetRef(astropy_type, {}, run="other_run")

809 quantum_id = uuid.uuid4()

810 provenance = DatasetProvenance(quantum_id=quantum_id)

811 provenance.add_input(input_ref)

812

813 put_ref = self.butler.put(tab1, self.datasetType, dataId={}, provenance=provenance)

814

815 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowAstropy")

816

817 expected = {

818 "LSST.BUTLER.ID": str(put_ref.id),

819 "LSST.BUTLER.RUN": "test_run",

820 "LSST.BUTLER.DATASETTYPE": "data",

821 "LSST.BUTLER.QUANTUM": str(quantum_id),

822 "LSST.BUTLER.N_INPUTS": 1,

823 "LSST.BUTLER.INPUT.0.ID": str(input_ref.id),

824 "LSST.BUTLER.INPUT.0.RUN": "other_run",

825 "LSST.BUTLER.INPUT.0.DATASETTYPE": "astropy_parquet",

826 }

827 self.assertEqual(tab2.meta, expected)

828

829 # Put the dataset again, with different provenance and ensure

830 # that the previous provenance was stripped.

831 self.butler.collections.register("new_run")

832 put_ref3 = self.butler.put(tab2, self.datasetType, dataId={}, run="new_run")

833

834 # tab2 will have been updated in place.

835 expected = {

836 "LSST.BUTLER.ID": str(put_ref3.id),

837 "LSST.BUTLER.RUN": "new_run",

838 "LSST.BUTLER.DATASETTYPE": "data",

839 "LSST.BUTLER.N_INPUTS": 0,

840 }

841 self.assertEqual(tab2.meta, expected)

842 null_prov, prov_ref = DatasetProvenance.from_flat_dict(tab2.meta, self.butler)

843 self.assertEqual(prov_ref, put_ref3)

844 self.assertEqual(null_prov, DatasetProvenance())

845

846 @unittest.skipUnless(np is not None, "Cannot test reading as numpy without numpy.")

847 def testWriteReadNumpyTableLossless(self):

848 tab1 = _makeSimpleNumpyTable(include_multidim=True)

849

850 self.butler.put(tab1, self.datasetType, dataId={})

851

852 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowNumpy")

853

854 _checkNumpyTableEquality(tab1, tab2)

855

856 @unittest.skipUnless(pa is not None, "Cannot test reading as arrow without pyarrow.")

857 def testMaskedNumpy(self):

858 tab1 = _makeSimpleArrowTable(include_multidim=False, include_masked=True)

859 tab1_np = arrow_to_numpy(tab1)

860 self.assertIsInstance(tab1_np, np.ma.MaskedArray)

861 # Stats on a masked column should ignore the nan in row 1.

862 col = tab1_np["m_f8"]

863 self.assertEqual(np.mean(col), 2.25, f"Column: {col}")

864

865 # Now without a mask.

866 tab1 = _makeSimpleArrowTable(include_multidim=False, include_masked=False)

867 tab1_np = arrow_to_numpy(tab1)

868 self.assertNotIsInstance(tab1_np, np.ma.MaskedArray)

869

870 @unittest.skipUnless(pa is not None, "Cannot test reading as arrow without pyarrow.")

871 def testWriteReadArrowTableLossless(self):

872 tab1 = _makeSimpleArrowTable(include_multidim=False, include_masked=True)

873

874 self.butler.put(tab1, self.datasetType, dataId={})

875

876 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowTable")

877

878 self.assertEqual(tab1.schema, tab2.schema)

879 tab1_np = arrow_to_numpy(tab1)

880 tab2_np = arrow_to_numpy(tab2)

881 for col in tab1.column_names:

882 np.testing.assert_array_equal(tab2_np[col], tab1_np[col])

883

884 @unittest.skipUnless(np is not None, "Cannot test reading as numpy dict without numpy.")

885 def testWriteReadNumpyDictLossless(self):

886 tab1 = _makeSimpleNumpyTable(include_multidim=True)

887 dict1 = _numpy_to_numpy_dict(tab1)

888

889 self.butler.put(tab1, self.datasetType, dataId={})

890

891 dict2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowNumpyDict")

892

893 _checkNumpyDictEquality(dict1, dict2)

894

895

896@unittest.skipUnless(pd is not None, "Cannot test InMemoryDatastore with DataFrames without pandas.")

897class InMemoryDataFrameDelegateTestCase(ParquetFormatterDataFrameTestCase):

898 """Tests for InMemoryDatastore, using ArrowTableDelegate with Dataframe."""

899

900 configFile = os.path.join(TESTDIR, "config/basic/butler-inmemory.yaml")

901

902 def testBadDataFrameColumnParquet(self):

903 # This test does not raise for an in-memory datastore.

904 pass

905

906 def testWriteMultiIndexDataFrameReadAsAstropyTable(self):

907 df1 = _makeMultiIndexDataFrame()

908

909 self.butler.put(df1, self.datasetType, dataId={})

910

911 with self.assertRaises(ValueError):

912 _ = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowAstropy")

913

914 def testLegacyDataFrame(self):

915 # This test does not work with an inMemoryDatastore.

916 pass

917

918 def testBadInput(self):

919 df1, _ = _makeSingleIndexDataFrame()

920 delegate = ArrowTableDelegate("DataFrame")

921

922 with self.assertRaises(ValueError):

923 delegate.handleParameters(inMemoryDataset="not_a_dataframe")

924

925 with self.assertRaises(AttributeError):

926 delegate.getComponent(composite=df1, componentName="nothing")

927

928 def testStorageClass(self):

929 df1, allColumns = _makeSingleIndexDataFrame()

930

931 factory = StorageClassFactory()

932 factory.addFromConfig(StorageClassConfig())

933

934 storageClass = factory.findStorageClass(type(df1), compare_types=False)

935 # Force the name lookup to do name matching.

936 storageClass._pytype = None

937 self.assertEqual(storageClass.name, "DataFrame")

938

939 storageClass = factory.findStorageClass(type(df1), compare_types=True)

940 # Force the name lookup to do name matching.

941 storageClass._pytype = None

942 self.assertEqual(storageClass.name, "DataFrame")

943

944

945@unittest.skipUnless(atable is not None, "Cannot test ParquetFormatterArrowAstropy without astropy.")

946@unittest.skipUnless(pa is not None, "Cannot test ParquetFormatterArrowAstropy without pyarrow.")

947class ParquetFormatterArrowAstropyTestCase(unittest.TestCase):

948 """Tests for ParquetFormatter, ArrowAstropy, using local file datastore."""

949

950 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")

951

952 def setUp(self):

953 """Create a new butler root for each test."""

954 self.root = makeTestTempDir(TESTDIR)

955 config = Config(self.configFile)

956 self.run = "test_run"

957 self.butler = Butler.from_config(

958 Butler.makeRepo(self.root, config=config), writeable=True, run=self.run

959 )

960 # No dimensions in dataset type so we don't have to worry about

961 # inserting dimension data or defining data IDs.

962 self.datasetType = DatasetType(

963 "data", dimensions=(), storageClass="ArrowAstropy", universe=self.butler.dimensions

964 )

965 self.butler.registry.registerDatasetType(self.datasetType)

966

967 def tearDown(self):

968 removeTestTempDir(self.root)

969

970 def testAstropyTable(self):

971 tab1 = _makeSimpleAstropyTable(include_multidim=True, include_masked=True)

972

973 self.butler.put(tab1, self.datasetType, dataId={})

974 # Read the whole Table.

975 tab2 = self.butler.get(self.datasetType, dataId={})

976 _checkAstropyTableEquality(tab1, tab2)

977 # Read the columns.

978 columns2 = self.butler.get(self.datasetType.componentTypeName("columns"), dataId={})

979 self.assertEqual(len(columns2), len(tab1.dtype.names))

980 for i, name in enumerate(tab1.dtype.names):

981 self.assertEqual(columns2[i], name)

982 # Read the rowcount.

983 rowcount = self.butler.get(self.datasetType.componentTypeName("rowcount"), dataId={})

984 self.assertEqual(rowcount, len(tab1))

985 # Read the schema.

986 schema = self.butler.get(self.datasetType.componentTypeName("schema"), dataId={})

987 self.assertEqual(schema, ArrowAstropySchema(tab1))

988 # Read just some columns a few different ways.

989 tab3 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["a", "c"]})

990 _checkAstropyTableEquality(tab1[("a", "c")], tab3)

991 tab4 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": "a"})

992 _checkAstropyTableEquality(tab1[("a",)], tab4)

993 tab5 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["index", "a"]})

994 _checkAstropyTableEquality(tab1[("index", "a")], tab5)

995 tab6 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": "ddd"})

996 _checkAstropyTableEquality(tab1[("ddd",)], tab6)

997 tab7 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["a", "a"]})

998 _checkAstropyTableEquality(tab1[("a",)], tab7)

999 tab8 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["d??"]})

1000 _checkAstropyTableEquality(tab1[("ddd", "dtn", "dtu")], tab8)

1001 tab9 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["d??", "a*"]})

1002 _checkAstropyTableEquality(tab1[("ddd", "dtn", "dtu", "a")], tab9)

1003 # Passing an unrecognized column should be a ValueError.

1004 with self.assertRaises(ValueError):

1005 self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["e"]})

1006

1007 def testAstropyTableBigEndian(self):

1008 tab1 = _makeSimpleAstropyTable(include_bigendian=True)

1009

1010 self.butler.put(tab1, self.datasetType, dataId={})

1011 # Read the whole Table.

1012 tab2 = self.butler.get(self.datasetType, dataId={})

1013 _checkAstropyTableEquality(tab1, tab2, has_bigendian=True)

1014

1015 def testAstropyTableWithMetadata(self):

1016 tab1 = _makeSimpleAstropyTable(include_multidim=True)

1017

1018 meta = {

1019 "meta_a": 5,

1020 "meta_b": 10.0,

1021 "meta_c": [1, 2, 3],

1022 "meta_d": True,

1023 "meta_e": "string",

1024 }

1025

1026 tab1.meta.update(meta)

1027

1028 self.butler.put(tab1, self.datasetType, dataId={})

1029 # Read the whole Table.

1030 tab2 = self.butler.get(self.datasetType, dataId={})

1031 # This will check that the metadata is equivalent as well.

1032 _checkAstropyTableEquality(tab1, tab2)

1033

1034 def testArrowAstropySchema(self):

1035 tab1 = _makeSimpleAstropyTable()

1036 tab1_arrow = astropy_to_arrow(tab1)

1037 schema = ArrowAstropySchema.from_arrow(tab1_arrow.schema)

1038

1039 self.assertIsInstance(schema.schema, atable.Table)

1040 self.assertEqual(repr(schema), repr(schema._schema))

1041 self.assertNotEqual(schema, "not_a_schema")

1042 self.assertEqual(schema, schema)

1043

1044 # Test various inequalities

1045 tab2 = tab1.copy()

1046 tab2.rename_column("index", "index2")

1047 schema2 = ArrowAstropySchema(tab2)

1048 self.assertNotEqual(schema2, schema)

1049

1050 tab2 = tab1.copy()

1051 tab2["index"].unit = units.micron

1052 schema2 = ArrowAstropySchema(tab2)

1053 self.assertNotEqual(schema2, schema)

1054

1055 tab2 = tab1.copy()

1056 tab2["index"].description = "Index column"

1057 schema2 = ArrowAstropySchema(tab2)

1058 self.assertNotEqual(schema2, schema)

1059

1060 tab2 = tab1.copy()

1061 tab2["index"].format = "%05d"

1062 schema2 = ArrowAstropySchema(tab2)

1063 self.assertNotEqual(schema2, schema)

1064

1065 def testAstropyParquet(self):

1066 tab1 = _makeSimpleAstropyTable()

1067

1068 # Remove datetime column which doesn't work with astropy currently.

1069 del tab1["dtn"]

1070 del tab1["dtu"]

1071

1072 fname = os.path.join(self.root, "test_astropy.parq")

1073 tab1.write(fname)

1074

1075 astropy_type = DatasetType(

1076 "astropy_parquet",

1077 dimensions=(),

1078 storageClass="ArrowAstropy",

1079 universe=self.butler.dimensions,

1080 )

1081 self.butler.registry.registerDatasetType(astropy_type)

1082

1083 data_id = {}

1084 ref = DatasetRef(astropy_type, data_id, run=self.run)

1085 dataset = FileDataset(path=fname, refs=[ref], formatter=ParquetFormatter)

1086

1087 self.butler.ingest(dataset, transfer="copy")

1088

1089 self.butler.put(tab1, self.datasetType, dataId={})

1090

1091 tab2a = self.butler.get(self.datasetType, dataId={})

1092 tab2b = self.butler.get("astropy_parquet", dataId={})

1093 _checkAstropyTableEquality(tab2a, tab2b)

1094

1095 columns2a = self.butler.get(self.datasetType.componentTypeName("columns"), dataId={})

1096 columns2b = self.butler.get("astropy_parquet.columns", dataId={})

1097 self.assertEqual(len(columns2b), len(columns2a))

1098 for i, name in enumerate(columns2a):

1099 self.assertEqual(columns2b[i], name)

1100

1101 rowcount2a = self.butler.get(self.datasetType.componentTypeName("rowcount"), dataId={})

1102 rowcount2b = self.butler.get("astropy_parquet.rowcount", dataId={})

1103 self.assertEqual(rowcount2a, rowcount2b)

1104

1105 schema2a = self.butler.get(self.datasetType.componentTypeName("schema"), dataId={})

1106 schema2b = self.butler.get("astropy_parquet.schema", dataId={})

1107 self.assertEqual(schema2a, schema2b)

1108

1109 @unittest.skipUnless(pa is not None, "Cannot test reading as arrow without pyarrow.")

1110 def testWriteAstropyReadAsArrowTable(self):

1111 # This astropy <-> arrow works fine with masked columns.

1112 tab1 = _makeSimpleAstropyTable(include_masked=True)

1113

1114 self.butler.put(tab1, self.datasetType, dataId={})

1115

1116 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowTable")

1117

1118 tab2_astropy = arrow_to_astropy(tab2)

1119 _checkAstropyTableEquality(tab1, tab2_astropy)

1120

1121 # Check reading the columns.

1122 columns = tab2.schema.names

1123 columns2 = self.butler.get(

1124 self.datasetType.componentTypeName("columns"), dataId={}, storageClass="ArrowColumnList"

1125 )

1126 self.assertEqual(columns2, columns)

1127

1128 # Check reading the schema.

1129 schema = tab2.schema

1130 schema2 = self.butler.get(

1131 self.datasetType.componentTypeName("schema"), dataId={}, storageClass="ArrowSchema"

1132 )

1133

1134 self.assertEqual(schema, schema2)

1135

1136 @unittest.skipUnless(pd is not None, "Cannot test reading as a dataframe without pandas.")

1137 def testWriteAstropyReadAsDataFrame(self):

1138 tab1 = _makeSimpleAstropyTable()

1139

1140 self.butler.put(tab1, self.datasetType, dataId={})

1141

1142 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="DataFrame")

1143

1144 # This is tricky because it loses the units and gains a bonus pandas

1145 # _index_ column, so we just test the dataframe form.

1146

1147 tab1_df = tab1.to_pandas()

1148 self.assertTrue(tab1_df.equals(tab2))

1149

1150 # Check reading the columns.

1151 columns = tab2.columns

1152 columns2 = self.butler.get(

1153 self.datasetType.componentTypeName("columns"), dataId={}, storageClass="DataFrameIndex"

1154 )

1155 self.assertTrue(columns.equals(columns2))

1156

1157 # Check reading the schema.

1158 schema = DataFrameSchema(tab2)

1159 schema2 = self.butler.get(

1160 self.datasetType.componentTypeName("schema"), dataId={}, storageClass="DataFrameSchema"

1161 )

1162

1163 self.assertEqual(schema2, schema)

1164

1165 @unittest.skipUnless(pd is not None, "Cannot test reading as a dataframe without pandas.")

1166 def testWriteAstropyWithMaskedColsReadAsDataFrame(self):

1167 # We need to special-case the write-as-astropy read-as-pandas code

1168 # with masks because pandas has multiple ways to use masked columns.

1169 # (When writing an astropy table with masked columns we get an object

1170 # column back, but each unmasked element has the correct type.)

1171 tab1 = _makeSimpleAstropyTable(include_masked=True)

1172

1173 self.butler.put(tab1, self.datasetType, dataId={})

1174

1175 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="DataFrame")

1176

1177 tab1_df = astropy_to_pandas(tab1)

1178

1179 self.assertTrue(tab1_df.columns.equals(tab2.columns))

1180 for name in tab2.columns:

1181 col1 = tab1_df[name]

1182 col2 = tab2[name]

1183

1184 if col1.hasnans:

1185 notNull = col1.notnull()

1186 self.assertTrue(notNull.equals(col2.notnull()))

1187 # Need to check value-by-value because column may

1188 # be made of objects, depending on what pandas decides.

1189 for index in notNull.values.nonzero()[0]:

1190 self.assertEqual(col1[index], col2[index])

1191 else:

1192 self.assertTrue(col1.equals(col2))

1193

1194 @unittest.skipUnless(pd is not None, "Cannot test writing as a dataframe without pandas.")

1195 def testWriteSingleIndexDataFrameWithMaskedColsReadAsAstropyTable(self):

1196 df1, allColumns = _makeSingleIndexDataFrame(include_masked=True)

1197

1198 self.butler.put(df1, self.datasetType, dataId={})

1199

1200 tab2 = self.butler.get(self.datasetType, dataId={})

1201

1202 df1_tab = pandas_to_astropy(df1)

1203

1204 _checkAstropyTableEquality(df1_tab, tab2)

1205

1206 @unittest.skipUnless(np is not None, "Cannot test reading as numpy without numpy.")

1207 def testWriteAstropyReadAsNumpyTable(self):

1208 tab1 = _makeSimpleAstropyTable()

1209 self.butler.put(tab1, self.datasetType, dataId={})

1210

1211 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowNumpy")

1212

1213 # This is tricky because it loses the units.

1214 tab2_astropy = atable.Table(tab2)

1215

1216 _checkAstropyTableEquality(tab1, tab2_astropy, skip_units=True)

1217

1218 # Check reading the columns.

1219 columns = list(tab2.dtype.names)

1220 columns2 = self.butler.get(

1221 self.datasetType.componentTypeName("columns"), dataId={}, storageClass="ArrowColumnList"

1222 )

1223 self.assertEqual(columns2, columns)

1224

1225 # Check reading the schema.

1226 schema = ArrowNumpySchema(tab2.dtype)

1227 schema2 = self.butler.get(

1228 self.datasetType.componentTypeName("schema"), dataId={}, storageClass="ArrowNumpySchema"

1229 )

1230

1231 self.assertEqual(schema2, schema)

1232

1233 @unittest.skipUnless(np is not None, "Cannot test reading as numpy without numpy.")

1234 def testWriteAstropyReadAsNumpyDict(self):

1235 tab1 = _makeSimpleAstropyTable()

1236 self.butler.put(tab1, self.datasetType, dataId={})

1237

1238 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowNumpyDict")

1239

1240 # This is tricky because it loses the units.

1241 tab2_astropy = atable.Table(tab2)

1242

1243 _checkAstropyTableEquality(tab1, tab2_astropy, skip_units=True)

1244

1245 def testBadAstropyColumnParquet(self):

1246 tab1 = _makeSimpleAstropyTable()

1247

1248 # Make a column with mixed type.

1249 bad_col1 = [0.0] * len(tab1)

1250 bad_col1[1] = 0.0 * units.nJy

1251 bad_tab = tab1.copy()

1252 bad_tab["bad_col1"] = bad_col1

1253

1254 # At the moment we cannot check that the correct note is added

1255 # to the exception, but that will be possible in the future.

1256 with self.assertRaises(RuntimeError):

1257 self.butler.put(bad_tab, self.datasetType, dataId={})

1258

1259 # Make a column with ragged size.

1260 bad_col2 = [[0]] * len(tab1)

1261 bad_col2[1] = [0, 0]

1262 bad_tab = tab1.copy()

1263 bad_tab["bad_col2"] = bad_col2

1264

1265 with self.assertRaises(RuntimeError):

1266 self.butler.put(bad_tab, self.datasetType, dataId={})

1267

1268 @unittest.skipUnless(pd is not None, "Cannot test ParquetFormatterDataFrame without pandas.")

1269 def testWriteAstropyTableWithPandasIndexHint(self, testStrip=True):

1270 tab1 = _makeSimpleAstropyTable()

1271

1272 add_pandas_index_to_astropy(tab1, "index")

1273

1274 self.butler.put(tab1, self.datasetType, dataId={})

1275

1276 # Read in as an astropy table and ensure index hint is still there.

1277 tab2 = self.butler.get(self.datasetType, dataId={})

1278

1279 self.assertIn(ASTROPY_PANDAS_INDEX_KEY, tab2.meta)

1280 self.assertEqual(tab2.meta[ASTROPY_PANDAS_INDEX_KEY], "index")

1281

1282 # Read as a dataframe and ensure index is set.

1283 df3 = self.butler.get(self.datasetType, dataId={}, storageClass="DataFrame")

1284

1285 self.assertEqual(df3.index.name, "index")

1286

1287 # Read as a dataframe without naming the index column.

1288 with self.assertLogs(level="WARNING") as cm:

1289 _ = self.butler.get(

1290 self.datasetType,

1291 dataId={},

1292 storageClass="DataFrame",

1293 parameters={"columns": ["a", "b"]},

1294 )

1295 self.assertIn("Index column ``index``", cm.output[0])

1296

1297 if testStrip:

1298 # Read as an astropy table without naming the index column.

1299 tab5 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["a", "b"]})

1300

1301 self.assertNotIn(ASTROPY_PANDAS_INDEX_KEY, tab5.meta)

1302

1303 with self.assertRaises(ValueError):

1304 add_pandas_index_to_astropy(tab1, "not_a_column")

1305

1306

1307@unittest.skipUnless(atable is not None, "Cannot test InMemoryDatastore with AstropyTable without astropy.")

1308class InMemoryArrowAstropyDelegateTestCase(ParquetFormatterArrowAstropyTestCase):

1309 """Tests for InMemoryDatastore, using ArrowTableDelegate with

1310 AstropyTable.

1311 """

1312

1313 configFile = os.path.join(TESTDIR, "config/basic/butler-inmemory.yaml")

1314

1315 def testAstropyParquet(self):

1316 # This test does not work with an inMemoryDatastore.

1317 pass

1318

1319 def testBadAstropyColumnParquet(self):

1320 # This test does not raise for an in-memory datastore.

1321 pass

1322

1323 def testBadInput(self):

1324 tab1 = _makeSimpleAstropyTable()

1325 delegate = ArrowTableDelegate("ArrowAstropy")

1326

1327 with self.assertRaises(ValueError):

1328 delegate.handleParameters(inMemoryDataset="not_an_astropy_table")

1329

1330 with self.assertRaises(NotImplementedError):

1331 delegate.handleParameters(inMemoryDataset=tab1, parameters={"columns": [("a", "b")]})

1332

1333 with self.assertRaises(AttributeError):

1334 delegate.getComponent(composite=tab1, componentName="nothing")

1335

1336 @unittest.skipUnless(pd is not None, "Cannot test ParquetFormatterDataFrame without pandas.")

1337 def testWriteAstropyTableWithPandasIndexHint(self):

1338 super().testWriteAstropyTableWithPandasIndexHint(testStrip=False)

1339

1340

1341@unittest.skipUnless(np is not None, "Cannot test ParquetFormatterArrowNumpy without numpy.")

1342@unittest.skipUnless(pa is not None, "Cannot test ParquetFormatterArrowNumpy without pyarrow.")

1343class ParquetFormatterArrowNumpyTestCase(unittest.TestCase):

1344 """Tests for ParquetFormatter, ArrowNumpy, using local file datastore."""

1345

1346 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")

1347

1348 def setUp(self):

1349 """Create a new butler root for each test."""

1350 self.root = makeTestTempDir(TESTDIR)

1351 config = Config(self.configFile)

1352 self.butler = Butler.from_config(

1353 Butler.makeRepo(self.root, config=config), writeable=True, run="test_run"

1354 )

1355 # No dimensions in dataset type so we don't have to worry about

1356 # inserting dimension data or defining data IDs.

1357 self.datasetType = DatasetType(

1358 "data", dimensions=(), storageClass="ArrowNumpy", universe=self.butler.dimensions

1359 )

1360 self.butler.registry.registerDatasetType(self.datasetType)

1361

1362 def tearDown(self):

1363 removeTestTempDir(self.root)

1364

1365 def testNumpyTable(self):

1366 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1367

1368 self.butler.put(tab1, self.datasetType, dataId={})

1369 # Read the whole Table.

1370 tab2 = self.butler.get(self.datasetType, dataId={})

1371 _checkNumpyTableEquality(tab1, tab2)

1372 # Read the columns.

1373 columns2 = self.butler.get(self.datasetType.componentTypeName("columns"), dataId={})

1374 self.assertEqual(len(columns2), len(tab1.dtype.names))

1375 for i, name in enumerate(tab1.dtype.names):

1376 self.assertEqual(columns2[i], name)

1377 # Read the rowcount.

1378 rowcount = self.butler.get(self.datasetType.componentTypeName("rowcount"), dataId={})

1379 self.assertEqual(rowcount, len(tab1))

1380 # Read the schema.

1381 schema = self.butler.get(self.datasetType.componentTypeName("schema"), dataId={})

1382 self.assertEqual(schema, ArrowNumpySchema(tab1.dtype))

1383 # Read just some columns a few different ways.

1384 tab3 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["a", "c"]})

1385 _checkNumpyTableEquality(tab1[["a", "c"]], tab3)

1386 tab4 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": "a"})

1387 _checkNumpyTableEquality(

1388 tab1[

1389 [

1390 "a",

1391 ]

1392 ],

1393 tab4,

1394 )

1395 tab5 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["index", "a"]})

1396 _checkNumpyTableEquality(tab1[["index", "a"]], tab5)

1397 tab6 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": "ddd"})

1398 _checkNumpyTableEquality(

1399 tab1[

1400 [

1401 "ddd",

1402 ]

1403 ],

1404 tab6,

1405 )

1406 tab7 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["a", "a"]})

1407 _checkNumpyTableEquality(

1408 tab1[

1409 [

1410 "a",

1411 ]

1412 ],

1413 tab7,

1414 )

1415 tab8 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["d??", "a*"]})

1416 _checkNumpyTableEquality(

1417 tab1[

1418 [

1419 "ddd",

1420 "dtn",

1421 "dtu",

1422 "a",

1423 ]

1424 ],

1425 tab8,

1426 )

1427 # Passing an unrecognized column should be a ValueError.

1428 with self.assertRaises(ValueError):

1429 self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["e"]})

1430

1431 def testNumpyTableBigEndian(self):

1432 tab1 = _makeSimpleNumpyTable(include_bigendian=True)

1433

1434 self.butler.put(tab1, self.datasetType, dataId={})

1435 # Read the whole Table.

1436 tab2 = self.butler.get(self.datasetType, dataId={})

1437 _checkNumpyTableEquality(tab1, tab2, has_bigendian=True)

1438

1439 def testArrowNumpySchema(self):

1440 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1441 tab1_arrow = numpy_to_arrow(tab1)

1442 schema = ArrowNumpySchema.from_arrow(tab1_arrow.schema)

1443

1444 self.assertIsInstance(schema.schema, np.dtype)

1445 self.assertEqual(repr(schema), repr(schema._dtype))

1446 self.assertNotEqual(schema, "not_a_schema")

1447 self.assertEqual(schema, schema)

1448

1449 # Test inequality

1450 tab2 = tab1.copy()

1451 names = list(tab2.dtype.names)

1452 names[0] = "index2"

1453 tab2.dtype.names = names

1454 schema2 = ArrowNumpySchema(tab2.dtype)

1455 self.assertNotEqual(schema2, schema)

1456

1457 @unittest.skipUnless(pa is not None, "Cannot test arrow conversions without pyarrow.")

1458 def testNumpyDictConversions(self):

1459 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1460

1461 # Verify that everything round-trips, including the schema.

1462 tab1_arrow = numpy_to_arrow(tab1)

1463 tab1_dict = arrow_to_numpy_dict(tab1_arrow)

1464 tab1_dict_arrow = numpy_dict_to_arrow(tab1_dict)

1465

1466 self.assertEqual(tab1_arrow.schema, tab1_dict_arrow.schema)

1467 self.assertEqual(tab1_arrow, tab1_dict_arrow)

1468

1469 @unittest.skipUnless(pa is not None, "Cannot test reading as arrow without pyarrow.")

1470 def testWriteNumpyTableReadAsArrowTable(self):

1471 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1472

1473 self.butler.put(tab1, self.datasetType, dataId={})

1474

1475 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowTable")

1476

1477 tab2_numpy = arrow_to_numpy(tab2)

1478

1479 _checkNumpyTableEquality(tab1, tab2_numpy)

1480

1481 # Check reading the columns.

1482 columns = tab2.schema.names

1483 columns2 = self.butler.get(

1484 self.datasetType.componentTypeName("columns"), dataId={}, storageClass="ArrowColumnList"

1485 )

1486 self.assertEqual(columns2, columns)

1487

1488 # Check reading the schema.

1489 schema = tab2.schema

1490 schema2 = self.butler.get(

1491 self.datasetType.componentTypeName("schema"), dataId={}, storageClass="ArrowSchema"

1492 )

1493 self.assertEqual(schema2, schema)

1494

1495 @unittest.skipUnless(pd is not None, "Cannot test reading as a dataframe without pandas.")

1496 def testWriteNumpyTableReadAsDataFrame(self):

1497 tab1 = _makeSimpleNumpyTable()

1498

1499 self.butler.put(tab1, self.datasetType, dataId={})

1500

1501 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="DataFrame")

1502

1503 # Converting this back to numpy gets confused with the index column

1504 # and changes the datatype of the string column.

1505

1506 tab1_df = pd.DataFrame(tab1)

1507

1508 self.assertTrue(tab1_df.equals(tab2))

1509

1510 # Check reading the columns.

1511 columns = tab2.columns

1512 columns2 = self.butler.get(

1513 self.datasetType.componentTypeName("columns"), dataId={}, storageClass="DataFrameIndex"

1514 )

1515 self.assertTrue(columns.equals(columns2))

1516

1517 # Check reading the schema.

1518 schema = DataFrameSchema(tab2)

1519 schema2 = self.butler.get(

1520 self.datasetType.componentTypeName("schema"), dataId={}, storageClass="DataFrameSchema"

1521 )

1522

1523 self.assertEqual(schema2, schema)

1524

1525 @unittest.skipUnless(atable is not None, "Cannot test reading as astropy without astropy.")

1526 def testWriteNumpyTableReadAsAstropyTable(self):

1527 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1528

1529 self.butler.put(tab1, self.datasetType, dataId={})

1530

1531 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowAstropy")

1532 tab2_numpy = tab2.as_array()

1533

1534 _checkNumpyTableEquality(tab1, tab2_numpy)

1535

1536 # Check reading the columns.

1537 columns = list(tab2.columns.keys())

1538 columns2 = self.butler.get(

1539 self.datasetType.componentTypeName("columns"), dataId={}, storageClass="ArrowColumnList"

1540 )

1541 self.assertEqual(columns2, columns)

1542

1543 # Check reading the schema.

1544 schema = ArrowAstropySchema(tab2)

1545 schema2 = self.butler.get(

1546 self.datasetType.componentTypeName("schema"), dataId={}, storageClass="ArrowAstropySchema"

1547 )

1548

1549 self.assertEqual(schema2, schema)

1550

1551 def testWriteNumpyTableReadAsNumpyDict(self):

1552 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1553

1554 self.butler.put(tab1, self.datasetType, dataId={})

1555

1556 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowNumpyDict")

1557 tab2_numpy = _numpy_dict_to_numpy(tab2)

1558

1559 _checkNumpyTableEquality(tab1, tab2_numpy)

1560

1561 def testBadNumpyColumnParquet(self):

1562 tab1 = _makeSimpleAstropyTable()

1563

1564 # Make a column with mixed type.

1565 bad_col1 = [0.0] * len(tab1)

1566 bad_col1[1] = 0.0 * units.nJy

1567 bad_tab = tab1.copy()

1568 bad_tab["bad_col1"] = bad_col1

1569

1570 bad_tab_np = bad_tab.as_array()

1571

1572 # At the moment we cannot check that the correct note is added

1573 # to the exception, but that will be possible in the future.

1574 with self.assertRaises(RuntimeError):

1575 self.butler.put(bad_tab_np, self.datasetType, dataId={})

1576

1577 # Make a column with ragged size.

1578 bad_col2 = [[0]] * len(tab1)

1579 bad_col2[1] = [0, 0]

1580 bad_tab = tab1.copy()

1581 bad_tab["bad_col2"] = bad_col2

1582

1583 bad_tab_np = bad_tab.as_array()

1584

1585 with self.assertRaises(RuntimeError):

1586 self.butler.put(bad_tab_np, self.datasetType, dataId={})

1587

1588 @unittest.skipUnless(atable is not None, "Cannot test reading as astropy without astropy.")

1589 def testWriteReadAstropyTableLossless(self):

1590 tab1 = _makeSimpleAstropyTable(include_multidim=True, include_masked=True)

1591

1592 self.butler.put(tab1, self.datasetType, dataId={})

1593

1594 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowAstropy")

1595

1596 _checkAstropyTableEquality(tab1, tab2)

1597

1598

1599@unittest.skipUnless(np is not None, "Cannot test ImMemoryDatastore with Numpy table without numpy.")

1600class InMemoryArrowNumpyDelegateTestCase(ParquetFormatterArrowNumpyTestCase):

1601 """Tests for InMemoryDatastore, using ArrowTableDelegate with

1602 Numpy table.

1603 """

1604

1605 configFile = os.path.join(TESTDIR, "config/basic/butler-inmemory.yaml")

1606

1607 def testBadNumpyColumnParquet(self):

1608 # This test does not raise for an in-memory datastore.

1609 pass

1610

1611 def testBadInput(self):

1612 tab1 = _makeSimpleNumpyTable()

1613 delegate = ArrowTableDelegate("ArrowNumpy")

1614

1615 with self.assertRaises(ValueError):

1616 delegate.handleParameters(inMemoryDataset="not_a_numpy_table")

1617

1618 with self.assertRaises(NotImplementedError):

1619 delegate.handleParameters(inMemoryDataset=tab1, parameters={"columns": [("a", "b")]})

1620

1621 with self.assertRaises(AttributeError):

1622 delegate.getComponent(composite=tab1, componentName="nothing")

1623

1624 def testStorageClass(self):

1625 tab1 = _makeSimpleNumpyTable()

1626

1627 factory = StorageClassFactory()

1628 factory.addFromConfig(StorageClassConfig())

1629

1630 storageClass = factory.findStorageClass(type(tab1), compare_types=False)

1631 # Force the name lookup to do name matching.

1632 storageClass._pytype = None

1633 self.assertEqual(storageClass.name, "ArrowNumpy")

1634

1635 storageClass = factory.findStorageClass(type(tab1), compare_types=True)

1636 # Force the name lookup to do name matching.

1637 storageClass._pytype = None

1638 self.assertEqual(storageClass.name, "ArrowNumpy")

1639

1640

1641@unittest.skipUnless(pa is not None, "Cannot test ParquetFormatterArrowTable without pyarrow.")

1642class ParquetFormatterArrowTableTestCase(unittest.TestCase):

1643 """Tests for ParquetFormatter, ArrowTable, using local file datastore."""

1644

1645 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")

1646

1647 def setUp(self):

1648 """Create a new butler root for each test."""

1649 self.root = makeTestTempDir(TESTDIR)

1650 config = Config(self.configFile)

1651 self.butler = Butler.from_config(

1652 Butler.makeRepo(self.root, config=config), writeable=True, run="test_run"

1653 )

1654 # No dimensions in dataset type so we don't have to worry about

1655 # inserting dimension data or defining data IDs.

1656 self.datasetType = DatasetType(

1657 "data", dimensions=(), storageClass="ArrowTable", universe=self.butler.dimensions

1658 )

1659 self.butler.registry.registerDatasetType(self.datasetType)

1660

1661 def tearDown(self):

1662 removeTestTempDir(self.root)

1663

1664 def testArrowTable(self):

1665 tab1 = _makeSimpleArrowTable(include_multidim=True, include_masked=True)

1666

1667 self.butler.put(tab1, self.datasetType, dataId={})

1668 # Read the whole Table.

1669 tab2 = self.butler.get(self.datasetType, dataId={})

1670 # We convert to use the numpy testing framework to handle nan

1671 # comparisons.

1672 self.assertEqual(tab1.schema, tab2.schema)

1673 tab1_np = arrow_to_numpy(tab1)

1674 tab2_np = arrow_to_numpy(tab2)

1675 for col in tab1.column_names:

1676 np.testing.assert_array_equal(tab2_np[col], tab1_np[col])

1677 # Read the columns.

1678 columns2 = self.butler.get(self.datasetType.componentTypeName("columns"), dataId={})

1679 self.assertEqual(len(columns2), len(tab1.schema.names))

1680 for i, name in enumerate(tab1.schema.names):

1681 self.assertEqual(columns2[i], name)

1682 # Read the rowcount.

1683 rowcount = self.butler.get(self.datasetType.componentTypeName("rowcount"), dataId={})

1684 self.assertEqual(rowcount, len(tab1))

1685 # Read the schema.

1686 schema = self.butler.get(self.datasetType.componentTypeName("schema"), dataId={})

1687 self.assertEqual(schema, tab1.schema)

1688 # Read just some columns a few different ways.

1689 tab3 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["a", "c"]})

1690 self.assertEqual(tab3, tab1.select(("a", "c")))

1691 tab4 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": "a"})

1692 self.assertEqual(tab4, tab1.select(("a",)))

1693 tab5 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["index", "a"]})

1694 self.assertEqual(tab5, tab1.select(("index", "a")))

1695 tab6 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": "ddd"})

1696 self.assertEqual(tab6, tab1.select(("ddd",)))

1697 tab7 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["a", "a"]})

1698 self.assertEqual(tab7, tab1.select(("a",)))

1699 tab8 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["a*", "d??"]})

1700 self.assertEqual(tab8, tab1.select(("a", "ddd", "dtn", "dtu")))

1701 # Passing an unrecognized column should be a ValueError.

1702 with self.assertRaises(ValueError):

1703 self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["e"]})

1704

1705 def testEmptyArrowTable(self):

1706 data = _makeSimpleNumpyTable()

1707 type_list = _numpy_dtype_to_arrow_types(data.dtype)

1708

1709 schema = pa.schema(type_list)

1710 arrays = [[]] * len(schema.names)

1711

1712 tab1 = pa.Table.from_arrays(arrays, schema=schema)

1713

1714 self.butler.put(tab1, self.datasetType, dataId={})

1715 tab2 = self.butler.get(self.datasetType, dataId={})

1716 self.assertEqual(tab2, tab1)

1717

1718 tab1_numpy = arrow_to_numpy(tab1)

1719 self.assertEqual(len(tab1_numpy), 0)

1720 tab1_numpy_arrow = numpy_to_arrow(tab1_numpy)

1721 self.assertEqual(tab1_numpy_arrow, tab1)

1722

1723 tab1_pandas = arrow_to_pandas(tab1)

1724 self.assertEqual(len(tab1_pandas), 0)

1725 tab1_pandas_arrow = pandas_to_arrow(tab1_pandas)

1726 # Unfortunately, string/byte columns get mangled when translated

1727 # through empty pandas dataframes.

1728 self.assertEqual(

1729 tab1_pandas_arrow.select(("index", "a", "b", "c", "ddd")),

1730 tab1.select(("index", "a", "b", "c", "ddd")),

1731 )

1732

1733 tab1_astropy = arrow_to_astropy(tab1)

1734 self.assertEqual(len(tab1_astropy), 0)

1735 tab1_astropy_arrow = astropy_to_arrow(tab1_astropy)

1736 self.assertEqual(tab1_astropy_arrow, tab1)

1737

1738 def testEmptyArrowTableMultidim(self):

1739 data = _makeSimpleNumpyTable(include_multidim=True)

1740 type_list = _numpy_dtype_to_arrow_types(data.dtype)

1741

1742 md = {}

1743 for name in data.dtype.names:

1744 _append_numpy_multidim_metadata(md, name, data.dtype[name])

1745

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

1747 arrays = [[]] * len(schema.names)

1748

1749 tab1 = pa.Table.from_arrays(arrays, schema=schema)

1750

1751 self.butler.put(tab1, self.datasetType, dataId={})

1752 tab2 = self.butler.get(self.datasetType, dataId={})

1753 self.assertEqual(tab2, tab1)

1754

1755 tab1_numpy = arrow_to_numpy(tab1)

1756 self.assertEqual(len(tab1_numpy), 0)

1757 tab1_numpy_arrow = numpy_to_arrow(tab1_numpy)

1758 self.assertEqual(tab1_numpy_arrow, tab1)

1759

1760 tab1_astropy = arrow_to_astropy(tab1)

1761 self.assertEqual(len(tab1_astropy), 0)

1762 tab1_astropy_arrow = astropy_to_arrow(tab1_astropy)

1763 self.assertEqual(tab1_astropy_arrow, tab1)

1764

1765 @unittest.skipUnless(pd is not None, "Cannot test reading as a dataframe without pandas.")

1766 def testWriteArrowTableReadAsSingleIndexDataFrame(self):

1767 df1, allColumns = _makeSingleIndexDataFrame()

1768

1769 self.butler.put(df1, self.datasetType, dataId={})

1770

1771 # Read back out as a dataframe.

1772 df2 = self.butler.get(self.datasetType, dataId={}, storageClass="DataFrame")

1773 self.assertTrue(df1.equals(df2))

1774

1775 # Read back out as an arrow table, convert to dataframe.

1776 tab3 = self.butler.get(self.datasetType, dataId={})

1777 df3 = arrow_to_pandas(tab3)

1778 self.assertTrue(df1.equals(df3))

1779

1780 # Check reading the columns.

1781 columns = df2.reset_index().columns

1782 columns2 = self.butler.get(

1783 self.datasetType.componentTypeName("columns"), dataId={}, storageClass="DataFrameIndex"

1784 )

1785 # We check the set because pandas reorders the columns.

1786 self.assertEqual(set(columns2.to_list()), set(columns.to_list()))

1787

1788 # Check reading the schema.

1789 schema = DataFrameSchema(df1)

1790 schema2 = self.butler.get(

1791 self.datasetType.componentTypeName("schema"), dataId={}, storageClass="DataFrameSchema"

1792 )

1793 self.assertEqual(schema2, schema)

1794

1795 @unittest.skipUnless(pd is not None, "Cannot test reading as a dataframe without pandas.")

1796 def testWriteArrowTableReadAsMultiIndexDataFrame(self):

1797 df1 = _makeMultiIndexDataFrame()

1798

1799 self.butler.put(df1, self.datasetType, dataId={})

1800

1801 # Read back out as a dataframe.

1802 df2 = self.butler.get(self.datasetType, dataId={}, storageClass="DataFrame")

1803 self.assertTrue(df1.equals(df2))

1804

1805 # Read back out as an arrow table, convert to dataframe.

1806 atab3 = self.butler.get(self.datasetType, dataId={})

1807 df3 = arrow_to_pandas(atab3)

1808 self.assertTrue(df1.equals(df3))

1809

1810 # Check reading the columns.

1811 columns = df2.columns

1812 columns2 = self.butler.get(

1813 self.datasetType.componentTypeName("columns"), dataId={}, storageClass="DataFrameIndex"

1814 )

1815 self.assertTrue(columns2.equals(columns))

1816

1817 # Check reading the schema.

1818 schema = DataFrameSchema(df1)

1819 schema2 = self.butler.get(

1820 self.datasetType.componentTypeName("schema"), dataId={}, storageClass="DataFrameSchema"

1821 )

1822 self.assertEqual(schema2, schema)

1823

1824 @unittest.skipUnless(atable is not None, "Cannot test reading as astropy without astropy.")

1825 def testWriteArrowTableReadAsAstropyTable(self):

1826 tab1 = _makeSimpleAstropyTable(include_multidim=True, include_masked=True)

1827

1828 self.butler.put(tab1, self.datasetType, dataId={})

1829

1830 # Read back out as an astropy table.

1831 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowAstropy")

1832 _checkAstropyTableEquality(tab1, tab2)

1833

1834 # Read back out as an arrow table, convert to astropy table.

1835 atab3 = self.butler.get(self.datasetType, dataId={})

1836 tab3 = arrow_to_astropy(atab3)

1837 _checkAstropyTableEquality(tab1, tab3)

1838

1839 # Check reading the columns.

1840 columns = list(tab2.columns.keys())

1841 columns2 = self.butler.get(

1842 self.datasetType.componentTypeName("columns"), dataId={}, storageClass="ArrowColumnList"

1843 )

1844 self.assertEqual(columns2, columns)

1845

1846 # Check reading the schema.

1847 schema = ArrowAstropySchema(tab1)

1848 schema2 = self.butler.get(

1849 self.datasetType.componentTypeName("schema"), dataId={}, storageClass="ArrowAstropySchema"

1850 )

1851 self.assertEqual(schema2, schema)

1852

1853 # Check the schema conversions and units.

1854 arrow_schema = schema.to_arrow_schema()

1855 for name in arrow_schema.names:

1856 field_metadata = arrow_schema.field(name).metadata

1857 if (

1858 b"description" in field_metadata

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

1860 ):

1861 self.assertEqual(schema2.schema[name].description, description)

1862 else:

1863 self.assertIsNone(schema2.schema[name].description)

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

1865 self.assertEqual(schema2.schema[name].unit, units.Unit(unit))

1866

1867 @unittest.skipUnless(np is not None, "Cannot test reading as numpy without numpy.")

1868 def testWriteArrowTableReadAsNumpyTable(self):

1869 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1870

1871 self.butler.put(tab1, self.datasetType, dataId={})

1872

1873 # Read back out as a numpy table.

1874 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowNumpy")

1875 _checkNumpyTableEquality(tab1, tab2)

1876

1877 # Read back out as an arrow table, convert to numpy table.

1878 atab3 = self.butler.get(self.datasetType, dataId={})

1879 tab3 = arrow_to_numpy(atab3)

1880 _checkNumpyTableEquality(tab1, tab3)

1881

1882 # Check reading the columns.

1883 columns = list(tab2.dtype.names)

1884 columns2 = self.butler.get(

1885 self.datasetType.componentTypeName("columns"), dataId={}, storageClass="ArrowColumnList"

1886 )

1887 self.assertEqual(columns2, columns)

1888

1889 # Check reading the schema.

1890 schema = ArrowNumpySchema(tab1.dtype)

1891 schema2 = self.butler.get(

1892 self.datasetType.componentTypeName("schema"), dataId={}, storageClass="ArrowNumpySchema"

1893 )

1894 self.assertEqual(schema2, schema)

1895

1896 @unittest.skipUnless(np is not None, "Cannot test reading as numpy without numpy.")

1897 def testWriteArrowTableReadAsNumpyDict(self):

1898 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1899

1900 self.butler.put(tab1, self.datasetType, dataId={})

1901

1902 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowNumpyDict")

1903 tab2_numpy = _numpy_dict_to_numpy(tab2)

1904 _checkNumpyTableEquality(tab1, tab2_numpy)

1905

1906 @unittest.skipUnless(atable is not None, "Cannot test reading as astropy without astropy.")

1907 def testWriteReadAstropyTableLossless(self):

1908 tab1 = _makeSimpleAstropyTable(include_multidim=True, include_masked=True)

1909

1910 self.butler.put(tab1, self.datasetType, dataId={})

1911

1912 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowAstropy")

1913

1914 _checkAstropyTableEquality(tab1, tab2)

1915

1916

1917@unittest.skipUnless(pa is not None, "Cannot test InMemoryDatastore with ArroWTable without pyarrow.")

1918class InMemoryArrowTableDelegateTestCase(ParquetFormatterArrowTableTestCase):

1919 """Tests for InMemoryDatastore, using ArrowTableDelegate."""

1920

1921 configFile = os.path.join(TESTDIR, "config/basic/butler-inmemory.yaml")

1922

1923 def testBadInput(self):

1924 tab1 = _makeSimpleArrowTable()

1925 delegate = ArrowTableDelegate("ArrowTable")

1926

1927 with self.assertRaises(ValueError):

1928 delegate.handleParameters(inMemoryDataset="not_an_arrow_table")

1929

1930 with self.assertRaises(NotImplementedError):

1931 delegate.handleParameters(inMemoryDataset=tab1, parameters={"columns": [("a", "b")]})

1932

1933 with self.assertRaises(AttributeError):

1934 delegate.getComponent(composite=tab1, componentName="nothing")

1935

1936 def testStorageClass(self):

1937 tab1 = _makeSimpleArrowTable()

1938

1939 factory = StorageClassFactory()

1940 factory.addFromConfig(StorageClassConfig())

1941

1942 storageClass = factory.findStorageClass(type(tab1), compare_types=False)

1943 # Force the name lookup to do name matching.

1944 storageClass._pytype = None

1945 self.assertEqual(storageClass.name, "ArrowTable")

1946

1947 storageClass = factory.findStorageClass(type(tab1), compare_types=True)

1948 # Force the name lookup to do name matching.

1949 storageClass._pytype = None

1950 self.assertEqual(storageClass.name, "ArrowTable")

1951

1952

1953@unittest.skipUnless(np is not None, "Cannot test ParquetFormatterArrowNumpy without numpy.")

1954@unittest.skipUnless(pa is not None, "Cannot test ParquetFormatterArrowNumpy without pyarrow.")

1955class ParquetFormatterArrowNumpyDictTestCase(unittest.TestCase):

1956 """Tests for ParquetFormatter, ArrowNumpyDict, using local file store."""

1957

1958 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")

1959

1960 def setUp(self):

1961 """Create a new butler root for each test."""

1962 self.root = makeTestTempDir(TESTDIR)

1963 config = Config(self.configFile)

1964 self.butler = Butler.from_config(

1965 Butler.makeRepo(self.root, config=config), writeable=True, run="test_run"

1966 )

1967 # No dimensions in dataset type so we don't have to worry about

1968 # inserting dimension data or defining data IDs.

1969 self.datasetType = DatasetType(

1970 "data", dimensions=(), storageClass="ArrowNumpyDict", universe=self.butler.dimensions

1971 )

1972 self.butler.registry.registerDatasetType(self.datasetType)

1973

1974 def tearDown(self):

1975 removeTestTempDir(self.root)

1976

1977 def testNumpyDict(self):

1978 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1979 dict1 = _numpy_to_numpy_dict(tab1)

1980

1981 self.butler.put(dict1, self.datasetType, dataId={})

1982 # Read the whole table.

1983 dict2 = self.butler.get(self.datasetType, dataId={})

1984 _checkNumpyDictEquality(dict1, dict2)

1985 # Read the columns.

1986 columns2 = self.butler.get(self.datasetType.componentTypeName("columns"), dataId={})

1987 self.assertEqual(len(columns2), len(dict1.keys()))

1988 for name in dict1:

1989 self.assertIn(name, columns2)

1990 # Read the rowcount.

1991 rowcount = self.butler.get(self.datasetType.componentTypeName("rowcount"), dataId={})

1992 self.assertEqual(rowcount, len(dict1["a"]))

1993 # Read the schema.

1994 schema = self.butler.get(self.datasetType.componentTypeName("schema"), dataId={})

1995 self.assertEqual(schema, ArrowNumpySchema(tab1.dtype))

1996 # Read just some columns a few different ways.

1997 tab3 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["a", "c"]})

1998 subdict = {key: dict1[key] for key in ["a", "c"]}

1999 _checkNumpyDictEquality(subdict, tab3)

2000 tab4 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": "a"})

2001 subdict = {key: dict1[key] for key in ["a"]}

2002 _checkNumpyDictEquality(subdict, tab4)

2003 tab5 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["index", "a"]})

2004 subdict = {key: dict1[key] for key in ["index", "a"]}

2005 _checkNumpyDictEquality(subdict, tab5)

2006 tab6 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": "ddd"})

2007 subdict = {key: dict1[key] for key in ["ddd"]}

2008 _checkNumpyDictEquality(subdict, tab6)

2009 tab7 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["a", "a"]})

2010 subdict = {key: dict1[key] for key in ["a"]}

2011 _checkNumpyDictEquality(subdict, tab7)

2012 tab8 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["d??", "a*"]})

2013 subdict = {key: dict1[key] for key in ["ddd", "dtn", "dtu", "a"]}

2014 _checkNumpyDictEquality(subdict, tab8)

2015 # Passing an unrecognized column should be a ValueError.

2016 with self.assertRaises(ValueError):

2017 self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["e"]})

2018

2019 @unittest.skipUnless(pa is not None, "Cannot test reading as arrow without pyarrow.")

2020 def testWriteNumpyDictReadAsArrowTable(self):

2021 tab1 = _makeSimpleNumpyTable(include_multidim=True)

2022 dict1 = _numpy_to_numpy_dict(tab1)

2023

2024 self.butler.put(dict1, self.datasetType, dataId={})

2025

2026 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowTable")

2027

2028 tab2_dict = arrow_to_numpy_dict(tab2)

2029

2030 _checkNumpyDictEquality(dict1, tab2_dict)

2031

2032 @unittest.skipUnless(pd is not None, "Cannot test reading as a dataframe without pandas.")

2033 def testWriteNumpyDictReadAsDataFrame(self):

2034 tab1 = _makeSimpleNumpyTable()

2035 dict1 = _numpy_to_numpy_dict(tab1)

2036

2037 self.butler.put(dict1, self.datasetType, dataId={})

2038

2039 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="DataFrame")

2040

2041 # The order of the dict may get mixed up, so we need to check column

2042 # by column. We also need to do this in dataframe form because pandas

2043 # changes the datatype of the string column.

2044 tab1_df = pd.DataFrame(tab1)

2045

2046 self.assertEqual(set(tab1_df.columns), set(tab2.columns))

2047 for col in tab1_df.columns:

2048 self.assertTrue(np.all(tab1_df[col].values == tab2[col].values))

2049

2050 @unittest.skipUnless(atable is not None, "Cannot test reading as astropy without astropy.")

2051 def testWriteNumpyDictReadAsAstropyTable(self):

2052 tab1 = _makeSimpleNumpyTable(include_multidim=True)

2053 dict1 = _numpy_to_numpy_dict(tab1)

2054

2055 self.butler.put(dict1, self.datasetType, dataId={})

2056

2057 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowAstropy")

2058 tab2_dict = _astropy_to_numpy_dict(tab2)

2059

2060 _checkNumpyDictEquality(dict1, tab2_dict)

2061

2062 def testWriteNumpyDictReadAsNumpyTable(self):

2063 tab1 = _makeSimpleNumpyTable(include_multidim=True)

2064 dict1 = _numpy_to_numpy_dict(tab1)

2065

2066 self.butler.put(dict1, self.datasetType, dataId={})

2067

2068 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowNumpy")

2069 tab2_dict = _numpy_to_numpy_dict(tab2)

2070

2071 _checkNumpyDictEquality(dict1, tab2_dict)

2072

2073 def testWriteNumpyDictBad(self):

2074 dict1 = {"a": 4, "b": np.ndarray([1])}

2075 with self.assertRaises(RuntimeError):

2076 self.butler.put(dict1, self.datasetType, dataId={})

2077

2078 dict2 = {"a": np.zeros(4), "b": np.zeros(5)}

2079 with self.assertRaises(RuntimeError):

2080 self.butler.put(dict2, self.datasetType, dataId={})

2081

2082 dict3 = {"a": [0] * 5, "b": np.zeros(5)}

2083 with self.assertRaises(RuntimeError):

2084 self.butler.put(dict3, self.datasetType, dataId={})

2085

2086 dict4 = {"a": np.zeros(4), "b": np.zeros(4, dtype="O")}

2087 with self.assertRaises(RuntimeError):

2088 self.butler.put(dict4, self.datasetType, dataId={})

2089

2090 @unittest.skipUnless(atable is not None, "Cannot test reading as astropy without astropy.")

2091 def testWriteReadAstropyTableLossless(self):

2092 tab1 = _makeSimpleAstropyTable(include_multidim=True, include_masked=True)

2093

2094 self.butler.put(tab1, self.datasetType, dataId={})

2095

2096 tab2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowAstropy")

2097

2098 _checkAstropyTableEquality(tab1, tab2)

2099

2100

2101@unittest.skipUnless(np is not None, "Cannot test InMemoryDatastore with NumpyDict without numpy.")

2102@unittest.skipUnless(pa is not None, "Cannot test InMemoryDatastore with NumpyDict without pyarrow.")

2103class InMemoryNumpyDictDelegateTestCase(ParquetFormatterArrowNumpyDictTestCase):

2104 """Tests for InMemoryDatastore, using ArrowTableDelegate with

2105 Numpy dict.

2106 """

2107

2108 configFile = os.path.join(TESTDIR, "config/basic/butler-inmemory.yaml")

2109

2110 def testWriteNumpyDictBad(self):

2111 # The sub-type checking is not done on in-memory datastore.

2112 pass

2113

2114

2115@unittest.skipUnless(pa is not None, "Cannot test ArrowSchema without pyarrow.")

2116class ParquetFormatterArrowSchemaTestCase(unittest.TestCase):

2117 """Tests for ParquetFormatter, ArrowSchema, using local file datastore."""

2118

2119 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml")

2120

2121 def setUp(self):

2122 """Create a new butler root for each test."""

2123 self.root = makeTestTempDir(TESTDIR)

2124 config = Config(self.configFile)

2125 self.butler = Butler.from_config(

2126 Butler.makeRepo(self.root, config=config), writeable=True, run="test_run"

2127 )

2128 # No dimensions in dataset type so we don't have to worry about

2129 # inserting dimension data or defining data IDs.

2130 self.datasetType = DatasetType(

2131 "data", dimensions=(), storageClass="ArrowSchema", universe=self.butler.dimensions

2132 )

2133 self.butler.registry.registerDatasetType(self.datasetType)

2134

2135 def tearDown(self):

2136 removeTestTempDir(self.root)

2137

2138 def _makeTestSchema(self):

2139 schema = pa.schema(

2140 [

2141 pa.field(

2142 "int32",

2143 pa.int32(),

2144 nullable=False,

2145 metadata={

2146 "description": "32-bit integer",

2147 "unit": "",

2148 },

2149 ),

2150 pa.field(

2151 "int64",

2152 pa.int64(),

2153 nullable=False,

2154 metadata={

2155 "description": "64-bit integer",

2156 "unit": "",

2157 },

2158 ),

2159 pa.field(

2160 "uint64",

2161 pa.uint64(),

2162 nullable=False,

2163 metadata={

2164 "description": "64-bit unsigned integer",

2165 "unit": "",

2166 },

2167 ),

2168 pa.field(

2169 "float32",

2170 pa.float32(),

2171 nullable=False,

2172 metadata={

2173 "description": "32-bit float",

2174 "unit": "count",

2175 },

2176 ),

2177 pa.field(

2178 "float64",

2179 pa.float64(),

2180 nullable=False,

2181 metadata={

2182 "description": "64-bit float",

2183 "unit": "nJy",

2184 },

2185 ),

2186 pa.field(

2187 "fixed_size_list",

2188 pa.list_(pa.float64(), list_size=10),

2189 nullable=False,

2190 metadata={

2191 "description": "Fixed size list of 64-bit floats.",

2192 "unit": "nJy",

2193 },

2194 ),

2195 pa.field(

2196 "variable_size_list",

2197 pa.list_(pa.float64()),

2198 nullable=False,

2199 metadata={

2200 "description": "Variable size list of 64-bit floats.",

2201 "unit": "nJy",

2202 },

2203 ),

2204 # One of these fields will have no description.

2205 pa.field(

2206 "string",

2207 pa.string(),

2208 nullable=False,

2209 metadata={

2210 "unit": "",

2211 },

2212 ),

2213 # One of these fields will have no metadata.

2214 pa.field(

2215 "binary",

2216 pa.binary(),

2217 nullable=False,

2218 ),

2219 ]

2220 )

2221

2222 return schema

2223

2224 def testArrowSchema(self):

2225 schema1 = self._makeTestSchema()

2226 self.butler.put(schema1, self.datasetType, dataId={})

2227

2228 schema2 = self.butler.get(self.datasetType, dataId={})

2229 self.assertEqual(schema2, schema1)

2230

2231 @unittest.skipUnless(pd is not None, "Cannot test reading as a dataframe schema without pandas.")

2232 def testWriteArrowSchemaReadAsDataFrameSchema(self):

2233 schema1 = self._makeTestSchema()

2234 self.butler.put(schema1, self.datasetType, dataId={})

2235

2236 df_schema1 = DataFrameSchema.from_arrow(schema1)

2237

2238 df_schema2 = self.butler.get(self.datasetType, dataId={}, storageClass="DataFrameSchema")

2239 self.assertEqual(df_schema2, df_schema1)

2240

2241 @unittest.skipUnless(atable is not None, "Cannot test reading as an astropy schema without astropy.")

2242 def testWriteArrowSchemaReadAsArrowAstropySchema(self):

2243 schema1 = self._makeTestSchema()

2244 self.butler.put(schema1, self.datasetType, dataId={})

2245

2246 ap_schema1 = ArrowAstropySchema.from_arrow(schema1)

2247

2248 ap_schema2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowAstropySchema")

2249 self.assertEqual(ap_schema2, ap_schema1)

2250

2251 # Confirm that the ap_schema2 has the unit/description we expect.

2252 for name in schema1.names:

2253 field_metadata = schema1.field(name).metadata

2254 if field_metadata is None:

2255 continue

2256 if (

2257 b"description" in field_metadata

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

2259 ):

2260 self.assertEqual(ap_schema2.schema[name].description, description)

2261 else:

2262 self.assertIsNone(ap_schema2.schema[name].description)

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

2264 self.assertEqual(ap_schema2.schema[name].unit, units.Unit(unit))

2265

2266 @unittest.skipUnless(atable is not None, "Cannot test reading as an numpy schema without numpy.")

2267 def testWriteArrowSchemaReadAsArrowNumpySchema(self):

2268 schema1 = self._makeTestSchema()

2269 self.butler.put(schema1, self.datasetType, dataId={})

2270

2271 np_schema1 = ArrowNumpySchema.from_arrow(schema1)

2272

2273 np_schema2 = self.butler.get(self.datasetType, dataId={}, storageClass="ArrowNumpySchema")

2274 self.assertEqual(np_schema2, np_schema1)

2275

2276

2277@unittest.skipUnless(pa is not None, "Cannot test InMemoryDatastore with ArrowSchema without pyarrow.")

2278class InMemoryArrowSchemaDelegateTestCase(ParquetFormatterArrowSchemaTestCase):

2279 """Tests for InMemoryDatastore and ArrowSchema."""

2280

2281 configFile = os.path.join(TESTDIR, "config/basic/butler-inmemory.yaml")

2282

2283

2284@unittest.skipUnless(pa is not None, "Cannot test S3 without pyarrow.")

2285@unittest.skipUnless(boto3 is not None, "Cannot test S3 without boto3.")

2286@unittest.skipUnless(fsspec is not None, "Cannot test S3 without fsspec.")

2287@unittest.skipUnless(s3fs is not None, "Cannot test S3 without s3fs.")

2288class ParquetFormatterArrowTableS3TestCase(unittest.TestCase):

2289 """Tests for arrow table/parquet with S3."""

2290

2291 # Code is adapted from test_butler.py

2292 configFile = os.path.join(TESTDIR, "config/basic/butler-s3store.yaml")

2293 fullConfigKey = None

2294 validationCanFail = True

2295

2296 bucketName = "anybucketname"

2297

2298 root = "butlerRoot/"

2299

2300 datastoreStr = [f"datastore={root}"]

2301

2302 datastoreName = ["FileDatastore@s3://{bucketName}/{root}"]

2303

2304 registryStr = "/gen3.sqlite3"

2305

2306 mock_aws = mock_aws()

2307

2308 def setUp(self):

2309 self.root = makeTestTempDir(TESTDIR)

2310

2311 config = Config(self.configFile)

2312 uri = ResourcePath(config[".datastore.datastore.root"])

2313 self.bucketName = uri.netloc

2314

2315 # Enable S3 mocking of tests.

2316 self.enterContext(clean_test_environment_for_s3())

2317 self.mock_aws.start()

2318

2319 rooturi = f"s3://{self.bucketName}/{self.root}"

2320 config.update({"datastore": {"datastore": {"root": rooturi}}})

2321

2322 # need local folder to store registry database

2323 self.reg_dir = makeTestTempDir(TESTDIR)

2324 config["registry", "db"] = f"sqlite:///{self.reg_dir}/gen3.sqlite3"

2325

2326 # MOTO needs to know that we expect Bucket bucketname to exist

2327 # (this used to be the class attribute bucketName)

2328 s3 = boto3.resource("s3")

2329 s3.create_bucket(Bucket=self.bucketName)

2330

2331 self.datastoreStr = [f"datastore='{rooturi}'"]

2332 self.datastoreName = [f"FileDatastore@{rooturi}"]

2333 Butler.makeRepo(rooturi, config=config, forceConfigRoot=False)

2334 self.tmpConfigFile = posixpath.join(rooturi, "butler.yaml")

2335

2336 self.butler = Butler(self.tmpConfigFile, writeable=True, run="test_run")

2337

2338 # No dimensions in dataset type so we don't have to worry about

2339 # inserting dimension data or defining data IDs.

2340 self.datasetType = DatasetType(

2341 "data", dimensions=(), storageClass="ArrowTable", universe=self.butler.dimensions

2342 )

2343 self.butler.registry.registerDatasetType(self.datasetType)

2344

2345 def tearDown(self):

2346 s3 = boto3.resource("s3")

2347 bucket = s3.Bucket(self.bucketName)

2348 try:

2349 bucket.objects.all().delete()

2350 except botocore.exceptions.ClientError as e:

2351 if e.response["Error"]["Code"] == "404":

2352 # the key was not reachable - pass

2353 pass

2354 else:

2355 raise

2356

2357 bucket = s3.Bucket(self.bucketName)

2358 bucket.delete()

2359

2360 # Stop the S3 mock.

2361 self.mock_aws.stop()

2362

2363 if self.reg_dir is not None and os.path.exists(self.reg_dir):

2364 shutil.rmtree(self.reg_dir, ignore_errors=True)

2365

2366 if os.path.exists(self.root):

2367 shutil.rmtree(self.root, ignore_errors=True)

2368

2369 def testArrowTableS3(self):

2370 tab1 = _makeSimpleArrowTable(include_multidim=True, include_masked=True)

2371

2372 self.butler.put(tab1, self.datasetType, dataId={})

2373

2374 # Read the whole Table.

2375 tab2 = self.butler.get(self.datasetType, dataId={})

2376 # We convert to use the numpy testing framework to handle nan

2377 # comparisons.

2378 self.assertEqual(tab1.schema, tab2.schema)

2379 tab1_np = arrow_to_numpy(tab1)

2380 tab2_np = arrow_to_numpy(tab2)

2381 for col in tab1.column_names:

2382 np.testing.assert_array_equal(tab2_np[col], tab1_np[col])

2383 # Read the columns.

2384 columns2 = self.butler.get(self.datasetType.componentTypeName("columns"), dataId={})

2385 self.assertEqual(len(columns2), len(tab1.schema.names))

2386 for i, name in enumerate(tab1.schema.names):

2387 self.assertEqual(columns2[i], name)

2388 # Read the rowcount.

2389 rowcount = self.butler.get(self.datasetType.componentTypeName("rowcount"), dataId={})

2390 self.assertEqual(rowcount, len(tab1))

2391 # Read the schema.

2392 schema = self.butler.get(self.datasetType.componentTypeName("schema"), dataId={})

2393 self.assertEqual(schema, tab1.schema)

2394 # Read just some columns a few different ways.

2395 tab3 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["a", "c"]})

2396 self.assertEqual(tab3, tab1.select(("a", "c")))

2397 tab4 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": "a"})

2398 self.assertEqual(tab4, tab1.select(("a",)))

2399 tab5 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["index", "a"]})

2400 self.assertEqual(tab5, tab1.select(("index", "a")))

2401 tab6 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": "ddd"})

2402 self.assertEqual(tab6, tab1.select(("ddd",)))

2403 tab7 = self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["a", "a"]})

2404 self.assertEqual(tab7, tab1.select(("a",)))

2405 # Passing an unrecognized column should be a ValueError.

2406 with self.assertRaises(ValueError):

2407 self.butler.get(self.datasetType, dataId={}, parameters={"columns": ["e"]})

2408

2409

2410@unittest.skipUnless(np is not None, "Cannot test compute_row_group_size without numpy.")

2411@unittest.skipUnless(pa is not None, "Cannot test compute_row_group_size without pyarrow.")

2412class ComputeRowGroupSizeTestCase(unittest.TestCase):

2413 """Tests for compute_row_group_size."""

2414

2415 def testRowGroupSizeNoMetadata(self):

2416 numpyTable = _makeSimpleNumpyTable(include_multidim=True)

2417

2418 # We can't use the numpy_to_arrow convenience function because

2419 # that adds metadata.

2420 type_list = _numpy_dtype_to_arrow_types(numpyTable.dtype)

2421 schema = pa.schema(type_list)

2422 arrays = _numpy_style_arrays_to_arrow_arrays(

2423 numpyTable.dtype,

2424 len(numpyTable),

2425 numpyTable,

2426 schema,

2427 )

2428 arrowTable = pa.Table.from_arrays(arrays, schema=schema)

2429

2430 row_group_size = compute_row_group_size(arrowTable.schema)

2431

2432 self.assertGreater(row_group_size, 1_000_000)

2433 self.assertLess(row_group_size, 2_000_000)

2434

2435 def testRowGroupSizeWithMetadata(self):

2436 numpyTable = _makeSimpleNumpyTable(include_multidim=True)

2437

2438 arrowTable = numpy_to_arrow(numpyTable)

2439

2440 row_group_size = compute_row_group_size(arrowTable.schema)

2441

2442 self.assertGreater(row_group_size, 1_000_000)

2443 self.assertLess(row_group_size, 2_000_000)

2444

2445 def testRowGroupSizeTinyTable(self):

2446 numpyTable = np.zeros(1, dtype=[("a", np.bool_)])

2447

2448 arrowTable = numpy_to_arrow(numpyTable)

2449

2450 row_group_size = compute_row_group_size(arrowTable.schema)

2451

2452 self.assertGreater(row_group_size, 1_000_000)

2453

2454 @unittest.skipUnless(pd is not None, "Cannot run testRowGroupSizeDataFrameWithLists without pandas.")

2455 def testRowGroupSizeDataFrameWithLists(self):

2456 df = pd.DataFrame({"a": np.zeros(10), "b": [[0, 0]] * 10, "c": [[0.0, 0.0]] * 10, "d": [[]] * 10})

2457 arrowTable = pandas_to_arrow(df)

2458 row_group_size = compute_row_group_size(arrowTable.schema)

2459

2460 self.assertGreater(row_group_size, 1_000_000)

2461

2462

2463def _checkAstropyTableEquality(table1, table2, skip_units=False, has_bigendian=False):

2464 """Check if two astropy tables have the same columns/values.

2465

2466 Parameters

2467 ----------

2468 table1 : `astropy.table.Table`

2469 table2 : `astropy.table.Table`

2470 skip_units : `bool`

2471 has_bigendian : `bool`

2472 """

2473 if not has_bigendian:

2474 assert table1.dtype == table2.dtype

2475 else:

2476 for name in table1.dtype.names:

2477 # Only check type matches, force to little-endian.

2478 assert table1.dtype[name].newbyteorder(">") == table2.dtype[name].newbyteorder(">")

2479

2480 # Strip provenance before comparison.

2481 DatasetProvenance.strip_provenance_from_flat_dict(table1.meta)

2482 DatasetProvenance.strip_provenance_from_flat_dict(table2.meta)

2483 assert table1.meta == table2.meta

2484 if not skip_units:

2485 for name in table1.columns:

2486 assert table1[name].unit == table2[name].unit

2487 assert table1[name].description == table2[name].description

2488 assert table1[name].format == table2[name].format

2489

2490 for name in table1.columns:

2491 # We need to check masked/regular columns after filling.

2492 has_masked = False

2493 if isinstance(table1[name], atable.column.MaskedColumn):

2494 c1 = table1[name].filled()

2495 has_masked = True

2496 else:

2497 c1 = np.array(table1[name])

2498 if has_masked:

2499 assert isinstance(table2[name], atable.column.MaskedColumn)

2500 c2 = table2[name].filled()

2501 else:

2502 assert not isinstance(table2[name], atable.column.MaskedColumn)

2503 c2 = np.array(table2[name])

2504 np.testing.assert_array_equal(c1, c2)

2505 # If we have a masked column then we test the underlying data.

2506 if has_masked:

2507 np.testing.assert_array_equal(np.array(c1), np.array(c2))

2508 np.testing.assert_array_equal(table1[name].mask, table2[name].mask)

2509

2510

2511def _checkNumpyTableEquality(table1, table2, has_bigendian=False):

2512 """Check if two numpy tables have the same columns/values

2513

2514 Parameters

2515 ----------

2516 table1 : `numpy.ndarray`

2517 table2 : `numpy.ndarray`

2518 has_bigendian : `bool`

2519 """

2520 assert table1.dtype.names == table2.dtype.names

2521 for name in table1.dtype.names:

2522 if not has_bigendian:

2523 assert table1.dtype[name] == table2.dtype[name]

2524 else:

2525 # Only check type matches, force to little-endian.

2526 assert table1.dtype[name].newbyteorder(">") == table2.dtype[name].newbyteorder(">")

2527 assert np.all(table1 == table2)

2528

2529

2530def _checkNumpyDictEquality(dict1, dict2):

2531 """Check if two numpy dicts have the same columns/values.

2532

2533 Parameters

2534 ----------

2535 dict1 : `dict` [`str`, `np.ndarray`]

2536 dict2 : `dict` [`str`, `np.ndarray`]

2537 """

2538 assert set(dict1.keys()) == set(dict2.keys())

2539 for name in dict1:

2540 assert dict1[name].dtype == dict2[name].dtype

2541 assert np.all(dict1[name] == dict2[name])

2542

2543

2544if __name__ == "__main__":

2545 unittest.main()