Coverage for tests/test_parquet.py: 17%

1# This file is part of daf_butler.

3# Developed for the LSST Data Management System.

4# This product includes software developed by the LSST Project

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

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

7# for details of code ownership.

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

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

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

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

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

14#

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

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

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

18# (at your option) any later version.

19#

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

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

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

23# GNU General Public License for more details.

24#

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

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

28"""Tests for ParquetFormatter.

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

31"""

33import datetime

34import os

35import posixpath

36import shutil

37import unittest

38import uuid

40try:

41 import pyarrow as pa

42except ImportError:

43 pa = None

44try:

45 import astropy.table as atable

46 from astropy import units

47except ImportError:

48 atable = None

49try:

50 import numpy as np

51except ImportError:

52 np = None

53try:

54 import pandas as pd

55except ImportError:

56 pd = None

58try:

59 import boto3

60 import botocore

62 from lsst.resources.s3utils import clean_test_environment_for_s3

64 try:

65 from moto import mock_aws # v5

66 except ImportError:

67 from moto import mock_s3 as mock_aws

68except ImportError:

69 boto3 = None

71try:

72 import fsspec

73except ImportError:

74 fsspec = None

76try:

77 import s3fs

78except ImportError:

79 s3fs = None

82from lsst.daf.butler import (

83 Butler,

84 Config,

85 DatasetProvenance,

86 DatasetRef,

87 DatasetType,

88 FileDataset,

89 StorageClassConfig,

90 StorageClassFactory,

91)

92from lsst.resources import ResourcePath

94try:

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

96except ImportError:

97 pa = None

99try:

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

101 ASTROPY_PANDAS_INDEX_KEY,

102 ArrowAstropySchema,

103 ArrowNumpySchema,

104 DataFrameSchema,

105 ParquetFormatter,

106 _append_numpy_multidim_metadata,

107 _astropy_to_numpy_dict,

108 _numpy_dict_to_numpy,

109 _numpy_dtype_to_arrow_types,

110 _numpy_style_arrays_to_arrow_arrays,

111 _numpy_to_numpy_dict,

112 add_pandas_index_to_astropy,

113 arrow_to_astropy,

114 arrow_to_numpy,

115 arrow_to_numpy_dict,

116 arrow_to_pandas,

117 astropy_to_arrow,

118 astropy_to_pandas,

119 compute_row_group_size,

120 numpy_dict_to_arrow,

121 numpy_to_arrow,

122 pandas_to_arrow,

123 pandas_to_astropy,

124 )

125except ImportError:

126 pa = None

127 pd = None

128 atable = None

129 np = None

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

131

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

133

134

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

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

137

138 Parameters

139 ----------

140 include_multidim : `bool`

141 Include multi-dimensional columns.

142 include_bigendian : `bool`

143 Include big-endian columns.

144

145 Returns

146 -------

147 numpyTable : `numpy.ndarray`

148 """

149 nrow = 5

150

151 dtype = [

152 ("index", "i4"),

153 ("a", "f8"),

154 ("b", "f8"),

155 ("c", "f8"),

156 ("ddd", "f8"),

157 ("f", "i8"),

158 ("strcol", "U10"),

159 ("bytecol", "S10"),

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

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

162 ]

163

164 if include_multidim:

165 dtype.extend(

166 [

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

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

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

170 ]

171 )

172

173 if include_bigendian:

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

175

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

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

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

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

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

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

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

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

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

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

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

187

188 if include_multidim:

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

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

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

192

193 if include_bigendian:

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

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

196

197 return data

198

199

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

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

202

203 Parameters

204 ----------

205 include_masked : `bool`

206 Include masked columns.

207 include_lists : `bool`

208 Include list columns.

209

210 Returns

211 -------

212 dataFrame : `~pandas.DataFrame`

213 The test dataframe.

214 allColumns : `list` [`str`]

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

216 """

217 data = _makeSimpleNumpyTable()

218 df = pd.DataFrame(data)

219 df = df.set_index("index")

220

221 if include_masked:

222 nrow = len(df)

223

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

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

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

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

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

229

230 if include_lists:

231 nrow = len(df)

232

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

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

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

236

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

238

239 return df, allColumns

240

241

242def _makeMultiIndexDataFrame():

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

244

245 Returns

246 -------

247 dataFrame : `~pandas.DataFrame`

248 The test dataframe.

249 """

250 columns = pd.MultiIndex.from_tuples(

251 [

252 ("g", "a"),

253 ("g", "b"),

254 ("g", "c"),

255 ("r", "a"),

256 ("r", "b"),

257 ("r", "c"),

258 ],

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

260 )

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

262

263 return df

264

265

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

267 """Make an astropy table for testing.

268

269 Parameters

270 ----------

271 include_multidim : `bool`

272 Include multi-dimensional columns.

273 include_masked : `bool`

274 Include masked columns.

275 include_bigendian : `bool`

276 Include big-endian columns.

277

278 Returns

279 -------

280 astropyTable : `astropy.table.Table`

281 The test table.

282 """

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

284 # Add a couple of units.

285 table = atable.Table(data)

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

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

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

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

290

291 # Add some masked columns.

292 if include_masked:

293 nrow = len(table)

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

295 mask[1] = True

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

297 # to be able test after serialization.

298

299 # Masked 64-bit integer.

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

301 arr[mask] = -1

302 arr[0] = 1649900760361600113

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

304 # Masked 32-bit float.

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

306 arr[mask] = np.nan

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

308 # Unmasked 32-bit float with NaNs.

309 table["um_f4"] = arr

310 # Masked 64-bit float.

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

312 arr[mask] = np.nan

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

314 # Unmasked 64-bit float with NaNs.

315 table["um_f8"] = arr

316 # Masked boolean.

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

318 arr[mask] = True

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

320 # Masked unsigned 32-bit unsigned int.

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

322 arr[mask] = 0

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

324 # Masked string.

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

326 # Masked bytes.

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

328

329 return table

330

331

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

333 """Make an arrow table for testing.

334

335 Parameters

336 ----------

337 include_multidim : `bool`

338 Include multi-dimensional columns.

339 include_masked : `bool`

340 Include masked columns.

341

342 Returns

343 -------

344 arrowTable : `pyarrow.Table`

345 The test table.

346 """

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

348 return astropy_to_arrow(data)

349

350

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

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

353class ParquetFormatterDataFrameTestCase(unittest.TestCase):

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

355

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

357

358 def setUp(self):

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

360 self.root = makeTestTempDir(TESTDIR)

361 config = Config(self.configFile)

362 self.run = "test_run"

363 self.butler = Butler.from_config(

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

365 )

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

367 # inserting dimension data or defining data IDs.

368 self.datasetType = DatasetType(

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

370 )

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

372

373 def tearDown(self):

374 removeTestTempDir(self.root)

375

376 def testSingleIndexDataFrame(self):

377 df1, allColumns = _makeSingleIndexDataFrame(include_masked=True)

378

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

380 # Read the whole DataFrame.

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

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

383 # Read just the column descriptions.

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

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

386 # Read the rowcount.

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

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

389 # Read the schema.

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

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

392 # Read just some columns a few different ways.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

407 # Passing an unrecognized column should be a ValueError.

408 with self.assertRaises(ValueError):

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

410

411 def testSingleIndexDataFrameWithLists(self):

412 df1, allColumns = _makeSingleIndexDataFrame(include_lists=True)

413

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

415 # Read the whole DataFrame.

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

417

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

419 # from lists to arrays.

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

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

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

423

424 def testMultiIndexDataFrame(self):

425 df1 = _makeMultiIndexDataFrame()

426

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

428 # Read the whole DataFrame.

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

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

431 # Read just the column descriptions.

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

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

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

435 # Read the rowcount.

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

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

438 # Read the schema.

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

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

441 # Read just some columns a few different ways.

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

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

444 df4 = self.butler.get(

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

446 )

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

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

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

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

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

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

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

454 # Passing an unrecognized column should be a ValueError.

455 with self.assertRaises(ValueError):

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

457

458 def testSingleIndexDataFrameEmptyString(self):

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

460 df1, _ = _makeSingleIndexDataFrame()

461

462 # Set one of the strings to None

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

464

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

466 # Read the whole DataFrame.

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

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

469

470 def testSingleIndexDataFrameAllEmptyStrings(self):

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

472 column.

473 """

474 df1, _ = _makeSingleIndexDataFrame()

475

476 # Set all of the strings to None

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

478

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

480 # Read the whole DataFrame.

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

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

483

484 def testLegacyDataFrame(self):

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

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

487 functionality.

488 """

489 df1, allColumns = _makeSingleIndexDataFrame()

490

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

492 df1.to_parquet(fname)

493

494 legacy_type = DatasetType(

495 "legacy_dataframe",

496 dimensions=(),

497 storageClass="DataFrame",

498 universe=self.butler.dimensions,

499 )

500 self.butler.registry.registerDatasetType(legacy_type)

501

502 data_id = {}

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

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

505

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

507

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

509

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

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

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

513

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

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

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

517

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

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

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

521

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

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

524 self.assertEqual(rowcount2a, rowcount2b)

525

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

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

528 self.assertEqual(schema2a, schema2b)

529

530 def testDataFrameSchema(self):

531 tab1 = _makeSimpleArrowTable()

532

533 schema = DataFrameSchema.from_arrow(tab1.schema)

534

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

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

537 self.assertNotEqual(schema, "not_a_schema")

538 self.assertEqual(schema, schema)

539

540 tab2 = _makeMultiIndexDataFrame()

541 schema2 = DataFrameSchema(tab2)

542

543 self.assertNotEqual(schema, schema2)

544

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

546 def testWriteSingleIndexDataFrameReadAsAstropyTable(self):

547 df1, allColumns = _makeSingleIndexDataFrame()

548

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

550

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

552

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

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

555

556 # Check reading the columns.

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

558 columns2 = self.butler.get(

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

560 )

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

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

563

564 # Check reading the schema.

565 schema = ArrowAstropySchema(tab2)

566 schema2 = self.butler.get(

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

568 )

569

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

571 # will be changed because of pandas indexing.

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

573 for name in schema.schema.columns:

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

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

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

577

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

579 def testWriteSingleIndexDataFrameWithMaskedColsReadAsAstropyTable(self):

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

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

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

583 # inconsistent.)

584 df1, allColumns = _makeSingleIndexDataFrame(include_masked=True)

585

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

587

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

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

590

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

592 for name in tab2_df.columns:

593 col1 = df1[name]

594 col2 = tab2_df[name]

595

596 if col1.hasnans:

597 notNull = col1.notnull()

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

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

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

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

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

603 else:

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

605

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

607 def testWriteMultiIndexDataFrameReadAsAstropyTable(self):

608 df1 = _makeMultiIndexDataFrame()

609

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

611

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

613

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

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

616 # recommended.

617

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

619 def testWriteAstropyTableWithMaskedColsReadAsSingleIndexDataFrame(self):

620 tab1 = _makeSimpleAstropyTable(include_masked=True)

621

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

623

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

625

626 tab1_df = astropy_to_pandas(tab1)

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

628

629 tab2_astropy = pandas_to_astropy(tab2)

630 for col in tab1.dtype.names:

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

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

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

634

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

636 def testWriteSingleIndexDataFrameReadAsArrowTable(self):

637 df1, allColumns = _makeSingleIndexDataFrame()

638

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

640

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

642

643 tab2_df = arrow_to_pandas(tab2)

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

645

646 # Check reading the columns.

647 columns = list(tab2.schema.names)

648 columns2 = self.butler.get(

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

650 )

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

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

653

654 # Override the component using a dataset type.

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

656 "ArrowColumnList"

657 )

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

659

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

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

662 # selects the component.

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

664 "columns"

665 )

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

667

668 # Check reading the schema.

669 schema = tab2.schema

670 schema2 = self.butler.get(

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

672 )

673

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

675 # information be maintained.

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

677 for name in schema.names:

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

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

680

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

682 def testWriteMultiIndexDataFrameReadAsArrowTable(self):

683 df1 = _makeMultiIndexDataFrame()

684

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

686

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

688

689 tab2_df = arrow_to_pandas(tab2)

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

691

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

693 def testWriteSingleIndexDataFrameReadAsNumpyTable(self):

694 df1, allColumns = _makeSingleIndexDataFrame()

695

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

697

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

699

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

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

702

703 # Check reading the columns.

704 columns = list(tab2.dtype.names)

705 columns2 = self.butler.get(

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

707 )

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

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

710

711 # Check reading the schema.

712 schema = ArrowNumpySchema(tab2.dtype)

713 schema2 = self.butler.get(

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

715 )

716

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

718 # will be changed because of pandas indexing.

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

720 for name in schema.schema.names:

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

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

723

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

725 def testWriteMultiIndexDataFrameReadAsNumpyTable(self):

726 df1 = _makeMultiIndexDataFrame()

727

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

729

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

731

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

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

734 # recommended.

735

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

737 def testWriteSingleIndexDataFrameReadAsNumpyDict(self):

738 df1, allColumns = _makeSingleIndexDataFrame()

739

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

741

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

743

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

745 # The column order is not maintained.

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

747 for col in df1.columns:

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

749

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

751 def testWriteMultiIndexDataFrameReadAsNumpyDict(self):

752 df1 = _makeMultiIndexDataFrame()

753

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

755

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

757

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

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

760 # recommended.

761

762 def testBadDataFrameColumnParquet(self):

763 df1, allColumns = _makeSingleIndexDataFrame()

764

765 # Make a column with mixed type.

766 bad_col1 = [0.0] * len(df1)

767 bad_col1[1] = 0.0 * units.nJy

768 bad_df = df1.copy()

769 bad_df["bad_col1"] = bad_col1

770

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

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

773 with self.assertRaises(RuntimeError):

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

775

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

777 def testWriteReadAstropyTableLossless(self):

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

779

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

781

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

783

784 # Check that minimal provenance was written by default.

785 expected = {

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

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

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

789 }

790

791 self.assertEqual(tab2.meta, expected)

792

793 _checkAstropyTableEquality(tab1, tab2)

794

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

796 def testWriteReadAstropyTableProvenance(self):

797 tab1 = _makeSimpleAstropyTable()

798

799 # Create a ref for provenance.

800 astropy_type = DatasetType(

801 "astropy_parquet",

802 dimensions=(),

803 storageClass="ArrowAstropy",

804 universe=self.butler.dimensions,

805 )

806 self.butler.registry.registerDatasetType(astropy_type)

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

808 quantum_id = uuid.uuid4()

809 provenance = DatasetProvenance(quantum_id=quantum_id)

810 provenance.add_input(input_ref)

811

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

813

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

815

816 expected = {

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

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

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

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

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

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

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

824 }

825 self.assertEqual(tab2.meta, expected)

826

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

828 # that the previous provenance was stripped.

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

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

831

832 # tab2 will have been updated in place.

833 expected = {

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

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

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

837 }

838 self.assertEqual(tab2.meta, expected)

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

840 self.assertEqual(prov_ref, put_ref3)

841 self.assertEqual(null_prov, DatasetProvenance())

842

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

844 def testWriteReadNumpyTableLossless(self):

845 tab1 = _makeSimpleNumpyTable(include_multidim=True)

846

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

848

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

850

851 _checkNumpyTableEquality(tab1, tab2)

852

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

854 def testMaskedNumpy(self):

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

856 tab1_np = arrow_to_numpy(tab1)

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

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

859 col = tab1_np["m_f8"]

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

861

862 # Now without a mask.

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

864 tab1_np = arrow_to_numpy(tab1)

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

866

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

868 def testWriteReadArrowTableLossless(self):

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

870

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

872

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

874

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

876 tab1_np = arrow_to_numpy(tab1)

877 tab2_np = arrow_to_numpy(tab2)

878 for col in tab1.column_names:

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

880

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

882 def testWriteReadNumpyDictLossless(self):

883 tab1 = _makeSimpleNumpyTable(include_multidim=True)

884 dict1 = _numpy_to_numpy_dict(tab1)

885

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

887

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

889

890 _checkNumpyDictEquality(dict1, dict2)

891

892

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

894class InMemoryDataFrameDelegateTestCase(ParquetFormatterDataFrameTestCase):

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

896

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

898

899 def testBadDataFrameColumnParquet(self):

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

901 pass

902

903 def testWriteMultiIndexDataFrameReadAsAstropyTable(self):

904 df1 = _makeMultiIndexDataFrame()

905

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

907

908 with self.assertRaises(ValueError):

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

910

911 def testLegacyDataFrame(self):

912 # This test does not work with an inMemoryDatastore.

913 pass

914

915 def testBadInput(self):

916 df1, _ = _makeSingleIndexDataFrame()

917 delegate = ArrowTableDelegate("DataFrame")

918

919 with self.assertRaises(ValueError):

920 delegate.handleParameters(inMemoryDataset="not_a_dataframe")

921

922 with self.assertRaises(AttributeError):

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

924

925 def testStorageClass(self):

926 df1, allColumns = _makeSingleIndexDataFrame()

927

928 factory = StorageClassFactory()

929 factory.addFromConfig(StorageClassConfig())

930

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

932 # Force the name lookup to do name matching.

933 storageClass._pytype = None

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

935

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

937 # Force the name lookup to do name matching.

938 storageClass._pytype = None

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

940

941

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

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

944class ParquetFormatterArrowAstropyTestCase(unittest.TestCase):

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

946

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

948

949 def setUp(self):

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

951 self.root = makeTestTempDir(TESTDIR)

952 config = Config(self.configFile)

953 self.run = "test_run"

954 self.butler = Butler.from_config(

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

956 )

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

958 # inserting dimension data or defining data IDs.

959 self.datasetType = DatasetType(

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

961 )

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

963

964 def tearDown(self):

965 removeTestTempDir(self.root)

966

967 def testAstropyTable(self):

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

969

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

971 # Read the whole Table.

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

973 _checkAstropyTableEquality(tab1, tab2)

974 # Read the columns.

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

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

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

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

979 # Read the rowcount.

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

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

982 # Read the schema.

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

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

985 # Read just some columns a few different ways.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

1000 # Passing an unrecognized column should be a ValueError.

1001 with self.assertRaises(ValueError):

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

1003

1004 def testAstropyTableBigEndian(self):

1005 tab1 = _makeSimpleAstropyTable(include_bigendian=True)

1006

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

1008 # Read the whole Table.

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

1010 _checkAstropyTableEquality(tab1, tab2, has_bigendian=True)

1011

1012 def testAstropyTableWithMetadata(self):

1013 tab1 = _makeSimpleAstropyTable(include_multidim=True)

1014

1015 meta = {

1016 "meta_a": 5,

1017 "meta_b": 10.0,

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

1019 "meta_d": True,

1020 "meta_e": "string",

1021 }

1022

1023 tab1.meta.update(meta)

1024

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

1026 # Read the whole Table.

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

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

1029 _checkAstropyTableEquality(tab1, tab2)

1030

1031 def testArrowAstropySchema(self):

1032 tab1 = _makeSimpleAstropyTable()

1033 tab1_arrow = astropy_to_arrow(tab1)

1034 schema = ArrowAstropySchema.from_arrow(tab1_arrow.schema)

1035

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

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

1038 self.assertNotEqual(schema, "not_a_schema")

1039 self.assertEqual(schema, schema)

1040

1041 # Test various inequalities

1042 tab2 = tab1.copy()

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

1044 schema2 = ArrowAstropySchema(tab2)

1045 self.assertNotEqual(schema2, schema)

1046

1047 tab2 = tab1.copy()

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

1049 schema2 = ArrowAstropySchema(tab2)

1050 self.assertNotEqual(schema2, schema)

1051

1052 tab2 = tab1.copy()

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

1054 schema2 = ArrowAstropySchema(tab2)

1055 self.assertNotEqual(schema2, schema)

1056

1057 tab2 = tab1.copy()

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

1059 schema2 = ArrowAstropySchema(tab2)

1060 self.assertNotEqual(schema2, schema)

1061

1062 def testAstropyParquet(self):

1063 tab1 = _makeSimpleAstropyTable()

1064

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

1066 del tab1["dtn"]

1067 del tab1["dtu"]

1068

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

1070 tab1.write(fname)

1071

1072 astropy_type = DatasetType(

1073 "astropy_parquet",

1074 dimensions=(),

1075 storageClass="ArrowAstropy",

1076 universe=self.butler.dimensions,

1077 )

1078 self.butler.registry.registerDatasetType(astropy_type)

1079

1080 data_id = {}

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

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

1083

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

1085

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

1087

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

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

1090 _checkAstropyTableEquality(tab2a, tab2b)

1091

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

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

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

1095 for i, name in enumerate(columns2a):

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

1097

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

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

1100 self.assertEqual(rowcount2a, rowcount2b)

1101

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

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

1104 self.assertEqual(schema2a, schema2b)

1105

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

1107 def testWriteAstropyReadAsArrowTable(self):

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

1109 tab1 = _makeSimpleAstropyTable(include_masked=True)

1110

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

1112

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

1114

1115 tab2_astropy = arrow_to_astropy(tab2)

1116 _checkAstropyTableEquality(tab1, tab2_astropy)

1117

1118 # Check reading the columns.

1119 columns = tab2.schema.names

1120 columns2 = self.butler.get(

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

1122 )

1123 self.assertEqual(columns2, columns)

1124

1125 # Check reading the schema.

1126 schema = tab2.schema

1127 schema2 = self.butler.get(

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

1129 )

1130

1131 self.assertEqual(schema, schema2)

1132

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

1134 def testWriteAstropyReadAsDataFrame(self):

1135 tab1 = _makeSimpleAstropyTable()

1136

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

1138

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

1140

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

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

1143

1144 tab1_df = tab1.to_pandas()

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

1146

1147 # Check reading the columns.

1148 columns = tab2.columns

1149 columns2 = self.butler.get(

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

1151 )

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

1153

1154 # Check reading the schema.

1155 schema = DataFrameSchema(tab2)

1156 schema2 = self.butler.get(

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

1158 )

1159

1160 self.assertEqual(schema2, schema)

1161

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

1163 def testWriteAstropyWithMaskedColsReadAsDataFrame(self):

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

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

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

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

1168 tab1 = _makeSimpleAstropyTable(include_masked=True)

1169

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

1171

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

1173

1174 tab1_df = astropy_to_pandas(tab1)

1175

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

1177 for name in tab2.columns:

1178 col1 = tab1_df[name]

1179 col2 = tab2[name]

1180

1181 if col1.hasnans:

1182 notNull = col1.notnull()

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

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

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

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

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

1188 else:

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

1190

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

1192 def testWriteSingleIndexDataFrameWithMaskedColsReadAsAstropyTable(self):

1193 df1, allColumns = _makeSingleIndexDataFrame(include_masked=True)

1194

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

1196

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

1198

1199 df1_tab = pandas_to_astropy(df1)

1200

1201 _checkAstropyTableEquality(df1_tab, tab2)

1202

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

1204 def testWriteAstropyReadAsNumpyTable(self):

1205 tab1 = _makeSimpleAstropyTable()

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

1207

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

1209

1210 # This is tricky because it loses the units.

1211 tab2_astropy = atable.Table(tab2)

1212

1213 _checkAstropyTableEquality(tab1, tab2_astropy, skip_units=True)

1214

1215 # Check reading the columns.

1216 columns = list(tab2.dtype.names)

1217 columns2 = self.butler.get(

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

1219 )

1220 self.assertEqual(columns2, columns)

1221

1222 # Check reading the schema.

1223 schema = ArrowNumpySchema(tab2.dtype)

1224 schema2 = self.butler.get(

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

1226 )

1227

1228 self.assertEqual(schema2, schema)

1229

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

1231 def testWriteAstropyReadAsNumpyDict(self):

1232 tab1 = _makeSimpleAstropyTable()

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

1234

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

1236

1237 # This is tricky because it loses the units.

1238 tab2_astropy = atable.Table(tab2)

1239

1240 _checkAstropyTableEquality(tab1, tab2_astropy, skip_units=True)

1241

1242 def testBadAstropyColumnParquet(self):

1243 tab1 = _makeSimpleAstropyTable()

1244

1245 # Make a column with mixed type.

1246 bad_col1 = [0.0] * len(tab1)

1247 bad_col1[1] = 0.0 * units.nJy

1248 bad_tab = tab1.copy()

1249 bad_tab["bad_col1"] = bad_col1

1250

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

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

1253 with self.assertRaises(RuntimeError):

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

1255

1256 # Make a column with ragged size.

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

1258 bad_col2[1] = [0, 0]

1259 bad_tab = tab1.copy()

1260 bad_tab["bad_col2"] = bad_col2

1261

1262 with self.assertRaises(RuntimeError):

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

1264

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

1266 def testWriteAstropyTableWithPandasIndexHint(self, testStrip=True):

1267 tab1 = _makeSimpleAstropyTable()

1268

1269 add_pandas_index_to_astropy(tab1, "index")

1270

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

1272

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

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

1275

1276 self.assertIn(ASTROPY_PANDAS_INDEX_KEY, tab2.meta)

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

1278

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

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

1281

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

1283

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

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

1286 _ = self.butler.get(

1287 self.datasetType,

1288 dataId={},

1289 storageClass="DataFrame",

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

1291 )

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

1293

1294 if testStrip:

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

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

1297

1298 self.assertNotIn(ASTROPY_PANDAS_INDEX_KEY, tab5.meta)

1299

1300 with self.assertRaises(ValueError):

1301 add_pandas_index_to_astropy(tab1, "not_a_column")

1302

1303

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

1305class InMemoryArrowAstropyDelegateTestCase(ParquetFormatterArrowAstropyTestCase):

1306 """Tests for InMemoryDatastore, using ArrowTableDelegate with

1307 AstropyTable.

1308 """

1309

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

1311

1312 def testAstropyParquet(self):

1313 # This test does not work with an inMemoryDatastore.

1314 pass

1315

1316 def testBadAstropyColumnParquet(self):

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

1318 pass

1319

1320 def testBadInput(self):

1321 tab1 = _makeSimpleAstropyTable()

1322 delegate = ArrowTableDelegate("ArrowAstropy")

1323

1324 with self.assertRaises(ValueError):

1325 delegate.handleParameters(inMemoryDataset="not_an_astropy_table")

1326

1327 with self.assertRaises(NotImplementedError):

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

1329

1330 with self.assertRaises(AttributeError):

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

1332

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

1334 def testWriteAstropyTableWithPandasIndexHint(self):

1335 super().testWriteAstropyTableWithPandasIndexHint(testStrip=False)

1336

1337

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

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

1340class ParquetFormatterArrowNumpyTestCase(unittest.TestCase):

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

1342

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

1344

1345 def setUp(self):

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

1347 self.root = makeTestTempDir(TESTDIR)

1348 config = Config(self.configFile)

1349 self.butler = Butler.from_config(

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

1351 )

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

1353 # inserting dimension data or defining data IDs.

1354 self.datasetType = DatasetType(

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

1356 )

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

1358

1359 def tearDown(self):

1360 removeTestTempDir(self.root)

1361

1362 def testNumpyTable(self):

1363 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1364

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

1366 # Read the whole Table.

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

1368 _checkNumpyTableEquality(tab1, tab2)

1369 # Read the columns.

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

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

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

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

1374 # Read the rowcount.

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

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

1377 # Read the schema.

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

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

1380 # Read just some columns a few different ways.

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

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

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

1384 _checkNumpyTableEquality(

1385 tab1[

1386 [

1387 "a",

1388 ]

1389 ],

1390 tab4,

1391 )

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

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

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

1395 _checkNumpyTableEquality(

1396 tab1[

1397 [

1398 "ddd",

1399 ]

1400 ],

1401 tab6,

1402 )

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

1404 _checkNumpyTableEquality(

1405 tab1[

1406 [

1407 "a",

1408 ]

1409 ],

1410 tab7,

1411 )

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

1413 _checkNumpyTableEquality(

1414 tab1[

1415 [

1416 "ddd",

1417 "dtn",

1418 "dtu",

1419 "a",

1420 ]

1421 ],

1422 tab8,

1423 )

1424 # Passing an unrecognized column should be a ValueError.

1425 with self.assertRaises(ValueError):

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

1427

1428 def testNumpyTableBigEndian(self):

1429 tab1 = _makeSimpleNumpyTable(include_bigendian=True)

1430

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

1432 # Read the whole Table.

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

1434 _checkNumpyTableEquality(tab1, tab2, has_bigendian=True)

1435

1436 def testArrowNumpySchema(self):

1437 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1438 tab1_arrow = numpy_to_arrow(tab1)

1439 schema = ArrowNumpySchema.from_arrow(tab1_arrow.schema)

1440

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

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

1443 self.assertNotEqual(schema, "not_a_schema")

1444 self.assertEqual(schema, schema)

1445

1446 # Test inequality

1447 tab2 = tab1.copy()

1448 names = list(tab2.dtype.names)

1449 names[0] = "index2"

1450 tab2.dtype.names = names

1451 schema2 = ArrowNumpySchema(tab2.dtype)

1452 self.assertNotEqual(schema2, schema)

1453

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

1455 def testNumpyDictConversions(self):

1456 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1457

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

1459 tab1_arrow = numpy_to_arrow(tab1)

1460 tab1_dict = arrow_to_numpy_dict(tab1_arrow)

1461 tab1_dict_arrow = numpy_dict_to_arrow(tab1_dict)

1462

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

1464 self.assertEqual(tab1_arrow, tab1_dict_arrow)

1465

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

1467 def testWriteNumpyTableReadAsArrowTable(self):

1468 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1469

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

1471

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

1473

1474 tab2_numpy = arrow_to_numpy(tab2)

1475

1476 _checkNumpyTableEquality(tab1, tab2_numpy)

1477

1478 # Check reading the columns.

1479 columns = tab2.schema.names

1480 columns2 = self.butler.get(

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

1482 )

1483 self.assertEqual(columns2, columns)

1484

1485 # Check reading the schema.

1486 schema = tab2.schema

1487 schema2 = self.butler.get(

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

1489 )

1490 self.assertEqual(schema2, schema)

1491

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

1493 def testWriteNumpyTableReadAsDataFrame(self):

1494 tab1 = _makeSimpleNumpyTable()

1495

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

1497

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

1499

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

1501 # and changes the datatype of the string column.

1502

1503 tab1_df = pd.DataFrame(tab1)

1504

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

1506

1507 # Check reading the columns.

1508 columns = tab2.columns

1509 columns2 = self.butler.get(

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

1511 )

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

1513

1514 # Check reading the schema.

1515 schema = DataFrameSchema(tab2)

1516 schema2 = self.butler.get(

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

1518 )

1519

1520 self.assertEqual(schema2, schema)

1521

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

1523 def testWriteNumpyTableReadAsAstropyTable(self):

1524 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1525

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

1527

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

1529 tab2_numpy = tab2.as_array()

1530

1531 _checkNumpyTableEquality(tab1, tab2_numpy)

1532

1533 # Check reading the columns.

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

1535 columns2 = self.butler.get(

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

1537 )

1538 self.assertEqual(columns2, columns)

1539

1540 # Check reading the schema.

1541 schema = ArrowAstropySchema(tab2)

1542 schema2 = self.butler.get(

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

1544 )

1545

1546 self.assertEqual(schema2, schema)

1547

1548 def testWriteNumpyTableReadAsNumpyDict(self):

1549 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1550

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

1552

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

1554 tab2_numpy = _numpy_dict_to_numpy(tab2)

1555

1556 _checkNumpyTableEquality(tab1, tab2_numpy)

1557

1558 def testBadNumpyColumnParquet(self):

1559 tab1 = _makeSimpleAstropyTable()

1560

1561 # Make a column with mixed type.

1562 bad_col1 = [0.0] * len(tab1)

1563 bad_col1[1] = 0.0 * units.nJy

1564 bad_tab = tab1.copy()

1565 bad_tab["bad_col1"] = bad_col1

1566

1567 bad_tab_np = bad_tab.as_array()

1568

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

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

1571 with self.assertRaises(RuntimeError):

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

1573

1574 # Make a column with ragged size.

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

1576 bad_col2[1] = [0, 0]

1577 bad_tab = tab1.copy()

1578 bad_tab["bad_col2"] = bad_col2

1579

1580 bad_tab_np = bad_tab.as_array()

1581

1582 with self.assertRaises(RuntimeError):

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

1584

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

1586 def testWriteReadAstropyTableLossless(self):

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

1588

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

1590

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

1592

1593 _checkAstropyTableEquality(tab1, tab2)

1594

1595

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

1597class InMemoryArrowNumpyDelegateTestCase(ParquetFormatterArrowNumpyTestCase):

1598 """Tests for InMemoryDatastore, using ArrowTableDelegate with

1599 Numpy table.

1600 """

1601

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

1603

1604 def testBadNumpyColumnParquet(self):

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

1606 pass

1607

1608 def testBadInput(self):

1609 tab1 = _makeSimpleNumpyTable()

1610 delegate = ArrowTableDelegate("ArrowNumpy")

1611

1612 with self.assertRaises(ValueError):

1613 delegate.handleParameters(inMemoryDataset="not_a_numpy_table")

1614

1615 with self.assertRaises(NotImplementedError):

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

1617

1618 with self.assertRaises(AttributeError):

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

1620

1621 def testStorageClass(self):

1622 tab1 = _makeSimpleNumpyTable()

1623

1624 factory = StorageClassFactory()

1625 factory.addFromConfig(StorageClassConfig())

1626

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

1628 # Force the name lookup to do name matching.

1629 storageClass._pytype = None

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

1631

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

1633 # Force the name lookup to do name matching.

1634 storageClass._pytype = None

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

1636

1637

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

1639class ParquetFormatterArrowTableTestCase(unittest.TestCase):

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

1641

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

1643

1644 def setUp(self):

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

1646 self.root = makeTestTempDir(TESTDIR)

1647 config = Config(self.configFile)

1648 self.butler = Butler.from_config(

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

1650 )

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

1652 # inserting dimension data or defining data IDs.

1653 self.datasetType = DatasetType(

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

1655 )

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

1657

1658 def tearDown(self):

1659 removeTestTempDir(self.root)

1660

1661 def testArrowTable(self):

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

1663

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

1665 # Read the whole Table.

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

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

1668 # comparisons.

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

1670 tab1_np = arrow_to_numpy(tab1)

1671 tab2_np = arrow_to_numpy(tab2)

1672 for col in tab1.column_names:

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

1674 # Read the columns.

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

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

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

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

1679 # Read the rowcount.

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

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

1682 # Read the schema.

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

1684 self.assertEqual(schema, tab1.schema)

1685 # Read just some columns a few different ways.

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

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

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

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

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

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

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

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

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

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

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

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

1698 # Passing an unrecognized column should be a ValueError.

1699 with self.assertRaises(ValueError):

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

1701

1702 def testEmptyArrowTable(self):

1703 data = _makeSimpleNumpyTable()

1704 type_list = _numpy_dtype_to_arrow_types(data.dtype)

1705

1706 schema = pa.schema(type_list)

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

1708

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

1710

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

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

1713 self.assertEqual(tab2, tab1)

1714

1715 tab1_numpy = arrow_to_numpy(tab1)

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

1717 tab1_numpy_arrow = numpy_to_arrow(tab1_numpy)

1718 self.assertEqual(tab1_numpy_arrow, tab1)

1719

1720 tab1_pandas = arrow_to_pandas(tab1)

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

1722 tab1_pandas_arrow = pandas_to_arrow(tab1_pandas)

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

1724 # through empty pandas dataframes.

1725 self.assertEqual(

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

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

1728 )

1729

1730 tab1_astropy = arrow_to_astropy(tab1)

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

1732 tab1_astropy_arrow = astropy_to_arrow(tab1_astropy)

1733 self.assertEqual(tab1_astropy_arrow, tab1)

1734

1735 def testEmptyArrowTableMultidim(self):

1736 data = _makeSimpleNumpyTable(include_multidim=True)

1737 type_list = _numpy_dtype_to_arrow_types(data.dtype)

1738

1739 md = {}

1740 for name in data.dtype.names:

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

1742

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

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

1745

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

1747

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

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

1750 self.assertEqual(tab2, tab1)

1751

1752 tab1_numpy = arrow_to_numpy(tab1)

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

1754 tab1_numpy_arrow = numpy_to_arrow(tab1_numpy)

1755 self.assertEqual(tab1_numpy_arrow, tab1)

1756

1757 tab1_astropy = arrow_to_astropy(tab1)

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

1759 tab1_astropy_arrow = astropy_to_arrow(tab1_astropy)

1760 self.assertEqual(tab1_astropy_arrow, tab1)

1761

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

1763 def testWriteArrowTableReadAsSingleIndexDataFrame(self):

1764 df1, allColumns = _makeSingleIndexDataFrame()

1765

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

1767

1768 # Read back out as a dataframe.

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

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

1771

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

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

1774 df3 = arrow_to_pandas(tab3)

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

1776

1777 # Check reading the columns.

1778 columns = df2.reset_index().columns

1779 columns2 = self.butler.get(

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

1781 )

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

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

1784

1785 # Check reading the schema.

1786 schema = DataFrameSchema(df1)

1787 schema2 = self.butler.get(

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

1789 )

1790 self.assertEqual(schema2, schema)

1791

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

1793 def testWriteArrowTableReadAsMultiIndexDataFrame(self):

1794 df1 = _makeMultiIndexDataFrame()

1795

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

1797

1798 # Read back out as a dataframe.

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

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

1801

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

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

1804 df3 = arrow_to_pandas(atab3)

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

1806

1807 # Check reading the columns.

1808 columns = df2.columns

1809 columns2 = self.butler.get(

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

1811 )

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

1813

1814 # Check reading the schema.

1815 schema = DataFrameSchema(df1)

1816 schema2 = self.butler.get(

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

1818 )

1819 self.assertEqual(schema2, schema)

1820

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

1822 def testWriteArrowTableReadAsAstropyTable(self):

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

1824

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

1826

1827 # Read back out as an astropy table.

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

1829 _checkAstropyTableEquality(tab1, tab2)

1830

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

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

1833 tab3 = arrow_to_astropy(atab3)

1834 _checkAstropyTableEquality(tab1, tab3)

1835

1836 # Check reading the columns.

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

1838 columns2 = self.butler.get(

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

1840 )

1841 self.assertEqual(columns2, columns)

1842

1843 # Check reading the schema.

1844 schema = ArrowAstropySchema(tab1)

1845 schema2 = self.butler.get(

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

1847 )

1848 self.assertEqual(schema2, schema)

1849

1850 # Check the schema conversions and units.

1851 arrow_schema = schema.to_arrow_schema()

1852 for name in arrow_schema.names:

1853 field_metadata = arrow_schema.field(name).metadata

1854 if (

1855 b"description" in field_metadata

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

1857 ):

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

1859 else:

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

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

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

1863

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

1865 def testWriteArrowTableReadAsNumpyTable(self):

1866 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1867

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

1869

1870 # Read back out as a numpy table.

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

1872 _checkNumpyTableEquality(tab1, tab2)

1873

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

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

1876 tab3 = arrow_to_numpy(atab3)

1877 _checkNumpyTableEquality(tab1, tab3)

1878

1879 # Check reading the columns.

1880 columns = list(tab2.dtype.names)

1881 columns2 = self.butler.get(

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

1883 )

1884 self.assertEqual(columns2, columns)

1885

1886 # Check reading the schema.

1887 schema = ArrowNumpySchema(tab1.dtype)

1888 schema2 = self.butler.get(

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

1890 )

1891 self.assertEqual(schema2, schema)

1892

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

1894 def testWriteArrowTableReadAsNumpyDict(self):

1895 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1896

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

1898

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

1900 tab2_numpy = _numpy_dict_to_numpy(tab2)

1901 _checkNumpyTableEquality(tab1, tab2_numpy)

1902

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

1904 def testWriteReadAstropyTableLossless(self):

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

1906

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

1908

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

1910

1911 _checkAstropyTableEquality(tab1, tab2)

1912

1913

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

1915class InMemoryArrowTableDelegateTestCase(ParquetFormatterArrowTableTestCase):

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

1917

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

1919

1920 def testBadInput(self):

1921 tab1 = _makeSimpleArrowTable()

1922 delegate = ArrowTableDelegate("ArrowTable")

1923

1924 with self.assertRaises(ValueError):

1925 delegate.handleParameters(inMemoryDataset="not_an_arrow_table")

1926

1927 with self.assertRaises(NotImplementedError):

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

1929

1930 with self.assertRaises(AttributeError):

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

1932

1933 def testStorageClass(self):

1934 tab1 = _makeSimpleArrowTable()

1935

1936 factory = StorageClassFactory()

1937 factory.addFromConfig(StorageClassConfig())

1938

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

1940 # Force the name lookup to do name matching.

1941 storageClass._pytype = None

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

1943

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

1945 # Force the name lookup to do name matching.

1946 storageClass._pytype = None

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

1948

1949

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

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

1952class ParquetFormatterArrowNumpyDictTestCase(unittest.TestCase):

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

1954

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

1956

1957 def setUp(self):

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

1959 self.root = makeTestTempDir(TESTDIR)

1960 config = Config(self.configFile)

1961 self.butler = Butler.from_config(

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

1963 )

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

1965 # inserting dimension data or defining data IDs.

1966 self.datasetType = DatasetType(

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

1968 )

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

1970

1971 def tearDown(self):

1972 removeTestTempDir(self.root)

1973

1974 def testNumpyDict(self):

1975 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1976 dict1 = _numpy_to_numpy_dict(tab1)

1977

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

1979 # Read the whole table.

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

1981 _checkNumpyDictEquality(dict1, dict2)

1982 # Read the columns.

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

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

1985 for name in dict1:

1986 self.assertIn(name, columns2)

1987 # Read the rowcount.

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

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

1990 # Read the schema.

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

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

1993 # Read just some columns a few different ways.

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

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

1996 _checkNumpyDictEquality(subdict, tab3)

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

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

1999 _checkNumpyDictEquality(subdict, tab4)

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

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

2002 _checkNumpyDictEquality(subdict, tab5)

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

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

2005 _checkNumpyDictEquality(subdict, tab6)

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

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

2008 _checkNumpyDictEquality(subdict, tab7)

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

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

2011 _checkNumpyDictEquality(subdict, tab8)

2012 # Passing an unrecognized column should be a ValueError.

2013 with self.assertRaises(ValueError):

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

2015

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

2017 def testWriteNumpyDictReadAsArrowTable(self):

2018 tab1 = _makeSimpleNumpyTable(include_multidim=True)

2019 dict1 = _numpy_to_numpy_dict(tab1)

2020

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

2022

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

2024

2025 tab2_dict = arrow_to_numpy_dict(tab2)

2026

2027 _checkNumpyDictEquality(dict1, tab2_dict)

2028

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

2030 def testWriteNumpyDictReadAsDataFrame(self):

2031 tab1 = _makeSimpleNumpyTable()

2032 dict1 = _numpy_to_numpy_dict(tab1)

2033

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

2035

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

2037

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

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

2040 # changes the datatype of the string column.

2041 tab1_df = pd.DataFrame(tab1)

2042

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

2044 for col in tab1_df.columns:

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

2046

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

2048 def testWriteNumpyDictReadAsAstropyTable(self):

2049 tab1 = _makeSimpleNumpyTable(include_multidim=True)

2050 dict1 = _numpy_to_numpy_dict(tab1)

2051

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

2053

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

2055 tab2_dict = _astropy_to_numpy_dict(tab2)

2056

2057 _checkNumpyDictEquality(dict1, tab2_dict)

2058

2059 def testWriteNumpyDictReadAsNumpyTable(self):

2060 tab1 = _makeSimpleNumpyTable(include_multidim=True)

2061 dict1 = _numpy_to_numpy_dict(tab1)

2062

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

2064

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

2066 tab2_dict = _numpy_to_numpy_dict(tab2)

2067

2068 _checkNumpyDictEquality(dict1, tab2_dict)

2069

2070 def testWriteNumpyDictBad(self):

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

2072 with self.assertRaises(RuntimeError):

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

2074

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

2076 with self.assertRaises(RuntimeError):

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

2078

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

2080 with self.assertRaises(RuntimeError):

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

2082

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

2084 with self.assertRaises(RuntimeError):

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

2086

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

2088 def testWriteReadAstropyTableLossless(self):

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

2090

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

2092

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

2094

2095 _checkAstropyTableEquality(tab1, tab2)

2096

2097

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

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

2100class InMemoryNumpyDictDelegateTestCase(ParquetFormatterArrowNumpyDictTestCase):

2101 """Tests for InMemoryDatastore, using ArrowTableDelegate with

2102 Numpy dict.

2103 """

2104

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

2106

2107 def testWriteNumpyDictBad(self):

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

2109 pass

2110

2111

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

2113class ParquetFormatterArrowSchemaTestCase(unittest.TestCase):

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

2115

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

2117

2118 def setUp(self):

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

2120 self.root = makeTestTempDir(TESTDIR)

2121 config = Config(self.configFile)

2122 self.butler = Butler.from_config(

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

2124 )

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

2126 # inserting dimension data or defining data IDs.

2127 self.datasetType = DatasetType(

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

2129 )

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

2131

2132 def tearDown(self):

2133 removeTestTempDir(self.root)

2134

2135 def _makeTestSchema(self):

2136 schema = pa.schema(

2137 [

2138 pa.field(

2139 "int32",

2140 pa.int32(),

2141 nullable=False,

2142 metadata={

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

2144 "unit": "",

2145 },

2146 ),

2147 pa.field(

2148 "int64",

2149 pa.int64(),

2150 nullable=False,

2151 metadata={

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

2153 "unit": "",

2154 },

2155 ),

2156 pa.field(

2157 "uint64",

2158 pa.uint64(),

2159 nullable=False,

2160 metadata={

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

2162 "unit": "",

2163 },

2164 ),

2165 pa.field(

2166 "float32",

2167 pa.float32(),

2168 nullable=False,

2169 metadata={

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

2171 "unit": "count",

2172 },

2173 ),

2174 pa.field(

2175 "float64",

2176 pa.float64(),

2177 nullable=False,

2178 metadata={

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

2180 "unit": "nJy",

2181 },

2182 ),

2183 pa.field(

2184 "fixed_size_list",

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

2186 nullable=False,

2187 metadata={

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

2189 "unit": "nJy",

2190 },

2191 ),

2192 pa.field(

2193 "variable_size_list",

2194 pa.list_(pa.float64()),

2195 nullable=False,

2196 metadata={

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

2198 "unit": "nJy",

2199 },

2200 ),

2201 # One of these fields will have no description.

2202 pa.field(

2203 "string",

2204 pa.string(),

2205 nullable=False,

2206 metadata={

2207 "unit": "",

2208 },

2209 ),

2210 # One of these fields will have no metadata.

2211 pa.field(

2212 "binary",

2213 pa.binary(),

2214 nullable=False,

2215 ),

2216 ]

2217 )

2218

2219 return schema

2220

2221 def testArrowSchema(self):

2222 schema1 = self._makeTestSchema()

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

2224

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

2226 self.assertEqual(schema2, schema1)

2227

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

2229 def testWriteArrowSchemaReadAsDataFrameSchema(self):

2230 schema1 = self._makeTestSchema()

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

2232

2233 df_schema1 = DataFrameSchema.from_arrow(schema1)

2234

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

2236 self.assertEqual(df_schema2, df_schema1)

2237

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

2239 def testWriteArrowSchemaReadAsArrowAstropySchema(self):

2240 schema1 = self._makeTestSchema()

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

2242

2243 ap_schema1 = ArrowAstropySchema.from_arrow(schema1)

2244

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

2246 self.assertEqual(ap_schema2, ap_schema1)

2247

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

2249 for name in schema1.names:

2250 field_metadata = schema1.field(name).metadata

2251 if field_metadata is None:

2252 continue

2253 if (

2254 b"description" in field_metadata

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

2256 ):

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

2258 else:

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

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

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

2262

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

2264 def testWriteArrowSchemaReadAsArrowNumpySchema(self):

2265 schema1 = self._makeTestSchema()

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

2267

2268 np_schema1 = ArrowNumpySchema.from_arrow(schema1)

2269

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

2271 self.assertEqual(np_schema2, np_schema1)

2272

2273

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

2275class InMemoryArrowSchemaDelegateTestCase(ParquetFormatterArrowSchemaTestCase):

2276 """Tests for InMemoryDatastore and ArrowSchema."""

2277

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

2279

2280

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

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

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

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

2285class ParquetFormatterArrowTableS3TestCase(unittest.TestCase):

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

2287

2288 # Code is adapted from test_butler.py

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

2290 fullConfigKey = None

2291 validationCanFail = True

2292

2293 bucketName = "anybucketname"

2294

2295 root = "butlerRoot/"

2296

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

2298

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

2300

2301 registryStr = "/gen3.sqlite3"

2302

2303 mock_aws = mock_aws()

2304

2305 def setUp(self):

2306 self.root = makeTestTempDir(TESTDIR)

2307

2308 config = Config(self.configFile)

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

2310 self.bucketName = uri.netloc

2311

2312 # Enable S3 mocking of tests.

2313 self.enterContext(clean_test_environment_for_s3())

2314 self.mock_aws.start()

2315

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

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

2318

2319 # need local folder to store registry database

2320 self.reg_dir = makeTestTempDir(TESTDIR)

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

2322

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

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

2325 s3 = boto3.resource("s3")

2326 s3.create_bucket(Bucket=self.bucketName)

2327

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

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

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

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

2332

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

2334

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

2336 # inserting dimension data or defining data IDs.

2337 self.datasetType = DatasetType(

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

2339 )

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

2341

2342 def tearDown(self):

2343 s3 = boto3.resource("s3")

2344 bucket = s3.Bucket(self.bucketName)

2345 try:

2346 bucket.objects.all().delete()

2347 except botocore.exceptions.ClientError as e:

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

2349 # the key was not reachable - pass

2350 pass

2351 else:

2352 raise

2353

2354 bucket = s3.Bucket(self.bucketName)

2355 bucket.delete()

2356

2357 # Stop the S3 mock.

2358 self.mock_aws.stop()

2359

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

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

2362

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

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

2365

2366 def testArrowTableS3(self):

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

2368

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

2370

2371 # Read the whole Table.

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

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

2374 # comparisons.

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

2376 tab1_np = arrow_to_numpy(tab1)

2377 tab2_np = arrow_to_numpy(tab2)

2378 for col in tab1.column_names:

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

2380 # Read the columns.

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

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

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

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

2385 # Read the rowcount.

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

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

2388 # Read the schema.

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

2390 self.assertEqual(schema, tab1.schema)

2391 # Read just some columns a few different ways.

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

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

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

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

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

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

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

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

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

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

2402 # Passing an unrecognized column should be a ValueError.

2403 with self.assertRaises(ValueError):

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

2405

2406

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

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

2409class ComputeRowGroupSizeTestCase(unittest.TestCase):

2410 """Tests for compute_row_group_size."""

2411

2412 def testRowGroupSizeNoMetadata(self):

2413 numpyTable = _makeSimpleNumpyTable(include_multidim=True)

2414

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

2416 # that adds metadata.

2417 type_list = _numpy_dtype_to_arrow_types(numpyTable.dtype)

2418 schema = pa.schema(type_list)

2419 arrays = _numpy_style_arrays_to_arrow_arrays(

2420 numpyTable.dtype,

2421 len(numpyTable),

2422 numpyTable,

2423 schema,

2424 )

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

2426

2427 row_group_size = compute_row_group_size(arrowTable.schema)

2428

2429 self.assertGreater(row_group_size, 1_000_000)

2430 self.assertLess(row_group_size, 2_000_000)

2431

2432 def testRowGroupSizeWithMetadata(self):

2433 numpyTable = _makeSimpleNumpyTable(include_multidim=True)

2434

2435 arrowTable = numpy_to_arrow(numpyTable)

2436

2437 row_group_size = compute_row_group_size(arrowTable.schema)

2438

2439 self.assertGreater(row_group_size, 1_000_000)

2440 self.assertLess(row_group_size, 2_000_000)

2441

2442 def testRowGroupSizeTinyTable(self):

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

2444

2445 arrowTable = numpy_to_arrow(numpyTable)

2446

2447 row_group_size = compute_row_group_size(arrowTable.schema)

2448

2449 self.assertGreater(row_group_size, 1_000_000)

2450

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

2452 def testRowGroupSizeDataFrameWithLists(self):

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

2454 arrowTable = pandas_to_arrow(df)

2455 row_group_size = compute_row_group_size(arrowTable.schema)

2456

2457 self.assertGreater(row_group_size, 1_000_000)

2458

2459

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

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

2462

2463 Parameters

2464 ----------

2465 table1 : `astropy.table.Table`

2466 table2 : `astropy.table.Table`

2467 skip_units : `bool`

2468 has_bigendian : `bool`

2469 """

2470 if not has_bigendian:

2471 assert table1.dtype == table2.dtype

2472 else:

2473 for name in table1.dtype.names:

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

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

2476

2477 # Strip provenance before comparison.

2478 DatasetProvenance.strip_provenance_from_flat_dict(table1.meta)

2479 DatasetProvenance.strip_provenance_from_flat_dict(table2.meta)

2480 assert table1.meta == table2.meta

2481 if not skip_units:

2482 for name in table1.columns:

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

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

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

2486

2487 for name in table1.columns:

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

2489 has_masked = False

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

2491 c1 = table1[name].filled()

2492 has_masked = True

2493 else:

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

2495 if has_masked:

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

2497 c2 = table2[name].filled()

2498 else:

2499 assert not isinstance(table2[name], atable.column.MaskedColumn)

2500 c2 = np.array(table2[name])

2501 np.testing.assert_array_equal(c1, c2)

2502 # If we have a masked column then we test the underlying data.

2503 if has_masked:

2504 np.testing.assert_array_equal(np.array(c1), np.array(c2))

2505 np.testing.assert_array_equal(table1[name].mask, table2[name].mask)

2506

2507

2508def _checkNumpyTableEquality(table1, table2, has_bigendian=False):

2509 """Check if two numpy tables have the same columns/values

2510

2511 Parameters

2512 ----------

2513 table1 : `numpy.ndarray`

2514 table2 : `numpy.ndarray`

2515 has_bigendian : `bool`

2516 """

2517 assert table1.dtype.names == table2.dtype.names

2518 for name in table1.dtype.names:

2519 if not has_bigendian:

2520 assert table1.dtype[name] == table2.dtype[name]

2521 else:

2522 # Only check type matches, force to little-endian.

2523 assert table1.dtype[name].newbyteorder(">") == table2.dtype[name].newbyteorder(">")

2524 assert np.all(table1 == table2)

2525

2526

2527def _checkNumpyDictEquality(dict1, dict2):

2528 """Check if two numpy dicts have the same columns/values.

2529

2530 Parameters

2531 ----------

2532 dict1 : `dict` [`str`, `np.ndarray`]

2533 dict2 : `dict` [`str`, `np.ndarray`]

2534 """

2535 assert set(dict1.keys()) == set(dict2.keys())

2536 for name in dict1:

2537 assert dict1[name].dtype == dict2[name].dtype

2538 assert np.all(dict1[name] == dict2[name])

2539

2540

2541if __name__ == "__main__":

2542 unittest.main()