Coverage for tests/test_parquet.py: 98%

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 self.enterContext(self.butler)

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

368 # inserting dimension data or defining data IDs.

369 self.datasetType = DatasetType(

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

371 )

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

373

374 def tearDown(self):

375 removeTestTempDir(self.root)

376

377 def testSingleIndexDataFrame(self):

378 df1, allColumns = _makeSingleIndexDataFrame(include_masked=True)

379

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

381 # Read the whole DataFrame.

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

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

384 # Read just the column descriptions.

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

386 self.assertEqual(set(allColumns), set(columns2))

387 # Read the rowcount.

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

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

390 # Read the schema.

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

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

393 # Read just some columns a few different ways.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

408 # Passing an unrecognized column should be a ValueError.

409 with self.assertRaises(ValueError):

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

411

412 def testSingleIndexDataFrameWithLists(self):

413 df1, allColumns = _makeSingleIndexDataFrame(include_lists=True)

414

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

416 # Read the whole DataFrame.

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

418

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

420 # from lists to arrays.

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

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

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

424

425 def testMultiIndexDataFrame(self):

426 df1 = _makeMultiIndexDataFrame()

427

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

429 # Read the whole DataFrame.

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

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

432 # Read just the column descriptions.

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

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

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

436 # Read the rowcount.

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

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

439 # Read the schema.

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

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

442 # Read just some columns a few different ways.

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

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

445 df4 = self.butler.get(

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

447 )

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

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

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

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

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

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

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

455 # Passing an unrecognized column should be a ValueError.

456 with self.assertRaises(ValueError):

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

458

459 def testSingleIndexDataFrameEmptyString(self):

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

461 df1, _ = _makeSingleIndexDataFrame()

462

463 # Set one of the strings to None

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

465

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

467 # Read the whole DataFrame.

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

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

470

471 def testSingleIndexDataFrameAllEmptyStrings(self):

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

473 column.

474 """

475 df1, _ = _makeSingleIndexDataFrame()

476

477 # Set all of the strings to None

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

479

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

481 # Read the whole DataFrame.

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

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

484

485 def testLegacyDataFrame(self):

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

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

488 functionality.

489 """

490 df1, allColumns = _makeSingleIndexDataFrame()

491

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

493 df1.to_parquet(fname)

494

495 legacy_type = DatasetType(

496 "legacy_dataframe",

497 dimensions=(),

498 storageClass="DataFrame",

499 universe=self.butler.dimensions,

500 )

501 self.butler.registry.registerDatasetType(legacy_type)

502

503 data_id = {}

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

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

506

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

508

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

510

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

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

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

514

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

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

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

518

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

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

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

522

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

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

525 self.assertEqual(rowcount2a, rowcount2b)

526

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

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

529 self.assertEqual(schema2a, schema2b)

530

531 def testDataFrameSchema(self):

532 tab1 = _makeSimpleArrowTable()

533

534 schema = DataFrameSchema.from_arrow(tab1.schema)

535

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

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

538 self.assertNotEqual(schema, "not_a_schema")

539 self.assertEqual(schema, schema)

540

541 tab2 = _makeMultiIndexDataFrame()

542 schema2 = DataFrameSchema(tab2)

543

544 self.assertNotEqual(schema, schema2)

545

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

547 def testWriteSingleIndexDataFrameReadAsAstropyTable(self):

548 df1, allColumns = _makeSingleIndexDataFrame()

549

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

551

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

553

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

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

556

557 # Check reading the columns.

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

559 columns2 = self.butler.get(

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

561 )

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

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

564

565 # Check reading the schema.

566 schema = ArrowAstropySchema(tab2)

567 schema2 = self.butler.get(

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

569 )

570

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

572 # will be changed because of pandas indexing.

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

574 for name in schema.schema.columns:

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

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

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

578

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

580 def testWriteSingleIndexDataFrameWithMaskedColsReadAsAstropyTable(self):

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

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

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

584 # inconsistent.)

585 df1, allColumns = _makeSingleIndexDataFrame(include_masked=True)

586

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

588

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

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

591

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

593 for name in tab2_df.columns:

594 col1 = df1[name]

595 col2 = tab2_df[name]

596

597 if col1.hasnans:

598 notNull = col1.notnull()

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

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

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

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

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

604 else:

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

606

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

608 def testWriteMultiIndexDataFrameReadAsAstropyTable(self):

609 df1 = _makeMultiIndexDataFrame()

610

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

612

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

614

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

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

617 # recommended.

618

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

620 def testWriteAstropyTableWithMaskedColsReadAsSingleIndexDataFrame(self):

621 tab1 = _makeSimpleAstropyTable(include_masked=True)

622

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

624

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

626

627 tab1_df = astropy_to_pandas(tab1)

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

629

630 tab2_astropy = pandas_to_astropy(tab2)

631 for col in tab1.dtype.names:

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

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

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

635

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

637 def testWriteSingleIndexDataFrameReadAsArrowTable(self):

638 df1, allColumns = _makeSingleIndexDataFrame()

639

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

641

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

643

644 tab2_df = arrow_to_pandas(tab2)

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

646

647 # Check reading the columns.

648 columns = list(tab2.schema.names)

649 columns2 = self.butler.get(

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

651 )

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

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

654

655 # Override the component using a dataset type.

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

657 "ArrowColumnList"

658 )

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

660

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

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

663 # selects the component.

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

665 "columns"

666 )

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

668

669 # Check reading the schema.

670 schema = tab2.schema

671 schema2 = self.butler.get(

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

673 )

674

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

676 # information be maintained.

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

678 for name in schema.names:

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

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

681

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

683 def testWriteMultiIndexDataFrameReadAsArrowTable(self):

684 df1 = _makeMultiIndexDataFrame()

685

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

687

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

689

690 tab2_df = arrow_to_pandas(tab2)

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

692

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

694 def testWriteSingleIndexDataFrameReadAsNumpyTable(self):

695 df1, allColumns = _makeSingleIndexDataFrame()

696

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

698

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

700

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

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

703

704 # Check reading the columns.

705 columns = list(tab2.dtype.names)

706 columns2 = self.butler.get(

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

708 )

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

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

711

712 # Check reading the schema.

713 schema = ArrowNumpySchema(tab2.dtype)

714 schema2 = self.butler.get(

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

716 )

717

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

719 # will be changed because of pandas indexing.

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

721 for name in schema.schema.names:

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

723 # It is not possible to properly track string columns via

724 # the schema consistently.

725 if schema.schema[name].type == np.dtype("O") or schema2.schema[name].type == np.dtype("O"):

726 continue

727 else:

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

729

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

731 def testWriteMultiIndexDataFrameReadAsNumpyTable(self):

732 df1 = _makeMultiIndexDataFrame()

733

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

735

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

737

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

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

740 # recommended.

741

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

743 def testWriteSingleIndexDataFrameReadAsNumpyDict(self):

744 df1, allColumns = _makeSingleIndexDataFrame()

745

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

747

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

749

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

751 # The column order is not maintained.

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

753 for col in df1.columns:

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

755

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

757 def testWriteMultiIndexDataFrameReadAsNumpyDict(self):

758 df1 = _makeMultiIndexDataFrame()

759

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

761

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

763

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

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

766 # recommended.

767

768 def testBadDataFrameColumnParquet(self):

769 df1, allColumns = _makeSingleIndexDataFrame()

770

771 # Make a column with mixed type.

772 bad_col1 = [0.0] * len(df1)

773 bad_col1[1] = 0.0 * units.nJy

774 bad_df = df1.copy()

775 bad_df["bad_col1"] = bad_col1

776

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

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

779 with self.assertRaises(RuntimeError):

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

781

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

783 def testWriteReadAstropyTableLossless(self):

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

785

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

787

788 tab2 = self.butler.get(

789 self.datasetType,

790 dataId={},

791 storageClass="ArrowAstropy",

792 parameters={"strip_astropy_meta_yaml": False},

793 )

794

795 # Check that minimal provenance was written by default.

796 expected = {

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

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

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

800 "LSST.BUTLER.N_INPUTS": 0,

801 }

802

803 self.assertEqual(tab2.meta, expected)

804

805 _checkAstropyTableEquality(tab1, tab2)

806

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

808 def testWriteReadAstropyTableProvenance(self):

809 tab1 = _makeSimpleAstropyTable()

810

811 # Create a ref for provenance.

812 astropy_type = DatasetType(

813 "astropy_parquet",

814 dimensions=(),

815 storageClass="ArrowAstropy",

816 universe=self.butler.dimensions,

817 )

818 self.butler.registry.registerDatasetType(astropy_type)

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

820 quantum_id = uuid.uuid4()

821 provenance = DatasetProvenance(quantum_id=quantum_id)

822 provenance.add_input(input_ref)

823

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

825

826 tab2 = self.butler.get(

827 self.datasetType,

828 dataId={},

829 storageClass="ArrowAstropy",

830 parameters={"strip_astropy_meta_yaml": False},

831 )

832

833 expected = {

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

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

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

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

838 "LSST.BUTLER.N_INPUTS": 1,

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

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

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

842 }

843 self.assertEqual(tab2.meta, expected)

844

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

846 # that the previous provenance was stripped.

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

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

849

850 # tab2 will have been updated in place.

851 expected = {

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

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

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

855 "LSST.BUTLER.N_INPUTS": 0,

856 }

857 self.assertEqual(tab2.meta, expected)

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

859 self.assertEqual(prov_ref, put_ref3)

860 self.assertEqual(null_prov, DatasetProvenance())

861

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

863 def testWriteReadNumpyTableLossless(self):

864 tab1 = _makeSimpleNumpyTable(include_multidim=True)

865

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

867

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

869

870 _checkNumpyTableEquality(tab1, tab2)

871

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

873 def testMaskedNumpy(self):

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

875 tab1_np = arrow_to_numpy(tab1)

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

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

878 col = tab1_np["m_f8"]

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

880

881 # Now without a mask.

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

883 tab1_np = arrow_to_numpy(tab1)

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

885

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

887 def testWriteReadArrowTableLossless(self):

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

889

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

891

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

893

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

895 tab1_np = arrow_to_numpy(tab1)

896 tab2_np = arrow_to_numpy(tab2)

897 for col in tab1.column_names:

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

899

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

901 def testWriteReadNumpyDictLossless(self):

902 tab1 = _makeSimpleNumpyTable(include_multidim=True)

903 dict1 = _numpy_to_numpy_dict(tab1)

904

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

906

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

908

909 _checkNumpyDictEquality(dict1, dict2)

910

911

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

913class InMemoryDataFrameDelegateTestCase(ParquetFormatterDataFrameTestCase):

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

915

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

917

918 def testBadDataFrameColumnParquet(self):

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

920 pass

921

922 def testWriteMultiIndexDataFrameReadAsAstropyTable(self):

923 df1 = _makeMultiIndexDataFrame()

924

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

926

927 with self.assertRaises(ValueError):

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

929

930 def testLegacyDataFrame(self):

931 # This test does not work with an inMemoryDatastore.

932 pass

933

934 def testBadInput(self):

935 df1, _ = _makeSingleIndexDataFrame()

936 delegate = ArrowTableDelegate("DataFrame")

937

938 with self.assertRaises(ValueError):

939 delegate.handleParameters(inMemoryDataset="not_a_dataframe")

940

941 with self.assertRaises(AttributeError):

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

943

944 def testStorageClass(self):

945 df1, allColumns = _makeSingleIndexDataFrame()

946

947 factory = StorageClassFactory()

948 factory.addFromConfig(StorageClassConfig())

949

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

951 # Force the name lookup to do name matching.

952 storageClass._pytype = None

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

954

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

956 # Force the name lookup to do name matching.

957 storageClass._pytype = None

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

959

960

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

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

963class ParquetFormatterArrowAstropyTestCase(unittest.TestCase):

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

965

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

967

968 def setUp(self):

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

970 self.root = makeTestTempDir(TESTDIR)

971 config = Config(self.configFile)

972 self.run = "test_run"

973 self.butler = Butler.from_config(

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

975 )

976 self.enterContext(self.butler)

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

978 # inserting dimension data or defining data IDs.

979 self.datasetType = DatasetType(

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

981 )

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

983

984 def tearDown(self):

985 removeTestTempDir(self.root)

986

987 def testAstropyTable(self):

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

989

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

991 # Read the whole Table.

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

993 _checkAstropyTableEquality(tab1, tab2)

994 # Read the columns.

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

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

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

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

999 # Read the rowcount.

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

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

1002 # Read the schema.

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

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

1005 # Read just some columns a few different ways.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

1020 # Passing an unrecognized column should be a ValueError.

1021 with self.assertRaises(ValueError):

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

1023

1024 def testAstropyTableBigEndian(self):

1025 tab1 = _makeSimpleAstropyTable(include_bigendian=True)

1026

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

1028 # Read the whole Table.

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

1030 _checkAstropyTableEquality(tab1, tab2, has_bigendian=True)

1031

1032 def testAstropyTableWithMetadata(self):

1033 tab1 = _makeSimpleAstropyTable(include_multidim=True)

1034

1035 meta = {

1036 "meta_a": 5,

1037 "meta_b": 10.0,

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

1039 "meta_d": True,

1040 "meta_e": "string",

1041 }

1042

1043 tab1.meta.update(meta)

1044

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

1046 # Read the whole Table.

1047 tab2 = self.butler.get(self.datasetType, dataId={}, parameters={"strip_astropy_meta_yaml": False})

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

1049 _checkAstropyTableEquality(tab1, tab2)

1050

1051 def testArrowAstropySchema(self):

1052 tab1 = _makeSimpleAstropyTable()

1053 tab1_arrow = astropy_to_arrow(tab1)

1054 schema = ArrowAstropySchema.from_arrow(tab1_arrow.schema)

1055

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

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

1058 self.assertNotEqual(schema, "not_a_schema")

1059 self.assertEqual(schema, schema)

1060

1061 # Test various inequalities

1062 tab2 = tab1.copy()

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

1064 schema2 = ArrowAstropySchema(tab2)

1065 self.assertNotEqual(schema2, schema)

1066

1067 tab2 = tab1.copy()

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

1069 schema2 = ArrowAstropySchema(tab2)

1070 self.assertNotEqual(schema2, schema)

1071

1072 tab2 = tab1.copy()

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

1074 schema2 = ArrowAstropySchema(tab2)

1075 self.assertNotEqual(schema2, schema)

1076

1077 tab2 = tab1.copy()

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

1079 schema2 = ArrowAstropySchema(tab2)

1080 self.assertNotEqual(schema2, schema)

1081

1082 def testAstropyParquet(self):

1083 tab1 = _makeSimpleAstropyTable()

1084

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

1086 del tab1["dtn"]

1087 del tab1["dtu"]

1088

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

1090 tab1.write(fname)

1091

1092 astropy_type = DatasetType(

1093 "astropy_parquet",

1094 dimensions=(),

1095 storageClass="ArrowAstropy",

1096 universe=self.butler.dimensions,

1097 )

1098 self.butler.registry.registerDatasetType(astropy_type)

1099

1100 data_id = {}

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

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

1103

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

1105

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

1107

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

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

1110 _checkAstropyTableEquality(tab2a, tab2b)

1111

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

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

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

1115 for i, name in enumerate(columns2a):

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

1117

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

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

1120 self.assertEqual(rowcount2a, rowcount2b)

1121

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

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

1124 self.assertEqual(schema2a, schema2b)

1125

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

1127 def testWriteAstropyReadAsArrowTable(self):

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

1129 tab1 = _makeSimpleAstropyTable(include_masked=True)

1130

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

1132

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

1134

1135 tab2_astropy = arrow_to_astropy(tab2)

1136 _checkAstropyTableEquality(tab1, tab2_astropy)

1137

1138 # Check reading the columns.

1139 columns = tab2.schema.names

1140 columns2 = self.butler.get(

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

1142 )

1143 self.assertEqual(columns2, columns)

1144

1145 # Check reading the schema.

1146 schema = tab2.schema

1147 schema2 = self.butler.get(

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

1149 )

1150

1151 self.assertEqual(schema, schema2)

1152

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

1154 def testWriteAstropyReadAsDataFrame(self):

1155 tab1 = _makeSimpleAstropyTable()

1156

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

1158

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

1160

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

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

1163

1164 tab1_df = tab1.to_pandas()

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

1166

1167 # Check reading the columns.

1168 columns = tab2.columns

1169 columns2 = self.butler.get(

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

1171 )

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

1173

1174 # Check reading the schema.

1175 schema = DataFrameSchema(tab2)

1176 schema2 = self.butler.get(

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

1178 )

1179

1180 self.assertEqual(schema2, schema)

1181

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

1183 def testWriteAstropyWithMaskedColsReadAsDataFrame(self):

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

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

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

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

1188 tab1 = _makeSimpleAstropyTable(include_masked=True)

1189

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

1191

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

1193

1194 tab1_df = astropy_to_pandas(tab1)

1195

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

1197 for name in tab2.columns:

1198 col1 = tab1_df[name]

1199 col2 = tab2[name]

1200

1201 if col1.hasnans:

1202 notNull = col1.notnull()

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

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

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

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

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

1208 else:

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

1210

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

1212 def testWriteSingleIndexDataFrameWithMaskedColsReadAsAstropyTable(self):

1213 df1, allColumns = _makeSingleIndexDataFrame(include_masked=True)

1214

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

1216

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

1218

1219 df1_tab = pandas_to_astropy(df1)

1220

1221 _checkAstropyTableEquality(df1_tab, tab2)

1222

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

1224 def testWriteAstropyReadAsNumpyTable(self):

1225 tab1 = _makeSimpleAstropyTable()

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

1227

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

1229

1230 # This is tricky because it loses the units.

1231 tab2_astropy = atable.Table(tab2)

1232

1233 _checkAstropyTableEquality(tab1, tab2_astropy, skip_units=True)

1234

1235 # Check reading the columns.

1236 columns = list(tab2.dtype.names)

1237 columns2 = self.butler.get(

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

1239 )

1240 self.assertEqual(columns2, columns)

1241

1242 # Check reading the schema.

1243 schema = ArrowNumpySchema(tab2.dtype)

1244 schema2 = self.butler.get(

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

1246 )

1247

1248 self.assertEqual(schema2, schema)

1249

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

1251 def testWriteAstropyReadAsNumpyDict(self):

1252 tab1 = _makeSimpleAstropyTable()

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

1254

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

1256

1257 # This is tricky because it loses the units.

1258 tab2_astropy = atable.Table(tab2)

1259

1260 _checkAstropyTableEquality(tab1, tab2_astropy, skip_units=True)

1261

1262 def testBadAstropyColumnParquet(self):

1263 tab1 = _makeSimpleAstropyTable()

1264

1265 # Make a column with mixed type.

1266 bad_col1 = [0.0] * len(tab1)

1267 bad_col1[1] = 0.0 * units.nJy

1268 bad_tab = tab1.copy()

1269 bad_tab["bad_col1"] = bad_col1

1270

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

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

1273 with self.assertRaises(RuntimeError):

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

1275

1276 # Make a column with ragged size.

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

1278 bad_col2[1] = [0, 0]

1279 bad_tab = tab1.copy()

1280 bad_tab["bad_col2"] = bad_col2

1281

1282 with self.assertRaises(RuntimeError):

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

1284

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

1286 def testWriteAstropyTableWithPandasIndexHint(self, testStrip=True):

1287 tab1 = _makeSimpleAstropyTable()

1288

1289 add_pandas_index_to_astropy(tab1, "index")

1290

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

1292

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

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

1295

1296 self.assertIn(ASTROPY_PANDAS_INDEX_KEY, tab2.meta)

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

1298

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

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

1301

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

1303

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

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

1306 _ = self.butler.get(

1307 self.datasetType,

1308 dataId={},

1309 storageClass="DataFrame",

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

1311 )

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

1313

1314 if testStrip:

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

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

1317

1318 self.assertNotIn(ASTROPY_PANDAS_INDEX_KEY, tab5.meta)

1319

1320 with self.assertRaises(ValueError):

1321 add_pandas_index_to_astropy(tab1, "not_a_column")

1322

1323

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

1325class InMemoryArrowAstropyDelegateTestCase(ParquetFormatterArrowAstropyTestCase):

1326 """Tests for InMemoryDatastore, using ArrowTableDelegate with

1327 AstropyTable.

1328 """

1329

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

1331

1332 def testAstropyParquet(self):

1333 # This test does not work with an inMemoryDatastore.

1334 pass

1335

1336 def testBadAstropyColumnParquet(self):

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

1338 pass

1339

1340 def testBadInput(self):

1341 tab1 = _makeSimpleAstropyTable()

1342 delegate = ArrowTableDelegate("ArrowAstropy")

1343

1344 with self.assertRaises(ValueError):

1345 delegate.handleParameters(inMemoryDataset="not_an_astropy_table")

1346

1347 with self.assertRaises(NotImplementedError):

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

1349

1350 with self.assertRaises(AttributeError):

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

1352

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

1354 def testWriteAstropyTableWithPandasIndexHint(self):

1355 super().testWriteAstropyTableWithPandasIndexHint(testStrip=False)

1356

1357

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

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

1360class ParquetFormatterArrowNumpyTestCase(unittest.TestCase):

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

1362

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

1364

1365 def setUp(self):

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

1367 self.root = makeTestTempDir(TESTDIR)

1368 config = Config(self.configFile)

1369 self.butler = Butler.from_config(

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

1371 )

1372 self.enterContext(self.butler)

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

1374 # inserting dimension data or defining data IDs.

1375 self.datasetType = DatasetType(

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

1377 )

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

1379

1380 def tearDown(self):

1381 removeTestTempDir(self.root)

1382

1383 def testNumpyTable(self):

1384 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1385

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

1387 # Read the whole Table.

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

1389 _checkNumpyTableEquality(tab1, tab2)

1390 # Read the columns.

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

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

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

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

1395 # Read the rowcount.

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

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

1398 # Read the schema.

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

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

1401 # Read just some columns a few different ways.

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

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

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

1405 _checkNumpyTableEquality(

1406 tab1[

1407 [

1408 "a",

1409 ]

1410 ],

1411 tab4,

1412 )

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

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

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

1416 _checkNumpyTableEquality(

1417 tab1[

1418 [

1419 "ddd",

1420 ]

1421 ],

1422 tab6,

1423 )

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

1425 _checkNumpyTableEquality(

1426 tab1[

1427 [

1428 "a",

1429 ]

1430 ],

1431 tab7,

1432 )

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

1434 _checkNumpyTableEquality(

1435 tab1[

1436 [

1437 "ddd",

1438 "dtn",

1439 "dtu",

1440 "a",

1441 ]

1442 ],

1443 tab8,

1444 )

1445 # Passing an unrecognized column should be a ValueError.

1446 with self.assertRaises(ValueError):

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

1448

1449 def testNumpyTableBigEndian(self):

1450 tab1 = _makeSimpleNumpyTable(include_bigendian=True)

1451

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

1453 # Read the whole Table.

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

1455 _checkNumpyTableEquality(tab1, tab2, has_bigendian=True)

1456

1457 def testArrowNumpySchema(self):

1458 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1459 tab1_arrow = numpy_to_arrow(tab1)

1460 schema = ArrowNumpySchema.from_arrow(tab1_arrow.schema)

1461

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

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

1464 self.assertNotEqual(schema, "not_a_schema")

1465 self.assertEqual(schema, schema)

1466

1467 # Test inequality

1468 tab2 = tab1.copy()

1469 names = list(tab2.dtype.names)

1470 names[0] = "index2"

1471 tab2.dtype.names = names

1472 schema2 = ArrowNumpySchema(tab2.dtype)

1473 self.assertNotEqual(schema2, schema)

1474

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

1476 def testNumpyDictConversions(self):

1477 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1478

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

1480 tab1_arrow = numpy_to_arrow(tab1)

1481 tab1_dict = arrow_to_numpy_dict(tab1_arrow)

1482 tab1_dict_arrow = numpy_dict_to_arrow(tab1_dict)

1483

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

1485 self.assertEqual(tab1_arrow, tab1_dict_arrow)

1486

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

1488 def testWriteNumpyTableReadAsArrowTable(self):

1489 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1490

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

1492

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

1494

1495 tab2_numpy = arrow_to_numpy(tab2)

1496

1497 _checkNumpyTableEquality(tab1, tab2_numpy)

1498

1499 # Check reading the columns.

1500 columns = tab2.schema.names

1501 columns2 = self.butler.get(

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

1503 )

1504 self.assertEqual(columns2, columns)

1505

1506 # Check reading the schema.

1507 schema = tab2.schema

1508 schema2 = self.butler.get(

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

1510 )

1511 self.assertEqual(schema2, schema)

1512

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

1514 def testWriteNumpyTableReadAsDataFrame(self):

1515 tab1 = _makeSimpleNumpyTable()

1516

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

1518

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

1520

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

1522 # and changes the datatype of the string column.

1523

1524 tab1_df = pd.DataFrame(tab1)

1525

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

1527

1528 # Check reading the columns.

1529 columns = tab2.columns

1530 columns2 = self.butler.get(

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

1532 )

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

1534

1535 # Check reading the schema.

1536 schema = DataFrameSchema(tab2)

1537 schema2 = self.butler.get(

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

1539 )

1540

1541 self.assertEqual(schema2, schema)

1542

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

1544 def testWriteNumpyTableReadAsAstropyTable(self):

1545 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1546

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

1548

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

1550 tab2_numpy = tab2.as_array()

1551

1552 _checkNumpyTableEquality(tab1, tab2_numpy)

1553

1554 # Check reading the columns.

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

1556 columns2 = self.butler.get(

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

1558 )

1559 self.assertEqual(columns2, columns)

1560

1561 # Check reading the schema.

1562 schema = ArrowAstropySchema(tab2)

1563 schema2 = self.butler.get(

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

1565 )

1566

1567 self.assertEqual(schema2, schema)

1568

1569 def testWriteNumpyTableReadAsNumpyDict(self):

1570 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1571

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

1573

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

1575 tab2_numpy = _numpy_dict_to_numpy(tab2)

1576

1577 _checkNumpyTableEquality(tab1, tab2_numpy)

1578

1579 def testBadNumpyColumnParquet(self):

1580 tab1 = _makeSimpleAstropyTable()

1581

1582 # Make a column with mixed type.

1583 bad_col1 = [0.0] * len(tab1)

1584 bad_col1[1] = 0.0 * units.nJy

1585 bad_tab = tab1.copy()

1586 bad_tab["bad_col1"] = bad_col1

1587

1588 bad_tab_np = bad_tab.as_array()

1589

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

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

1592 with self.assertRaises(RuntimeError):

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

1594

1595 # Make a column with ragged size.

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

1597 bad_col2[1] = [0, 0]

1598 bad_tab = tab1.copy()

1599 bad_tab["bad_col2"] = bad_col2

1600

1601 bad_tab_np = bad_tab.as_array()

1602

1603 with self.assertRaises(RuntimeError):

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

1605

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

1607 def testWriteReadAstropyTableLossless(self):

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

1609

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

1611

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

1613

1614 _checkAstropyTableEquality(tab1, tab2)

1615

1616

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

1618class InMemoryArrowNumpyDelegateTestCase(ParquetFormatterArrowNumpyTestCase):

1619 """Tests for InMemoryDatastore, using ArrowTableDelegate with

1620 Numpy table.

1621 """

1622

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

1624

1625 def testBadNumpyColumnParquet(self):

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

1627 pass

1628

1629 def testBadInput(self):

1630 tab1 = _makeSimpleNumpyTable()

1631 delegate = ArrowTableDelegate("ArrowNumpy")

1632

1633 with self.assertRaises(ValueError):

1634 delegate.handleParameters(inMemoryDataset="not_a_numpy_table")

1635

1636 with self.assertRaises(NotImplementedError):

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

1638

1639 with self.assertRaises(AttributeError):

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

1641

1642 def testStorageClass(self):

1643 tab1 = _makeSimpleNumpyTable()

1644

1645 factory = StorageClassFactory()

1646 factory.addFromConfig(StorageClassConfig())

1647

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

1649 # Force the name lookup to do name matching.

1650 storageClass._pytype = None

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

1652

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

1654 # Force the name lookup to do name matching.

1655 storageClass._pytype = None

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

1657

1658

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

1660class ParquetFormatterArrowTableTestCase(unittest.TestCase):

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

1662

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

1664

1665 def setUp(self):

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

1667 self.root = makeTestTempDir(TESTDIR)

1668 config = Config(self.configFile)

1669 self.butler = Butler.from_config(

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

1671 )

1672 self.enterContext(self.butler)

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

1674 # inserting dimension data or defining data IDs.

1675 self.datasetType = DatasetType(

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

1677 )

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

1679

1680 def tearDown(self):

1681 removeTestTempDir(self.root)

1682

1683 def testArrowTable(self):

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

1685

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

1687 # Read the whole Table.

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

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

1690 # comparisons.

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

1692 tab1_np = arrow_to_numpy(tab1)

1693 tab2_np = arrow_to_numpy(tab2)

1694 for col in tab1.column_names:

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

1696 # Read the columns.

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

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

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

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

1701 # Read the rowcount.

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

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

1704 # Read the schema.

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

1706 self.assertEqual(schema, tab1.schema)

1707 # Read just some columns a few different ways.

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

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

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

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

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

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

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

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

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

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

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

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

1720 # Passing an unrecognized column should be a ValueError.

1721 with self.assertRaises(ValueError):

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

1723

1724 def testEmptyArrowTable(self):

1725 data = _makeSimpleNumpyTable()

1726 type_list = _numpy_dtype_to_arrow_types(data.dtype)

1727

1728 schema = pa.schema(type_list)

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

1730

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

1732

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

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

1735 self.assertEqual(tab2, tab1)

1736

1737 tab1_numpy = arrow_to_numpy(tab1)

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

1739 tab1_numpy_arrow = numpy_to_arrow(tab1_numpy)

1740 self.assertEqual(tab1_numpy_arrow, tab1)

1741

1742 tab1_pandas = arrow_to_pandas(tab1)

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

1744 tab1_pandas_arrow = pandas_to_arrow(tab1_pandas)

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

1746 # through empty pandas dataframes.

1747 self.assertEqual(

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

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

1750 )

1751

1752 tab1_astropy = arrow_to_astropy(tab1)

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

1754 tab1_astropy_arrow = astropy_to_arrow(tab1_astropy)

1755 self.assertEqual(tab1_astropy_arrow, tab1)

1756

1757 def testEmptyArrowTableMultidim(self):

1758 data = _makeSimpleNumpyTable(include_multidim=True)

1759 type_list = _numpy_dtype_to_arrow_types(data.dtype)

1760

1761 md = {}

1762 for name in data.dtype.names:

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

1764

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

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

1767

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

1769

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

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

1772 self.assertEqual(tab2, tab1)

1773

1774 tab1_numpy = arrow_to_numpy(tab1)

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

1776 tab1_numpy_arrow = numpy_to_arrow(tab1_numpy)

1777 self.assertEqual(tab1_numpy_arrow, tab1)

1778

1779 tab1_astropy = arrow_to_astropy(tab1)

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

1781 tab1_astropy_arrow = astropy_to_arrow(tab1_astropy)

1782 self.assertEqual(tab1_astropy_arrow, tab1)

1783

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

1785 def testWriteArrowTableReadAsSingleIndexDataFrame(self):

1786 df1, allColumns = _makeSingleIndexDataFrame()

1787

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

1789

1790 # Read back out as a dataframe.

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

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

1793

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

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

1796 df3 = arrow_to_pandas(tab3)

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

1798

1799 # Check reading the columns.

1800 columns = df2.reset_index().columns

1801 columns2 = self.butler.get(

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

1803 )

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

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

1806

1807 # Check reading the schema.

1808 schema = DataFrameSchema(df1)

1809 schema2 = self.butler.get(

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

1811 )

1812 self.assertEqual(schema2, schema)

1813

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

1815 def testWriteArrowTableReadAsMultiIndexDataFrame(self):

1816 df1 = _makeMultiIndexDataFrame()

1817

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

1819

1820 # Read back out as a dataframe.

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

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

1823

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

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

1826 df3 = arrow_to_pandas(atab3)

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

1828

1829 # Check reading the columns.

1830 columns = df2.columns

1831 columns2 = self.butler.get(

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

1833 )

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

1835

1836 # Check reading the schema.

1837 schema = DataFrameSchema(df1)

1838 schema2 = self.butler.get(

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

1840 )

1841 self.assertEqual(schema2, schema)

1842

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

1844 def testWriteArrowTableReadAsAstropyTable(self):

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

1846

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

1848

1849 # Read back out as an astropy table.

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

1851 _checkAstropyTableEquality(tab1, tab2)

1852

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

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

1855 tab3 = arrow_to_astropy(atab3)

1856 _checkAstropyTableEquality(tab1, tab3)

1857

1858 # Check reading the columns.

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

1860 columns2 = self.butler.get(

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

1862 )

1863 self.assertEqual(columns2, columns)

1864

1865 # Check reading the schema.

1866 schema = ArrowAstropySchema(tab1)

1867 schema2 = self.butler.get(

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

1869 )

1870 self.assertEqual(schema2, schema)

1871

1872 # Check the schema conversions and units.

1873 arrow_schema = schema.to_arrow_schema()

1874 for name in arrow_schema.names:

1875 field_metadata = arrow_schema.field(name).metadata

1876 if (

1877 b"description" in field_metadata

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

1879 ):

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

1881 else:

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

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

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

1885

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

1887 def testWriteArrowTableReadAsNumpyTable(self):

1888 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1889

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

1891

1892 # Read back out as a numpy table.

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

1894 _checkNumpyTableEquality(tab1, tab2)

1895

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

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

1898 tab3 = arrow_to_numpy(atab3)

1899 _checkNumpyTableEquality(tab1, tab3)

1900

1901 # Check reading the columns.

1902 columns = list(tab2.dtype.names)

1903 columns2 = self.butler.get(

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

1905 )

1906 self.assertEqual(columns2, columns)

1907

1908 # Check reading the schema.

1909 schema = ArrowNumpySchema(tab1.dtype)

1910 schema2 = self.butler.get(

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

1912 )

1913 self.assertEqual(schema2, schema)

1914

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

1916 def testWriteArrowTableReadAsNumpyDict(self):

1917 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1918

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

1920

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

1922 tab2_numpy = _numpy_dict_to_numpy(tab2)

1923 _checkNumpyTableEquality(tab1, tab2_numpy)

1924

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

1926 def testWriteReadAstropyTableLossless(self):

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

1928

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

1930

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

1932

1933 _checkAstropyTableEquality(tab1, tab2)

1934

1935

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

1937class InMemoryArrowTableDelegateTestCase(ParquetFormatterArrowTableTestCase):

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

1939

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

1941

1942 def testBadInput(self):

1943 tab1 = _makeSimpleArrowTable()

1944 delegate = ArrowTableDelegate("ArrowTable")

1945

1946 with self.assertRaises(ValueError):

1947 delegate.handleParameters(inMemoryDataset="not_an_arrow_table")

1948

1949 with self.assertRaises(NotImplementedError):

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

1951

1952 with self.assertRaises(AttributeError):

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

1954

1955 def testStorageClass(self):

1956 tab1 = _makeSimpleArrowTable()

1957

1958 factory = StorageClassFactory()

1959 factory.addFromConfig(StorageClassConfig())

1960

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

1962 # Force the name lookup to do name matching.

1963 storageClass._pytype = None

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

1965

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

1967 # Force the name lookup to do name matching.

1968 storageClass._pytype = None

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

1970

1971

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

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

1974class ParquetFormatterArrowNumpyDictTestCase(unittest.TestCase):

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

1976

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

1978

1979 def setUp(self):

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

1981 self.root = makeTestTempDir(TESTDIR)

1982 config = Config(self.configFile)

1983 self.butler = Butler.from_config(

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

1985 )

1986 self.enterContext(self.butler)

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

1988 # inserting dimension data or defining data IDs.

1989 self.datasetType = DatasetType(

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

1991 )

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

1993

1994 def tearDown(self):

1995 removeTestTempDir(self.root)

1996

1997 def testNumpyDict(self):

1998 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1999 dict1 = _numpy_to_numpy_dict(tab1)

2000

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

2002 # Read the whole table.

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

2004 _checkNumpyDictEquality(dict1, dict2)

2005 # Read the columns.

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

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

2008 for name in dict1:

2009 self.assertIn(name, columns2)

2010 # Read the rowcount.

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

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

2013 # Read the schema.

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

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

2016 # Read just some columns a few different ways.

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

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

2019 _checkNumpyDictEquality(subdict, tab3)

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

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

2022 _checkNumpyDictEquality(subdict, tab4)

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

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

2025 _checkNumpyDictEquality(subdict, tab5)

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

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

2028 _checkNumpyDictEquality(subdict, tab6)

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

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

2031 _checkNumpyDictEquality(subdict, tab7)

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

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

2034 _checkNumpyDictEquality(subdict, tab8)

2035 # Passing an unrecognized column should be a ValueError.

2036 with self.assertRaises(ValueError):

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

2038

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

2040 def testWriteNumpyDictReadAsArrowTable(self):

2041 tab1 = _makeSimpleNumpyTable(include_multidim=True)

2042 dict1 = _numpy_to_numpy_dict(tab1)

2043

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

2045

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

2047

2048 tab2_dict = arrow_to_numpy_dict(tab2)

2049

2050 _checkNumpyDictEquality(dict1, tab2_dict)

2051

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

2053 def testWriteNumpyDictReadAsDataFrame(self):

2054 tab1 = _makeSimpleNumpyTable()

2055 dict1 = _numpy_to_numpy_dict(tab1)

2056

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

2058

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

2060

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

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

2063 # changes the datatype of the string column.

2064 tab1_df = pd.DataFrame(tab1)

2065

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

2067 for col in tab1_df.columns:

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

2069

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

2071 def testWriteNumpyDictReadAsAstropyTable(self):

2072 tab1 = _makeSimpleNumpyTable(include_multidim=True)

2073 dict1 = _numpy_to_numpy_dict(tab1)

2074

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

2076

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

2078 tab2_dict = _astropy_to_numpy_dict(tab2)

2079

2080 _checkNumpyDictEquality(dict1, tab2_dict)

2081

2082 def testWriteNumpyDictReadAsNumpyTable(self):

2083 tab1 = _makeSimpleNumpyTable(include_multidim=True)

2084 dict1 = _numpy_to_numpy_dict(tab1)

2085

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

2087

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

2089 tab2_dict = _numpy_to_numpy_dict(tab2)

2090

2091 _checkNumpyDictEquality(dict1, tab2_dict)

2092

2093 def testWriteNumpyDictBad(self):

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

2095 with self.assertRaises(RuntimeError):

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

2097

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

2099 with self.assertRaises(RuntimeError):

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

2101

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

2103 with self.assertRaises(RuntimeError):

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

2105

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

2107 with self.assertRaises(RuntimeError):

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

2109

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

2111 def testWriteReadAstropyTableLossless(self):

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

2113

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

2115

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

2117

2118 _checkAstropyTableEquality(tab1, tab2)

2119

2120

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

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

2123class InMemoryNumpyDictDelegateTestCase(ParquetFormatterArrowNumpyDictTestCase):

2124 """Tests for InMemoryDatastore, using ArrowTableDelegate with

2125 Numpy dict.

2126 """

2127

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

2129

2130 def testWriteNumpyDictBad(self):

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

2132 pass

2133

2134

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

2136class ParquetFormatterArrowSchemaTestCase(unittest.TestCase):

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

2138

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

2140

2141 def setUp(self):

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

2143 self.root = makeTestTempDir(TESTDIR)

2144 config = Config(self.configFile)

2145 self.butler = Butler.from_config(

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

2147 )

2148 self.enterContext(self.butler)

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

2150 # inserting dimension data or defining data IDs.

2151 self.datasetType = DatasetType(

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

2153 )

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

2155

2156 def tearDown(self):

2157 removeTestTempDir(self.root)

2158

2159 def _makeTestSchema(self):

2160 schema = pa.schema(

2161 [

2162 pa.field(

2163 "int32",

2164 pa.int32(),

2165 nullable=False,

2166 metadata={

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

2168 "unit": "",

2169 },

2170 ),

2171 pa.field(

2172 "int64",

2173 pa.int64(),

2174 nullable=False,

2175 metadata={

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

2177 "unit": "",

2178 },

2179 ),

2180 pa.field(

2181 "uint64",

2182 pa.uint64(),

2183 nullable=False,

2184 metadata={

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

2186 "unit": "",

2187 },

2188 ),

2189 pa.field(

2190 "float32",

2191 pa.float32(),

2192 nullable=False,

2193 metadata={

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

2195 "unit": "count",

2196 },

2197 ),

2198 pa.field(

2199 "float64",

2200 pa.float64(),

2201 nullable=False,

2202 metadata={

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

2204 "unit": "nJy",

2205 },

2206 ),

2207 pa.field(

2208 "fixed_size_list",

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

2210 nullable=False,

2211 metadata={

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

2213 "unit": "nJy",

2214 },

2215 ),

2216 pa.field(

2217 "variable_size_list",

2218 pa.list_(pa.float64()),

2219 nullable=False,

2220 metadata={

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

2222 "unit": "nJy",

2223 },

2224 ),

2225 # One of these fields will have no description.

2226 pa.field(

2227 "string",

2228 pa.string(),

2229 nullable=False,

2230 metadata={

2231 "unit": "",

2232 },

2233 ),

2234 # One of these fields will have no metadata.

2235 pa.field(

2236 "binary",

2237 pa.binary(),

2238 nullable=False,

2239 ),

2240 ]

2241 )

2242

2243 return schema

2244

2245 def testArrowSchema(self):

2246 schema1 = self._makeTestSchema()

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

2248

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

2250 self.assertEqual(schema2, schema1)

2251

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

2253 def testWriteArrowSchemaReadAsDataFrameSchema(self):

2254 schema1 = self._makeTestSchema()

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

2256

2257 df_schema1 = DataFrameSchema.from_arrow(schema1)

2258

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

2260 self.assertEqual(df_schema2, df_schema1)

2261

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

2263 def testWriteArrowSchemaReadAsArrowAstropySchema(self):

2264 schema1 = self._makeTestSchema()

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

2266

2267 ap_schema1 = ArrowAstropySchema.from_arrow(schema1)

2268

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

2270 self.assertEqual(ap_schema2, ap_schema1)

2271

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

2273 for name in schema1.names:

2274 field_metadata = schema1.field(name).metadata

2275 if field_metadata is None:

2276 continue

2277 if (

2278 b"description" in field_metadata

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

2280 ):

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

2282 else:

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

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

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

2286

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

2288 def testWriteArrowSchemaReadAsArrowNumpySchema(self):

2289 schema1 = self._makeTestSchema()

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

2291

2292 np_schema1 = ArrowNumpySchema.from_arrow(schema1)

2293

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

2295 self.assertEqual(np_schema2, np_schema1)

2296

2297

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

2299class InMemoryArrowSchemaDelegateTestCase(ParquetFormatterArrowSchemaTestCase):

2300 """Tests for InMemoryDatastore and ArrowSchema."""

2301

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

2303

2304

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

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

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

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

2309class ParquetFormatterArrowTableS3TestCase(unittest.TestCase):

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

2311

2312 # Code is adapted from test_butler.py

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

2314 fullConfigKey = None

2315 validationCanFail = True

2316

2317 bucketName = "anybucketname"

2318

2319 root = "butlerRoot/"

2320

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

2322

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

2324

2325 registryStr = "/gen3.sqlite3"

2326

2327 mock_aws = mock_aws()

2328

2329 def setUp(self):

2330 self.root = makeTestTempDir(TESTDIR)

2331

2332 config = Config(self.configFile)

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

2334 self.bucketName = uri.netloc

2335

2336 # Enable S3 mocking of tests.

2337 self.enterContext(clean_test_environment_for_s3())

2338 self.mock_aws.start()

2339

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

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

2342

2343 # need local folder to store registry database

2344 self.reg_dir = makeTestTempDir(TESTDIR)

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

2346

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

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

2349 s3 = boto3.resource("s3")

2350 s3.create_bucket(Bucket=self.bucketName)

2351

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

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

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

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

2356

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

2358 self.enterContext(self.butler)

2359

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

2361 # inserting dimension data or defining data IDs.

2362 self.datasetType = DatasetType(

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

2364 )

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

2366

2367 def tearDown(self):

2368 s3 = boto3.resource("s3")

2369 bucket = s3.Bucket(self.bucketName)

2370 try:

2371 bucket.objects.all().delete()

2372 except botocore.exceptions.ClientError as e:

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

2374 # the key was not reachable - pass

2375 pass

2376 else:

2377 raise

2378

2379 bucket = s3.Bucket(self.bucketName)

2380 bucket.delete()

2381

2382 # Stop the S3 mock.

2383 self.mock_aws.stop()

2384

2385 if self.reg_dir is not None and os.path.exists(self.reg_dir): 2385 ↛ 2388line 2385 didn't jump to line 2388 because the condition on line 2385 was always true

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

2387

2388 if os.path.exists(self.root): 2388 ↛ exitline 2388 didn't return from function 'tearDown' because the condition on line 2388 was always true

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

2390

2391 def testArrowTableS3(self):

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

2393

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

2395

2396 # Read the whole Table.

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

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

2399 # comparisons.

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

2401 tab1_np = arrow_to_numpy(tab1)

2402 tab2_np = arrow_to_numpy(tab2)

2403 for col in tab1.column_names:

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

2405 # Read the columns.

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

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

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

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

2410 # Read the rowcount.

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

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

2413 # Read the schema.

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

2415 self.assertEqual(schema, tab1.schema)

2416 # Read just some columns a few different ways.

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

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

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

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

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

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

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

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

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

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

2427 # Passing an unrecognized column should be a ValueError.

2428 with self.assertRaises(ValueError):

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

2430

2431

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

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

2434class ComputeRowGroupSizeTestCase(unittest.TestCase):

2435 """Tests for compute_row_group_size."""

2436

2437 def testRowGroupSizeNoMetadata(self):

2438 numpyTable = _makeSimpleNumpyTable(include_multidim=True)

2439

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

2441 # that adds metadata.

2442 type_list = _numpy_dtype_to_arrow_types(numpyTable.dtype)

2443 schema = pa.schema(type_list)

2444 arrays = _numpy_style_arrays_to_arrow_arrays(

2445 numpyTable.dtype,

2446 len(numpyTable),

2447 numpyTable,

2448 schema,

2449 )

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

2451

2452 row_group_size = compute_row_group_size(arrowTable.schema)

2453

2454 self.assertGreater(row_group_size, 1_000_000)

2455 self.assertLess(row_group_size, 2_000_000)

2456

2457 def testRowGroupSizeWithMetadata(self):

2458 numpyTable = _makeSimpleNumpyTable(include_multidim=True)

2459

2460 arrowTable = numpy_to_arrow(numpyTable)

2461

2462 row_group_size = compute_row_group_size(arrowTable.schema)

2463

2464 self.assertGreater(row_group_size, 1_000_000)

2465 self.assertLess(row_group_size, 2_000_000)

2466

2467 def testRowGroupSizeTinyTable(self):

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

2469

2470 arrowTable = numpy_to_arrow(numpyTable)

2471

2472 row_group_size = compute_row_group_size(arrowTable.schema)

2473

2474 self.assertGreater(row_group_size, 1_000_000)

2475

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

2477 def testRowGroupSizeDataFrameWithLists(self):

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

2479 arrowTable = pandas_to_arrow(df)

2480 row_group_size = compute_row_group_size(arrowTable.schema)

2481

2482 self.assertGreater(row_group_size, 1_000_000)

2483

2484

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

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

2487

2488 Parameters

2489 ----------

2490 table1 : `astropy.table.Table`

2491 table2 : `astropy.table.Table`

2492 skip_units : `bool`

2493 has_bigendian : `bool`

2494 """

2495 if not has_bigendian:

2496 assert table1.dtype == table2.dtype

2497 else:

2498 for name in table1.dtype.names:

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

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

2501

2502 # Strip provenance before comparison.

2503 DatasetProvenance.strip_provenance_from_flat_dict(table1.meta)

2504 DatasetProvenance.strip_provenance_from_flat_dict(table2.meta)

2505 assert table1.meta == table2.meta

2506 if not skip_units:

2507 for name in table1.columns:

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

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

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

2511

2512 for name in table1.columns:

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

2514 has_masked = False

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

2516 c1 = table1[name].filled()

2517 has_masked = True

2518 else:

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

2520 if has_masked:

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

2522 c2 = table2[name].filled()

2523 else:

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

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

2526 np.testing.assert_array_equal(c1, c2)

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

2528 if has_masked:

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

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

2531

2532

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

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

2535

2536 Parameters

2537 ----------

2538 table1 : `numpy.ndarray`

2539 table2 : `numpy.ndarray`

2540 has_bigendian : `bool`

2541 """

2542 assert table1.dtype.names == table2.dtype.names

2543 for name in table1.dtype.names:

2544 if not has_bigendian:

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

2546 else:

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

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

2549 assert np.all(table1 == table2)

2550

2551

2552def _checkNumpyDictEquality(dict1, dict2):

2553 """Check if two numpy dicts have the same columns/values.

2554

2555 Parameters

2556 ----------

2557 dict1 : `dict` [`str`, `np.ndarray`]

2558 dict2 : `dict` [`str`, `np.ndarray`]

2559 """

2560 assert set(dict1.keys()) == set(dict2.keys())

2561 for name in dict1:

2562 assert dict1[name].dtype == dict2[name].dtype

2563 assert np.all(dict1[name] == dict2[name])

2564

2565

2566if __name__ == "__main__":

2567 unittest.main()