Coverage for tests/test_parquet.py: 98%

1# This file is part of daf_butler.

3# Developed for the LSST Data Management System.

4# This product includes software developed by the LSST Project

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

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

7# for details of code ownership.

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

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

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

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

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

14#

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

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

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

18# (at your option) any later version.

19#

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

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

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

23# GNU General Public License for more details.

24#

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

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

28"""Tests for ParquetFormatter.

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

31"""

33import datetime

34import os

35import posixpath

36import shutil

37import unittest

38import uuid

40try:

41 import pyarrow as pa

42except ImportError:

43 pa = None

44try:

45 import astropy.table as atable

46 from astropy import units

47except ImportError:

48 atable = None

49try:

50 import numpy as np

51except ImportError:

52 np = None

53try:

54 import pandas as pd

55except ImportError:

56 pd = None

58try:

59 import boto3

60 import botocore

62 from lsst.resources.s3utils import clean_test_environment_for_s3

64 try:

65 from moto import mock_aws # v5

66 except ImportError:

67 from moto import mock_s3 as mock_aws

68except ImportError:

69 boto3 = None

71try:

72 import fsspec

73except ImportError:

74 fsspec = None

76try:

77 import s3fs

78except ImportError:

79 s3fs = None

82from lsst.daf.butler import (

83 Butler,

84 Config,

85 DatasetProvenance,

86 DatasetRef,

87 DatasetType,

88 FileDataset,

89 StorageClassConfig,

90 StorageClassFactory,

91)

92from lsst.resources import ResourcePath

94try:

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

96except ImportError:

97 pa = None

99try:

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

101 ASTROPY_PANDAS_INDEX_KEY,

102 ArrowAstropySchema,

103 ArrowNumpySchema,

104 DataFrameSchema,

105 ParquetFormatter,

106 _append_numpy_multidim_metadata,

107 _astropy_to_numpy_dict,

108 _numpy_dict_to_numpy,

109 _numpy_dtype_to_arrow_types,

110 _numpy_style_arrays_to_arrow_arrays,

111 _numpy_to_numpy_dict,

112 add_pandas_index_to_astropy,

113 arrow_to_astropy,

114 arrow_to_numpy,

115 arrow_to_numpy_dict,

116 arrow_to_pandas,

117 astropy_to_arrow,

118 astropy_to_pandas,

119 compute_row_group_size,

120 numpy_dict_to_arrow,

121 numpy_to_arrow,

122 pandas_to_arrow,

123 pandas_to_astropy,

124 )

125except ImportError:

126 pa = None

127 pd = None

128 atable = None

129 np = None

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

131

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

133

134

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

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

137

138 Parameters

139 ----------

140 include_multidim : `bool`

141 Include multi-dimensional columns.

142 include_bigendian : `bool`

143 Include big-endian columns.

144

145 Returns

146 -------

147 numpyTable : `numpy.ndarray`

148 """

149 nrow = 5

150

151 dtype = [

152 ("index", "i4"),

153 ("a", "f8"),

154 ("b", "f8"),

155 ("c", "f8"),

156 ("ddd", "f8"),

157 ("f", "i8"),

158 ("strcol", "U10"),

159 ("bytecol", "S10"),

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

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

162 ]

163

164 if include_multidim:

165 dtype.extend(

166 [

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

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

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

170 ]

171 )

172

173 if include_bigendian:

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

175

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

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

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

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

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

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

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

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

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

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

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

187

188 if include_multidim:

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

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

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

192

193 if include_bigendian:

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

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

196

197 return data

198

199

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

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

202

203 Parameters

204 ----------

205 include_masked : `bool`

206 Include masked columns.

207 include_lists : `bool`

208 Include list columns.

209

210 Returns

211 -------

212 dataFrame : `~pandas.DataFrame`

213 The test dataframe.

214 allColumns : `list` [`str`]

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

216 """

217 data = _makeSimpleNumpyTable()

218 df = pd.DataFrame(data)

219 df = df.set_index("index")

220

221 if include_masked:

222 nrow = len(df)

223

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

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

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

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

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

229

230 if include_lists:

231 nrow = len(df)

232

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

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

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

236

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

238

239 return df, allColumns

240

241

242def _makeMultiIndexDataFrame():

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

244

245 Returns

246 -------

247 dataFrame : `~pandas.DataFrame`

248 The test dataframe.

249 """

250 columns = pd.MultiIndex.from_tuples(

251 [

252 ("g", "a"),

253 ("g", "b"),

254 ("g", "c"),

255 ("r", "a"),

256 ("r", "b"),

257 ("r", "c"),

258 ],

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

260 )

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

262

263 return df

264

265

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

267 """Make an astropy table for testing.

268

269 Parameters

270 ----------

271 include_multidim : `bool`

272 Include multi-dimensional columns.

273 include_masked : `bool`

274 Include masked columns.

275 include_bigendian : `bool`

276 Include big-endian columns.

277

278 Returns

279 -------

280 astropyTable : `astropy.table.Table`

281 The test table.

282 """

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

284 # Add a couple of units.

285 table = atable.Table(data)

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

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

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

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

290

291 # Add some masked columns.

292 if include_masked:

293 nrow = len(table)

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

295 mask[1] = True

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

297 # to be able test after serialization.

298

299 # Masked 64-bit integer.

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

301 arr[mask] = -1

302 arr[0] = 1649900760361600113

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

304 # Masked 32-bit float.

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

306 arr[mask] = np.nan

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

308 # Unmasked 32-bit float with NaNs.

309 table["um_f4"] = arr

310 # Masked 64-bit float.

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

312 arr[mask] = np.nan

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

314 # Unmasked 64-bit float with NaNs.

315 table["um_f8"] = arr

316 # Masked boolean.

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

318 arr[mask] = True

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

320 # Masked unsigned 32-bit unsigned int.

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

322 arr[mask] = 0

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

324 # Masked string.

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

326 # Masked bytes.

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

328

329 return table

330

331

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

333 """Make an arrow table for testing.

334

335 Parameters

336 ----------

337 include_multidim : `bool`

338 Include multi-dimensional columns.

339 include_masked : `bool`

340 Include masked columns.

341

342 Returns

343 -------

344 arrowTable : `pyarrow.Table`

345 The test table.

346 """

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

348 return astropy_to_arrow(data)

349

350

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

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

353class ParquetFormatterDataFrameTestCase(unittest.TestCase):

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

355

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

357

358 def setUp(self):

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

360 self.root = makeTestTempDir(TESTDIR)

361 config = Config(self.configFile)

362 self.run = "test_run"

363 self.butler = Butler.from_config(

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

365 )

366 self.enterContext(self.butler)

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

368 # inserting dimension data or defining data IDs.

369 self.datasetType = DatasetType(

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

371 )

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

373

374 def tearDown(self):

375 removeTestTempDir(self.root)

376

377 def testSingleIndexDataFrame(self):

378 df1, allColumns = _makeSingleIndexDataFrame(include_masked=True)

379

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

381 # Read the whole DataFrame.

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

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

384 # Read just the column descriptions.

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

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

387 # Read the rowcount.

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

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

390 # Read the schema.

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

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

393 # Read just some columns a few different ways.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

408 # Passing an unrecognized column should be a ValueError.

409 with self.assertRaises(ValueError):

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

411

412 def testSingleIndexDataFrameWithLists(self):

413 df1, allColumns = _makeSingleIndexDataFrame(include_lists=True)

414

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

416 # Read the whole DataFrame.

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

418

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

420 # from lists to arrays.

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

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

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

424

425 def testMultiIndexDataFrame(self):

426 df1 = _makeMultiIndexDataFrame()

427

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

429 # Read the whole DataFrame.

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

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

432 # Read just the column descriptions.

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

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

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

436 # Read the rowcount.

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

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

439 # Read the schema.

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

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

442 # Read just some columns a few different ways.

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

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

445 df4 = self.butler.get(

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

447 )

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

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

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

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

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

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

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

455 # Passing an unrecognized column should be a ValueError.

456 with self.assertRaises(ValueError):

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

458

459 def testSingleIndexDataFrameEmptyString(self):

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

461 df1, _ = _makeSingleIndexDataFrame()

462

463 # Set one of the strings to None

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

465

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

467 # Read the whole DataFrame.

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

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

470

471 def testSingleIndexDataFrameAllEmptyStrings(self):

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

473 column.

474 """

475 df1, _ = _makeSingleIndexDataFrame()

476

477 # Set all of the strings to None

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

479

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

481 # Read the whole DataFrame.

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

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

484

485 def testLegacyDataFrame(self):

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

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

488 functionality.

489 """

490 df1, allColumns = _makeSingleIndexDataFrame()

491

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

493 df1.to_parquet(fname)

494

495 legacy_type = DatasetType(

496 "legacy_dataframe",

497 dimensions=(),

498 storageClass="DataFrame",

499 universe=self.butler.dimensions,

500 )

501 self.butler.registry.registerDatasetType(legacy_type)

502

503 data_id = {}

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

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

506

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

508

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

510

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

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

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

514

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

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

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

518

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

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

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

522

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

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

525 self.assertEqual(rowcount2a, rowcount2b)

526

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

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

529 self.assertEqual(schema2a, schema2b)

530

531 def testDataFrameSchema(self):

532 tab1 = _makeSimpleArrowTable()

533

534 schema = DataFrameSchema.from_arrow(tab1.schema)

535

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

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

538 self.assertNotEqual(schema, "not_a_schema")

539 self.assertEqual(schema, schema)

540

541 tab2 = _makeMultiIndexDataFrame()

542 schema2 = DataFrameSchema(tab2)

543

544 self.assertNotEqual(schema, schema2)

545

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

547 def testWriteSingleIndexDataFrameReadAsAstropyTable(self):

548 df1, allColumns = _makeSingleIndexDataFrame()

549

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

551

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

553

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

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

556

557 # Check reading the columns.

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

559 columns2 = self.butler.get(

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

561 )

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

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

564

565 # Check reading the schema.

566 schema = ArrowAstropySchema(tab2)

567 schema2 = self.butler.get(

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

569 )

570

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

572 # will be changed because of pandas indexing.

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

574 for name in schema.schema.columns:

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

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

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

578

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

580 def testWriteSingleIndexDataFrameWithMaskedColsReadAsAstropyTable(self):

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

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

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

584 # inconsistent.)

585 df1, allColumns = _makeSingleIndexDataFrame(include_masked=True)

586

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

588

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

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

591

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

593 for name in tab2_df.columns:

594 col1 = df1[name]

595 col2 = tab2_df[name]

596

597 if col1.hasnans:

598 notNull = col1.notnull()

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

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

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

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

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

604 else:

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

606

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

608 def testWriteMultiIndexDataFrameReadAsAstropyTable(self):

609 df1 = _makeMultiIndexDataFrame()

610

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

612

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

614

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

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

617 # recommended.

618

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

620 def testWriteAstropyTableWithMaskedColsReadAsSingleIndexDataFrame(self):

621 tab1 = _makeSimpleAstropyTable(include_masked=True)

622

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

624

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

626

627 tab1_df = astropy_to_pandas(tab1)

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

629

630 tab2_astropy = pandas_to_astropy(tab2)

631 for col in tab1.dtype.names:

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

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

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

635

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

637 def testWriteSingleIndexDataFrameReadAsArrowTable(self):

638 df1, allColumns = _makeSingleIndexDataFrame()

639

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

641

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

643

644 tab2_df = arrow_to_pandas(tab2)

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

646

647 # Check reading the columns.

648 columns = list(tab2.schema.names)

649 columns2 = self.butler.get(

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

651 )

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

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

654

655 # Override the component using a dataset type.

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

657 "ArrowColumnList"

658 )

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

660

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

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

663 # selects the component.

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

665 "columns"

666 )

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

668

669 # Check reading the schema.

670 schema = tab2.schema

671 schema2 = self.butler.get(

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

673 )

674

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

676 # information be maintained.

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

678 for name in schema.names:

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

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

681

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

683 def testWriteMultiIndexDataFrameReadAsArrowTable(self):

684 df1 = _makeMultiIndexDataFrame()

685

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

687

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

689

690 tab2_df = arrow_to_pandas(tab2)

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

692

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

694 def testWriteSingleIndexDataFrameReadAsNumpyTable(self):

695 df1, allColumns = _makeSingleIndexDataFrame()

696

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

698

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

700

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

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

703

704 # Check reading the columns.

705 columns = list(tab2.dtype.names)

706 columns2 = self.butler.get(

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

708 )

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

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

711

712 # Check reading the schema.

713 schema = ArrowNumpySchema(tab2.dtype)

714 schema2 = self.butler.get(

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

716 )

717

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

719 # will be changed because of pandas indexing.

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

721 for name in schema.schema.names:

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

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

724

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

726 def testWriteMultiIndexDataFrameReadAsNumpyTable(self):

727 df1 = _makeMultiIndexDataFrame()

728

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

730

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

732

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

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

735 # recommended.

736

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

738 def testWriteSingleIndexDataFrameReadAsNumpyDict(self):

739 df1, allColumns = _makeSingleIndexDataFrame()

740

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

742

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

744

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

746 # The column order is not maintained.

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

748 for col in df1.columns:

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

750

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

752 def testWriteMultiIndexDataFrameReadAsNumpyDict(self):

753 df1 = _makeMultiIndexDataFrame()

754

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

756

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

758

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

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

761 # recommended.

762

763 def testBadDataFrameColumnParquet(self):

764 df1, allColumns = _makeSingleIndexDataFrame()

765

766 # Make a column with mixed type.

767 bad_col1 = [0.0] * len(df1)

768 bad_col1[1] = 0.0 * units.nJy

769 bad_df = df1.copy()

770 bad_df["bad_col1"] = bad_col1

771

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

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

774 with self.assertRaises(RuntimeError):

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

776

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

778 def testWriteReadAstropyTableLossless(self):

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

780

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

782

783 tab2 = self.butler.get(

784 self.datasetType,

785 dataId={},

786 storageClass="ArrowAstropy",

787 parameters={"strip_astropy_meta_yaml": False},

788 )

789

790 # Check that minimal provenance was written by default.

791 expected = {

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

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

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

795 "LSST.BUTLER.N_INPUTS": 0,

796 }

797

798 self.assertEqual(tab2.meta, expected)

799

800 _checkAstropyTableEquality(tab1, tab2)

801

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

803 def testWriteReadAstropyTableProvenance(self):

804 tab1 = _makeSimpleAstropyTable()

805

806 # Create a ref for provenance.

807 astropy_type = DatasetType(

808 "astropy_parquet",

809 dimensions=(),

810 storageClass="ArrowAstropy",

811 universe=self.butler.dimensions,

812 )

813 self.butler.registry.registerDatasetType(astropy_type)

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

815 quantum_id = uuid.uuid4()

816 provenance = DatasetProvenance(quantum_id=quantum_id)

817 provenance.add_input(input_ref)

818

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

820

821 tab2 = self.butler.get(

822 self.datasetType,

823 dataId={},

824 storageClass="ArrowAstropy",

825 parameters={"strip_astropy_meta_yaml": False},

826 )

827

828 expected = {

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

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

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

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

833 "LSST.BUTLER.N_INPUTS": 1,

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

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

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

837 }

838 self.assertEqual(tab2.meta, expected)

839

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

841 # that the previous provenance was stripped.

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

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

844

845 # tab2 will have been updated in place.

846 expected = {

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

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

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

850 "LSST.BUTLER.N_INPUTS": 0,

851 }

852 self.assertEqual(tab2.meta, expected)

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

854 self.assertEqual(prov_ref, put_ref3)

855 self.assertEqual(null_prov, DatasetProvenance())

856

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

858 def testWriteReadNumpyTableLossless(self):

859 tab1 = _makeSimpleNumpyTable(include_multidim=True)

860

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

862

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

864

865 _checkNumpyTableEquality(tab1, tab2)

866

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

868 def testMaskedNumpy(self):

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

870 tab1_np = arrow_to_numpy(tab1)

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

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

873 col = tab1_np["m_f8"]

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

875

876 # Now without a mask.

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

878 tab1_np = arrow_to_numpy(tab1)

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

880

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

882 def testWriteReadArrowTableLossless(self):

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

884

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

886

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

888

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

890 tab1_np = arrow_to_numpy(tab1)

891 tab2_np = arrow_to_numpy(tab2)

892 for col in tab1.column_names:

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

894

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

896 def testWriteReadNumpyDictLossless(self):

897 tab1 = _makeSimpleNumpyTable(include_multidim=True)

898 dict1 = _numpy_to_numpy_dict(tab1)

899

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

901

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

903

904 _checkNumpyDictEquality(dict1, dict2)

905

906

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

908class InMemoryDataFrameDelegateTestCase(ParquetFormatterDataFrameTestCase):

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

910

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

912

913 def testBadDataFrameColumnParquet(self):

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

915 pass

916

917 def testWriteMultiIndexDataFrameReadAsAstropyTable(self):

918 df1 = _makeMultiIndexDataFrame()

919

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

921

922 with self.assertRaises(ValueError):

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

924

925 def testLegacyDataFrame(self):

926 # This test does not work with an inMemoryDatastore.

927 pass

928

929 def testBadInput(self):

930 df1, _ = _makeSingleIndexDataFrame()

931 delegate = ArrowTableDelegate("DataFrame")

932

933 with self.assertRaises(ValueError):

934 delegate.handleParameters(inMemoryDataset="not_a_dataframe")

935

936 with self.assertRaises(AttributeError):

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

938

939 def testStorageClass(self):

940 df1, allColumns = _makeSingleIndexDataFrame()

941

942 factory = StorageClassFactory()

943 factory.addFromConfig(StorageClassConfig())

944

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

946 # Force the name lookup to do name matching.

947 storageClass._pytype = None

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

949

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

951 # Force the name lookup to do name matching.

952 storageClass._pytype = None

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

954

955

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

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

958class ParquetFormatterArrowAstropyTestCase(unittest.TestCase):

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

960

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

962

963 def setUp(self):

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

965 self.root = makeTestTempDir(TESTDIR)

966 config = Config(self.configFile)

967 self.run = "test_run"

968 self.butler = Butler.from_config(

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

970 )

971 self.enterContext(self.butler)

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

973 # inserting dimension data or defining data IDs.

974 self.datasetType = DatasetType(

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

976 )

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

978

979 def tearDown(self):

980 removeTestTempDir(self.root)

981

982 def testAstropyTable(self):

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

984

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

986 # Read the whole Table.

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

988 _checkAstropyTableEquality(tab1, tab2)

989 # Read the columns.

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

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

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

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

994 # Read the rowcount.

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

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

997 # Read the schema.

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

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

1000 # Read just some columns a few different ways.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

1015 # Passing an unrecognized column should be a ValueError.

1016 with self.assertRaises(ValueError):

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

1018

1019 def testAstropyTableBigEndian(self):

1020 tab1 = _makeSimpleAstropyTable(include_bigendian=True)

1021

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

1023 # Read the whole Table.

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

1025 _checkAstropyTableEquality(tab1, tab2, has_bigendian=True)

1026

1027 def testAstropyTableWithMetadata(self):

1028 tab1 = _makeSimpleAstropyTable(include_multidim=True)

1029

1030 meta = {

1031 "meta_a": 5,

1032 "meta_b": 10.0,

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

1034 "meta_d": True,

1035 "meta_e": "string",

1036 }

1037

1038 tab1.meta.update(meta)

1039

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

1041 # Read the whole Table.

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

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

1044 _checkAstropyTableEquality(tab1, tab2)

1045

1046 def testArrowAstropySchema(self):

1047 tab1 = _makeSimpleAstropyTable()

1048 tab1_arrow = astropy_to_arrow(tab1)

1049 schema = ArrowAstropySchema.from_arrow(tab1_arrow.schema)

1050

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

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

1053 self.assertNotEqual(schema, "not_a_schema")

1054 self.assertEqual(schema, schema)

1055

1056 # Test various inequalities

1057 tab2 = tab1.copy()

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

1059 schema2 = ArrowAstropySchema(tab2)

1060 self.assertNotEqual(schema2, schema)

1061

1062 tab2 = tab1.copy()

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

1064 schema2 = ArrowAstropySchema(tab2)

1065 self.assertNotEqual(schema2, schema)

1066

1067 tab2 = tab1.copy()

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

1069 schema2 = ArrowAstropySchema(tab2)

1070 self.assertNotEqual(schema2, schema)

1071

1072 tab2 = tab1.copy()

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

1074 schema2 = ArrowAstropySchema(tab2)

1075 self.assertNotEqual(schema2, schema)

1076

1077 def testAstropyParquet(self):

1078 tab1 = _makeSimpleAstropyTable()

1079

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

1081 del tab1["dtn"]

1082 del tab1["dtu"]

1083

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

1085 tab1.write(fname)

1086

1087 astropy_type = DatasetType(

1088 "astropy_parquet",

1089 dimensions=(),

1090 storageClass="ArrowAstropy",

1091 universe=self.butler.dimensions,

1092 )

1093 self.butler.registry.registerDatasetType(astropy_type)

1094

1095 data_id = {}

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

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

1098

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

1100

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

1102

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

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

1105 _checkAstropyTableEquality(tab2a, tab2b)

1106

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

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

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

1110 for i, name in enumerate(columns2a):

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

1112

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

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

1115 self.assertEqual(rowcount2a, rowcount2b)

1116

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

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

1119 self.assertEqual(schema2a, schema2b)

1120

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

1122 def testWriteAstropyReadAsArrowTable(self):

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

1124 tab1 = _makeSimpleAstropyTable(include_masked=True)

1125

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

1127

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

1129

1130 tab2_astropy = arrow_to_astropy(tab2)

1131 _checkAstropyTableEquality(tab1, tab2_astropy)

1132

1133 # Check reading the columns.

1134 columns = tab2.schema.names

1135 columns2 = self.butler.get(

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

1137 )

1138 self.assertEqual(columns2, columns)

1139

1140 # Check reading the schema.

1141 schema = tab2.schema

1142 schema2 = self.butler.get(

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

1144 )

1145

1146 self.assertEqual(schema, schema2)

1147

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

1149 def testWriteAstropyReadAsDataFrame(self):

1150 tab1 = _makeSimpleAstropyTable()

1151

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

1153

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

1155

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

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

1158

1159 tab1_df = tab1.to_pandas()

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

1161

1162 # Check reading the columns.

1163 columns = tab2.columns

1164 columns2 = self.butler.get(

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

1166 )

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

1168

1169 # Check reading the schema.

1170 schema = DataFrameSchema(tab2)

1171 schema2 = self.butler.get(

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

1173 )

1174

1175 self.assertEqual(schema2, schema)

1176

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

1178 def testWriteAstropyWithMaskedColsReadAsDataFrame(self):

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

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

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

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

1183 tab1 = _makeSimpleAstropyTable(include_masked=True)

1184

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

1186

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

1188

1189 tab1_df = astropy_to_pandas(tab1)

1190

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

1192 for name in tab2.columns:

1193 col1 = tab1_df[name]

1194 col2 = tab2[name]

1195

1196 if col1.hasnans:

1197 notNull = col1.notnull()

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

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

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

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

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

1203 else:

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

1205

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

1207 def testWriteSingleIndexDataFrameWithMaskedColsReadAsAstropyTable(self):

1208 df1, allColumns = _makeSingleIndexDataFrame(include_masked=True)

1209

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

1211

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

1213

1214 df1_tab = pandas_to_astropy(df1)

1215

1216 _checkAstropyTableEquality(df1_tab, tab2)

1217

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

1219 def testWriteAstropyReadAsNumpyTable(self):

1220 tab1 = _makeSimpleAstropyTable()

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

1222

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

1224

1225 # This is tricky because it loses the units.

1226 tab2_astropy = atable.Table(tab2)

1227

1228 _checkAstropyTableEquality(tab1, tab2_astropy, skip_units=True)

1229

1230 # Check reading the columns.

1231 columns = list(tab2.dtype.names)

1232 columns2 = self.butler.get(

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

1234 )

1235 self.assertEqual(columns2, columns)

1236

1237 # Check reading the schema.

1238 schema = ArrowNumpySchema(tab2.dtype)

1239 schema2 = self.butler.get(

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

1241 )

1242

1243 self.assertEqual(schema2, schema)

1244

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

1246 def testWriteAstropyReadAsNumpyDict(self):

1247 tab1 = _makeSimpleAstropyTable()

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

1249

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

1251

1252 # This is tricky because it loses the units.

1253 tab2_astropy = atable.Table(tab2)

1254

1255 _checkAstropyTableEquality(tab1, tab2_astropy, skip_units=True)

1256

1257 def testBadAstropyColumnParquet(self):

1258 tab1 = _makeSimpleAstropyTable()

1259

1260 # Make a column with mixed type.

1261 bad_col1 = [0.0] * len(tab1)

1262 bad_col1[1] = 0.0 * units.nJy

1263 bad_tab = tab1.copy()

1264 bad_tab["bad_col1"] = bad_col1

1265

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

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

1268 with self.assertRaises(RuntimeError):

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

1270

1271 # Make a column with ragged size.

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

1273 bad_col2[1] = [0, 0]

1274 bad_tab = tab1.copy()

1275 bad_tab["bad_col2"] = bad_col2

1276

1277 with self.assertRaises(RuntimeError):

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

1279

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

1281 def testWriteAstropyTableWithPandasIndexHint(self, testStrip=True):

1282 tab1 = _makeSimpleAstropyTable()

1283

1284 add_pandas_index_to_astropy(tab1, "index")

1285

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

1287

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

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

1290

1291 self.assertIn(ASTROPY_PANDAS_INDEX_KEY, tab2.meta)

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

1293

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

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

1296

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

1298

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

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

1301 _ = self.butler.get(

1302 self.datasetType,

1303 dataId={},

1304 storageClass="DataFrame",

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

1306 )

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

1308

1309 if testStrip:

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

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

1312

1313 self.assertNotIn(ASTROPY_PANDAS_INDEX_KEY, tab5.meta)

1314

1315 with self.assertRaises(ValueError):

1316 add_pandas_index_to_astropy(tab1, "not_a_column")

1317

1318

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

1320class InMemoryArrowAstropyDelegateTestCase(ParquetFormatterArrowAstropyTestCase):

1321 """Tests for InMemoryDatastore, using ArrowTableDelegate with

1322 AstropyTable.

1323 """

1324

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

1326

1327 def testAstropyParquet(self):

1328 # This test does not work with an inMemoryDatastore.

1329 pass

1330

1331 def testBadAstropyColumnParquet(self):

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

1333 pass

1334

1335 def testBadInput(self):

1336 tab1 = _makeSimpleAstropyTable()

1337 delegate = ArrowTableDelegate("ArrowAstropy")

1338

1339 with self.assertRaises(ValueError):

1340 delegate.handleParameters(inMemoryDataset="not_an_astropy_table")

1341

1342 with self.assertRaises(NotImplementedError):

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

1344

1345 with self.assertRaises(AttributeError):

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

1347

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

1349 def testWriteAstropyTableWithPandasIndexHint(self):

1350 super().testWriteAstropyTableWithPandasIndexHint(testStrip=False)

1351

1352

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

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

1355class ParquetFormatterArrowNumpyTestCase(unittest.TestCase):

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

1357

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

1359

1360 def setUp(self):

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

1362 self.root = makeTestTempDir(TESTDIR)

1363 config = Config(self.configFile)

1364 self.butler = Butler.from_config(

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

1366 )

1367 self.enterContext(self.butler)

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

1369 # inserting dimension data or defining data IDs.

1370 self.datasetType = DatasetType(

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

1372 )

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

1374

1375 def tearDown(self):

1376 removeTestTempDir(self.root)

1377

1378 def testNumpyTable(self):

1379 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1380

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

1382 # Read the whole Table.

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

1384 _checkNumpyTableEquality(tab1, tab2)

1385 # Read the columns.

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

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

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

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

1390 # Read the rowcount.

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

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

1393 # Read the schema.

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

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

1396 # Read just some columns a few different ways.

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

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

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

1400 _checkNumpyTableEquality(

1401 tab1[

1402 [

1403 "a",

1404 ]

1405 ],

1406 tab4,

1407 )

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

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

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

1411 _checkNumpyTableEquality(

1412 tab1[

1413 [

1414 "ddd",

1415 ]

1416 ],

1417 tab6,

1418 )

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

1420 _checkNumpyTableEquality(

1421 tab1[

1422 [

1423 "a",

1424 ]

1425 ],

1426 tab7,

1427 )

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

1429 _checkNumpyTableEquality(

1430 tab1[

1431 [

1432 "ddd",

1433 "dtn",

1434 "dtu",

1435 "a",

1436 ]

1437 ],

1438 tab8,

1439 )

1440 # Passing an unrecognized column should be a ValueError.

1441 with self.assertRaises(ValueError):

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

1443

1444 def testNumpyTableBigEndian(self):

1445 tab1 = _makeSimpleNumpyTable(include_bigendian=True)

1446

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

1448 # Read the whole Table.

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

1450 _checkNumpyTableEquality(tab1, tab2, has_bigendian=True)

1451

1452 def testArrowNumpySchema(self):

1453 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1454 tab1_arrow = numpy_to_arrow(tab1)

1455 schema = ArrowNumpySchema.from_arrow(tab1_arrow.schema)

1456

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

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

1459 self.assertNotEqual(schema, "not_a_schema")

1460 self.assertEqual(schema, schema)

1461

1462 # Test inequality

1463 tab2 = tab1.copy()

1464 names = list(tab2.dtype.names)

1465 names[0] = "index2"

1466 tab2.dtype.names = names

1467 schema2 = ArrowNumpySchema(tab2.dtype)

1468 self.assertNotEqual(schema2, schema)

1469

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

1471 def testNumpyDictConversions(self):

1472 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1473

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

1475 tab1_arrow = numpy_to_arrow(tab1)

1476 tab1_dict = arrow_to_numpy_dict(tab1_arrow)

1477 tab1_dict_arrow = numpy_dict_to_arrow(tab1_dict)

1478

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

1480 self.assertEqual(tab1_arrow, tab1_dict_arrow)

1481

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

1483 def testWriteNumpyTableReadAsArrowTable(self):

1484 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1485

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

1487

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

1489

1490 tab2_numpy = arrow_to_numpy(tab2)

1491

1492 _checkNumpyTableEquality(tab1, tab2_numpy)

1493

1494 # Check reading the columns.

1495 columns = tab2.schema.names

1496 columns2 = self.butler.get(

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

1498 )

1499 self.assertEqual(columns2, columns)

1500

1501 # Check reading the schema.

1502 schema = tab2.schema

1503 schema2 = self.butler.get(

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

1505 )

1506 self.assertEqual(schema2, schema)

1507

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

1509 def testWriteNumpyTableReadAsDataFrame(self):

1510 tab1 = _makeSimpleNumpyTable()

1511

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

1513

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

1515

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

1517 # and changes the datatype of the string column.

1518

1519 tab1_df = pd.DataFrame(tab1)

1520

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

1522

1523 # Check reading the columns.

1524 columns = tab2.columns

1525 columns2 = self.butler.get(

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

1527 )

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

1529

1530 # Check reading the schema.

1531 schema = DataFrameSchema(tab2)

1532 schema2 = self.butler.get(

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

1534 )

1535

1536 self.assertEqual(schema2, schema)

1537

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

1539 def testWriteNumpyTableReadAsAstropyTable(self):

1540 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1541

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

1543

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

1545 tab2_numpy = tab2.as_array()

1546

1547 _checkNumpyTableEquality(tab1, tab2_numpy)

1548

1549 # Check reading the columns.

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

1551 columns2 = self.butler.get(

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

1553 )

1554 self.assertEqual(columns2, columns)

1555

1556 # Check reading the schema.

1557 schema = ArrowAstropySchema(tab2)

1558 schema2 = self.butler.get(

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

1560 )

1561

1562 self.assertEqual(schema2, schema)

1563

1564 def testWriteNumpyTableReadAsNumpyDict(self):

1565 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1566

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

1568

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

1570 tab2_numpy = _numpy_dict_to_numpy(tab2)

1571

1572 _checkNumpyTableEquality(tab1, tab2_numpy)

1573

1574 def testBadNumpyColumnParquet(self):

1575 tab1 = _makeSimpleAstropyTable()

1576

1577 # Make a column with mixed type.

1578 bad_col1 = [0.0] * len(tab1)

1579 bad_col1[1] = 0.0 * units.nJy

1580 bad_tab = tab1.copy()

1581 bad_tab["bad_col1"] = bad_col1

1582

1583 bad_tab_np = bad_tab.as_array()

1584

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

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

1587 with self.assertRaises(RuntimeError):

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

1589

1590 # Make a column with ragged size.

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

1592 bad_col2[1] = [0, 0]

1593 bad_tab = tab1.copy()

1594 bad_tab["bad_col2"] = bad_col2

1595

1596 bad_tab_np = bad_tab.as_array()

1597

1598 with self.assertRaises(RuntimeError):

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

1600

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

1602 def testWriteReadAstropyTableLossless(self):

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

1604

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

1606

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

1608

1609 _checkAstropyTableEquality(tab1, tab2)

1610

1611

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

1613class InMemoryArrowNumpyDelegateTestCase(ParquetFormatterArrowNumpyTestCase):

1614 """Tests for InMemoryDatastore, using ArrowTableDelegate with

1615 Numpy table.

1616 """

1617

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

1619

1620 def testBadNumpyColumnParquet(self):

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

1622 pass

1623

1624 def testBadInput(self):

1625 tab1 = _makeSimpleNumpyTable()

1626 delegate = ArrowTableDelegate("ArrowNumpy")

1627

1628 with self.assertRaises(ValueError):

1629 delegate.handleParameters(inMemoryDataset="not_a_numpy_table")

1630

1631 with self.assertRaises(NotImplementedError):

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

1633

1634 with self.assertRaises(AttributeError):

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

1636

1637 def testStorageClass(self):

1638 tab1 = _makeSimpleNumpyTable()

1639

1640 factory = StorageClassFactory()

1641 factory.addFromConfig(StorageClassConfig())

1642

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

1644 # Force the name lookup to do name matching.

1645 storageClass._pytype = None

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

1647

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

1649 # Force the name lookup to do name matching.

1650 storageClass._pytype = None

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

1652

1653

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

1655class ParquetFormatterArrowTableTestCase(unittest.TestCase):

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

1657

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

1659

1660 def setUp(self):

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

1662 self.root = makeTestTempDir(TESTDIR)

1663 config = Config(self.configFile)

1664 self.butler = Butler.from_config(

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

1666 )

1667 self.enterContext(self.butler)

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

1669 # inserting dimension data or defining data IDs.

1670 self.datasetType = DatasetType(

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

1672 )

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

1674

1675 def tearDown(self):

1676 removeTestTempDir(self.root)

1677

1678 def testArrowTable(self):

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

1680

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

1682 # Read the whole Table.

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

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

1685 # comparisons.

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

1687 tab1_np = arrow_to_numpy(tab1)

1688 tab2_np = arrow_to_numpy(tab2)

1689 for col in tab1.column_names:

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

1691 # Read the columns.

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

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

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

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

1696 # Read the rowcount.

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

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

1699 # Read the schema.

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

1701 self.assertEqual(schema, tab1.schema)

1702 # Read just some columns a few different ways.

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

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

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

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

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

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

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

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

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

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

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

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

1715 # Passing an unrecognized column should be a ValueError.

1716 with self.assertRaises(ValueError):

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

1718

1719 def testEmptyArrowTable(self):

1720 data = _makeSimpleNumpyTable()

1721 type_list = _numpy_dtype_to_arrow_types(data.dtype)

1722

1723 schema = pa.schema(type_list)

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

1725

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

1727

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

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

1730 self.assertEqual(tab2, tab1)

1731

1732 tab1_numpy = arrow_to_numpy(tab1)

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

1734 tab1_numpy_arrow = numpy_to_arrow(tab1_numpy)

1735 self.assertEqual(tab1_numpy_arrow, tab1)

1736

1737 tab1_pandas = arrow_to_pandas(tab1)

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

1739 tab1_pandas_arrow = pandas_to_arrow(tab1_pandas)

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

1741 # through empty pandas dataframes.

1742 self.assertEqual(

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

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

1745 )

1746

1747 tab1_astropy = arrow_to_astropy(tab1)

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

1749 tab1_astropy_arrow = astropy_to_arrow(tab1_astropy)

1750 self.assertEqual(tab1_astropy_arrow, tab1)

1751

1752 def testEmptyArrowTableMultidim(self):

1753 data = _makeSimpleNumpyTable(include_multidim=True)

1754 type_list = _numpy_dtype_to_arrow_types(data.dtype)

1755

1756 md = {}

1757 for name in data.dtype.names:

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

1759

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

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

1762

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

1764

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

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

1767 self.assertEqual(tab2, tab1)

1768

1769 tab1_numpy = arrow_to_numpy(tab1)

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

1771 tab1_numpy_arrow = numpy_to_arrow(tab1_numpy)

1772 self.assertEqual(tab1_numpy_arrow, tab1)

1773

1774 tab1_astropy = arrow_to_astropy(tab1)

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

1776 tab1_astropy_arrow = astropy_to_arrow(tab1_astropy)

1777 self.assertEqual(tab1_astropy_arrow, tab1)

1778

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

1780 def testWriteArrowTableReadAsSingleIndexDataFrame(self):

1781 df1, allColumns = _makeSingleIndexDataFrame()

1782

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

1784

1785 # Read back out as a dataframe.

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

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

1788

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

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

1791 df3 = arrow_to_pandas(tab3)

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

1793

1794 # Check reading the columns.

1795 columns = df2.reset_index().columns

1796 columns2 = self.butler.get(

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

1798 )

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

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

1801

1802 # Check reading the schema.

1803 schema = DataFrameSchema(df1)

1804 schema2 = self.butler.get(

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

1806 )

1807 self.assertEqual(schema2, schema)

1808

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

1810 def testWriteArrowTableReadAsMultiIndexDataFrame(self):

1811 df1 = _makeMultiIndexDataFrame()

1812

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

1814

1815 # Read back out as a dataframe.

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

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

1818

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

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

1821 df3 = arrow_to_pandas(atab3)

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

1823

1824 # Check reading the columns.

1825 columns = df2.columns

1826 columns2 = self.butler.get(

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

1828 )

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

1830

1831 # Check reading the schema.

1832 schema = DataFrameSchema(df1)

1833 schema2 = self.butler.get(

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

1835 )

1836 self.assertEqual(schema2, schema)

1837

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

1839 def testWriteArrowTableReadAsAstropyTable(self):

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

1841

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

1843

1844 # Read back out as an astropy table.

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

1846 _checkAstropyTableEquality(tab1, tab2)

1847

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

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

1850 tab3 = arrow_to_astropy(atab3)

1851 _checkAstropyTableEquality(tab1, tab3)

1852

1853 # Check reading the columns.

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

1855 columns2 = self.butler.get(

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

1857 )

1858 self.assertEqual(columns2, columns)

1859

1860 # Check reading the schema.

1861 schema = ArrowAstropySchema(tab1)

1862 schema2 = self.butler.get(

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

1864 )

1865 self.assertEqual(schema2, schema)

1866

1867 # Check the schema conversions and units.

1868 arrow_schema = schema.to_arrow_schema()

1869 for name in arrow_schema.names:

1870 field_metadata = arrow_schema.field(name).metadata

1871 if (

1872 b"description" in field_metadata

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

1874 ):

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

1876 else:

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

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

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

1880

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

1882 def testWriteArrowTableReadAsNumpyTable(self):

1883 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1884

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

1886

1887 # Read back out as a numpy table.

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

1889 _checkNumpyTableEquality(tab1, tab2)

1890

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

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

1893 tab3 = arrow_to_numpy(atab3)

1894 _checkNumpyTableEquality(tab1, tab3)

1895

1896 # Check reading the columns.

1897 columns = list(tab2.dtype.names)

1898 columns2 = self.butler.get(

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

1900 )

1901 self.assertEqual(columns2, columns)

1902

1903 # Check reading the schema.

1904 schema = ArrowNumpySchema(tab1.dtype)

1905 schema2 = self.butler.get(

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

1907 )

1908 self.assertEqual(schema2, schema)

1909

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

1911 def testWriteArrowTableReadAsNumpyDict(self):

1912 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1913

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

1915

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

1917 tab2_numpy = _numpy_dict_to_numpy(tab2)

1918 _checkNumpyTableEquality(tab1, tab2_numpy)

1919

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

1921 def testWriteReadAstropyTableLossless(self):

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

1923

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

1925

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

1927

1928 _checkAstropyTableEquality(tab1, tab2)

1929

1930

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

1932class InMemoryArrowTableDelegateTestCase(ParquetFormatterArrowTableTestCase):

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

1934

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

1936

1937 def testBadInput(self):

1938 tab1 = _makeSimpleArrowTable()

1939 delegate = ArrowTableDelegate("ArrowTable")

1940

1941 with self.assertRaises(ValueError):

1942 delegate.handleParameters(inMemoryDataset="not_an_arrow_table")

1943

1944 with self.assertRaises(NotImplementedError):

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

1946

1947 with self.assertRaises(AttributeError):

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

1949

1950 def testStorageClass(self):

1951 tab1 = _makeSimpleArrowTable()

1952

1953 factory = StorageClassFactory()

1954 factory.addFromConfig(StorageClassConfig())

1955

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

1957 # Force the name lookup to do name matching.

1958 storageClass._pytype = None

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

1960

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

1962 # Force the name lookup to do name matching.

1963 storageClass._pytype = None

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

1965

1966

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

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

1969class ParquetFormatterArrowNumpyDictTestCase(unittest.TestCase):

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

1971

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

1973

1974 def setUp(self):

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

1976 self.root = makeTestTempDir(TESTDIR)

1977 config = Config(self.configFile)

1978 self.butler = Butler.from_config(

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

1980 )

1981 self.enterContext(self.butler)

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

1983 # inserting dimension data or defining data IDs.

1984 self.datasetType = DatasetType(

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

1986 )

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

1988

1989 def tearDown(self):

1990 removeTestTempDir(self.root)

1991

1992 def testNumpyDict(self):

1993 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1994 dict1 = _numpy_to_numpy_dict(tab1)

1995

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

1997 # Read the whole table.

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

1999 _checkNumpyDictEquality(dict1, dict2)

2000 # Read the columns.

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

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

2003 for name in dict1:

2004 self.assertIn(name, columns2)

2005 # Read the rowcount.

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

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

2008 # Read the schema.

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

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

2011 # Read just some columns a few different ways.

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

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

2014 _checkNumpyDictEquality(subdict, tab3)

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

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

2017 _checkNumpyDictEquality(subdict, tab4)

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

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

2020 _checkNumpyDictEquality(subdict, tab5)

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

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

2023 _checkNumpyDictEquality(subdict, tab6)

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

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

2026 _checkNumpyDictEquality(subdict, tab7)

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

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

2029 _checkNumpyDictEquality(subdict, tab8)

2030 # Passing an unrecognized column should be a ValueError.

2031 with self.assertRaises(ValueError):

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

2033

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

2035 def testWriteNumpyDictReadAsArrowTable(self):

2036 tab1 = _makeSimpleNumpyTable(include_multidim=True)

2037 dict1 = _numpy_to_numpy_dict(tab1)

2038

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

2040

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

2042

2043 tab2_dict = arrow_to_numpy_dict(tab2)

2044

2045 _checkNumpyDictEquality(dict1, tab2_dict)

2046

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

2048 def testWriteNumpyDictReadAsDataFrame(self):

2049 tab1 = _makeSimpleNumpyTable()

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="DataFrame")

2055

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

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

2058 # changes the datatype of the string column.

2059 tab1_df = pd.DataFrame(tab1)

2060

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

2062 for col in tab1_df.columns:

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

2064

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

2066 def testWriteNumpyDictReadAsAstropyTable(self):

2067 tab1 = _makeSimpleNumpyTable(include_multidim=True)

2068 dict1 = _numpy_to_numpy_dict(tab1)

2069

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

2071

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

2073 tab2_dict = _astropy_to_numpy_dict(tab2)

2074

2075 _checkNumpyDictEquality(dict1, tab2_dict)

2076

2077 def testWriteNumpyDictReadAsNumpyTable(self):

2078 tab1 = _makeSimpleNumpyTable(include_multidim=True)

2079 dict1 = _numpy_to_numpy_dict(tab1)

2080

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

2082

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

2084 tab2_dict = _numpy_to_numpy_dict(tab2)

2085

2086 _checkNumpyDictEquality(dict1, tab2_dict)

2087

2088 def testWriteNumpyDictBad(self):

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

2090 with self.assertRaises(RuntimeError):

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

2092

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

2094 with self.assertRaises(RuntimeError):

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

2096

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

2098 with self.assertRaises(RuntimeError):

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

2100

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

2102 with self.assertRaises(RuntimeError):

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

2104

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

2106 def testWriteReadAstropyTableLossless(self):

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

2108

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

2110

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

2112

2113 _checkAstropyTableEquality(tab1, tab2)

2114

2115

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

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

2118class InMemoryNumpyDictDelegateTestCase(ParquetFormatterArrowNumpyDictTestCase):

2119 """Tests for InMemoryDatastore, using ArrowTableDelegate with

2120 Numpy dict.

2121 """

2122

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

2124

2125 def testWriteNumpyDictBad(self):

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

2127 pass

2128

2129

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

2131class ParquetFormatterArrowSchemaTestCase(unittest.TestCase):

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

2133

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

2135

2136 def setUp(self):

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

2138 self.root = makeTestTempDir(TESTDIR)

2139 config = Config(self.configFile)

2140 self.butler = Butler.from_config(

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

2142 )

2143 self.enterContext(self.butler)

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

2145 # inserting dimension data or defining data IDs.

2146 self.datasetType = DatasetType(

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

2148 )

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

2150

2151 def tearDown(self):

2152 removeTestTempDir(self.root)

2153

2154 def _makeTestSchema(self):

2155 schema = pa.schema(

2156 [

2157 pa.field(

2158 "int32",

2159 pa.int32(),

2160 nullable=False,

2161 metadata={

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

2163 "unit": "",

2164 },

2165 ),

2166 pa.field(

2167 "int64",

2168 pa.int64(),

2169 nullable=False,

2170 metadata={

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

2172 "unit": "",

2173 },

2174 ),

2175 pa.field(

2176 "uint64",

2177 pa.uint64(),

2178 nullable=False,

2179 metadata={

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

2181 "unit": "",

2182 },

2183 ),

2184 pa.field(

2185 "float32",

2186 pa.float32(),

2187 nullable=False,

2188 metadata={

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

2190 "unit": "count",

2191 },

2192 ),

2193 pa.field(

2194 "float64",

2195 pa.float64(),

2196 nullable=False,

2197 metadata={

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

2199 "unit": "nJy",

2200 },

2201 ),

2202 pa.field(

2203 "fixed_size_list",

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

2205 nullable=False,

2206 metadata={

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

2208 "unit": "nJy",

2209 },

2210 ),

2211 pa.field(

2212 "variable_size_list",

2213 pa.list_(pa.float64()),

2214 nullable=False,

2215 metadata={

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

2217 "unit": "nJy",

2218 },

2219 ),

2220 # One of these fields will have no description.

2221 pa.field(

2222 "string",

2223 pa.string(),

2224 nullable=False,

2225 metadata={

2226 "unit": "",

2227 },

2228 ),

2229 # One of these fields will have no metadata.

2230 pa.field(

2231 "binary",

2232 pa.binary(),

2233 nullable=False,

2234 ),

2235 ]

2236 )

2237

2238 return schema

2239

2240 def testArrowSchema(self):

2241 schema1 = self._makeTestSchema()

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

2243

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

2245 self.assertEqual(schema2, schema1)

2246

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

2248 def testWriteArrowSchemaReadAsDataFrameSchema(self):

2249 schema1 = self._makeTestSchema()

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

2251

2252 df_schema1 = DataFrameSchema.from_arrow(schema1)

2253

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

2255 self.assertEqual(df_schema2, df_schema1)

2256

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

2258 def testWriteArrowSchemaReadAsArrowAstropySchema(self):

2259 schema1 = self._makeTestSchema()

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

2261

2262 ap_schema1 = ArrowAstropySchema.from_arrow(schema1)

2263

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

2265 self.assertEqual(ap_schema2, ap_schema1)

2266

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

2268 for name in schema1.names:

2269 field_metadata = schema1.field(name).metadata

2270 if field_metadata is None:

2271 continue

2272 if (

2273 b"description" in field_metadata

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

2275 ):

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

2277 else:

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

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

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

2281

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

2283 def testWriteArrowSchemaReadAsArrowNumpySchema(self):

2284 schema1 = self._makeTestSchema()

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

2286

2287 np_schema1 = ArrowNumpySchema.from_arrow(schema1)

2288

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

2290 self.assertEqual(np_schema2, np_schema1)

2291

2292

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

2294class InMemoryArrowSchemaDelegateTestCase(ParquetFormatterArrowSchemaTestCase):

2295 """Tests for InMemoryDatastore and ArrowSchema."""

2296

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

2298

2299

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

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

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

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

2304class ParquetFormatterArrowTableS3TestCase(unittest.TestCase):

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

2306

2307 # Code is adapted from test_butler.py

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

2309 fullConfigKey = None

2310 validationCanFail = True

2311

2312 bucketName = "anybucketname"

2313

2314 root = "butlerRoot/"

2315

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

2317

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

2319

2320 registryStr = "/gen3.sqlite3"

2321

2322 mock_aws = mock_aws()

2323

2324 def setUp(self):

2325 self.root = makeTestTempDir(TESTDIR)

2326

2327 config = Config(self.configFile)

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

2329 self.bucketName = uri.netloc

2330

2331 # Enable S3 mocking of tests.

2332 self.enterContext(clean_test_environment_for_s3())

2333 self.mock_aws.start()

2334

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

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

2337

2338 # need local folder to store registry database

2339 self.reg_dir = makeTestTempDir(TESTDIR)

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

2341

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

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

2344 s3 = boto3.resource("s3")

2345 s3.create_bucket(Bucket=self.bucketName)

2346

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

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

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

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

2351

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

2353 self.enterContext(self.butler)

2354

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

2356 # inserting dimension data or defining data IDs.

2357 self.datasetType = DatasetType(

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

2359 )

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

2361

2362 def tearDown(self):

2363 s3 = boto3.resource("s3")

2364 bucket = s3.Bucket(self.bucketName)

2365 try:

2366 bucket.objects.all().delete()

2367 except botocore.exceptions.ClientError as e:

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

2369 # the key was not reachable - pass

2370 pass

2371 else:

2372 raise

2373

2374 bucket = s3.Bucket(self.bucketName)

2375 bucket.delete()

2376

2377 # Stop the S3 mock.

2378 self.mock_aws.stop()

2379

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

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

2382

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

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

2385

2386 def testArrowTableS3(self):

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

2388

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

2390

2391 # Read the whole Table.

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

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

2394 # comparisons.

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

2396 tab1_np = arrow_to_numpy(tab1)

2397 tab2_np = arrow_to_numpy(tab2)

2398 for col in tab1.column_names:

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

2400 # Read the columns.

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

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

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

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

2405 # Read the rowcount.

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

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

2408 # Read the schema.

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

2410 self.assertEqual(schema, tab1.schema)

2411 # Read just some columns a few different ways.

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

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

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

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

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

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

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

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

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

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

2422 # Passing an unrecognized column should be a ValueError.

2423 with self.assertRaises(ValueError):

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

2425

2426

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

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

2429class ComputeRowGroupSizeTestCase(unittest.TestCase):

2430 """Tests for compute_row_group_size."""

2431

2432 def testRowGroupSizeNoMetadata(self):

2433 numpyTable = _makeSimpleNumpyTable(include_multidim=True)

2434

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

2436 # that adds metadata.

2437 type_list = _numpy_dtype_to_arrow_types(numpyTable.dtype)

2438 schema = pa.schema(type_list)

2439 arrays = _numpy_style_arrays_to_arrow_arrays(

2440 numpyTable.dtype,

2441 len(numpyTable),

2442 numpyTable,

2443 schema,

2444 )

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

2446

2447 row_group_size = compute_row_group_size(arrowTable.schema)

2448

2449 self.assertGreater(row_group_size, 1_000_000)

2450 self.assertLess(row_group_size, 2_000_000)

2451

2452 def testRowGroupSizeWithMetadata(self):

2453 numpyTable = _makeSimpleNumpyTable(include_multidim=True)

2454

2455 arrowTable = numpy_to_arrow(numpyTable)

2456

2457 row_group_size = compute_row_group_size(arrowTable.schema)

2458

2459 self.assertGreater(row_group_size, 1_000_000)

2460 self.assertLess(row_group_size, 2_000_000)

2461

2462 def testRowGroupSizeTinyTable(self):

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

2464

2465 arrowTable = numpy_to_arrow(numpyTable)

2466

2467 row_group_size = compute_row_group_size(arrowTable.schema)

2468

2469 self.assertGreater(row_group_size, 1_000_000)

2470

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

2472 def testRowGroupSizeDataFrameWithLists(self):

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

2474 arrowTable = pandas_to_arrow(df)

2475 row_group_size = compute_row_group_size(arrowTable.schema)

2476

2477 self.assertGreater(row_group_size, 1_000_000)

2478

2479

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

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

2482

2483 Parameters

2484 ----------

2485 table1 : `astropy.table.Table`

2486 table2 : `astropy.table.Table`

2487 skip_units : `bool`

2488 has_bigendian : `bool`

2489 """

2490 if not has_bigendian:

2491 assert table1.dtype == table2.dtype

2492 else:

2493 for name in table1.dtype.names:

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

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

2496

2497 # Strip provenance before comparison.

2498 DatasetProvenance.strip_provenance_from_flat_dict(table1.meta)

2499 DatasetProvenance.strip_provenance_from_flat_dict(table2.meta)

2500 assert table1.meta == table2.meta

2501 if not skip_units:

2502 for name in table1.columns:

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

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

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

2506

2507 for name in table1.columns:

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

2509 has_masked = False

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

2511 c1 = table1[name].filled()

2512 has_masked = True

2513 else:

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

2515 if has_masked:

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

2517 c2 = table2[name].filled()

2518 else:

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

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

2521 np.testing.assert_array_equal(c1, c2)

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

2523 if has_masked:

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

2525 np.testing.assert_array_equal(table1[name].mask, table2[name].mask)

2526

2527

2528def _checkNumpyTableEquality(table1, table2, has_bigendian=False):

2529 """Check if two numpy tables have the same columns/values

2530

2531 Parameters

2532 ----------

2533 table1 : `numpy.ndarray`

2534 table2 : `numpy.ndarray`

2535 has_bigendian : `bool`

2536 """

2537 assert table1.dtype.names == table2.dtype.names

2538 for name in table1.dtype.names:

2539 if not has_bigendian:

2540 assert table1.dtype[name] == table2.dtype[name]

2541 else:

2542 # Only check type matches, force to little-endian.

2543 assert table1.dtype[name].newbyteorder(">") == table2.dtype[name].newbyteorder(">")

2544 assert np.all(table1 == table2)

2545

2546

2547def _checkNumpyDictEquality(dict1, dict2):

2548 """Check if two numpy dicts have the same columns/values.

2549

2550 Parameters

2551 ----------

2552 dict1 : `dict` [`str`, `np.ndarray`]

2553 dict2 : `dict` [`str`, `np.ndarray`]

2554 """

2555 assert set(dict1.keys()) == set(dict2.keys())

2556 for name in dict1:

2557 assert dict1[name].dtype == dict2[name].dtype

2558 assert np.all(dict1[name] == dict2[name])

2559

2560

2561if __name__ == "__main__":

2562 unittest.main()