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 if isinstance(df1.index, pd.RangeIndex): 492 ↛ 498line 492 didn't jump to line 498 because the condition on line 492 was never true

493 # Turn the RangeIndex into a regular index or it won't

494 # give us all the column names. This is necessary for pandas v3.

495 # Unfortunately, parquet files serialized directly with

496 # pandas v3 will not report their index column names if

497 # they are sequential integers.

498 df1.index = pd.Index(df1.index.to_numpy(), name=df1.index.name)

499

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

501 df1.to_parquet(fname)

502

503 legacy_type = DatasetType(

504 "legacy_dataframe",

505 dimensions=(),

506 storageClass="DataFrame",

507 universe=self.butler.dimensions,

508 )

509 self.butler.registry.registerDatasetType(legacy_type)

510

511 data_id = {}

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

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

514

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

516

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

518

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

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

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

522

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

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

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

526

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

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

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

530

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

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

533 self.assertEqual(rowcount2a, rowcount2b)

534

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

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

537 self.assertEqual(schema2a, schema2b)

538

539 def testDataFrameSchema(self):

540 tab1 = _makeSimpleArrowTable()

541

542 schema = DataFrameSchema.from_arrow(tab1.schema)

543

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

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

546 self.assertNotEqual(schema, "not_a_schema")

547 self.assertEqual(schema, schema)

548

549 tab2 = _makeMultiIndexDataFrame()

550 schema2 = DataFrameSchema(tab2)

551

552 self.assertNotEqual(schema, schema2)

553

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

555 def testWriteSingleIndexDataFrameReadAsAstropyTable(self):

556 df1, allColumns = _makeSingleIndexDataFrame()

557

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

559

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

561

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

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

564

565 # Check reading the columns.

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

567 columns2 = self.butler.get(

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

569 )

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

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

572

573 # Check reading the schema.

574 schema = ArrowAstropySchema(tab2)

575 schema2 = self.butler.get(

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

577 )

578

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

580 # will be changed because of pandas indexing.

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

582 for name in schema.schema.columns:

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

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

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

586

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

588 def testWriteSingleIndexDataFrameWithMaskedColsReadAsAstropyTable(self):

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

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

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

592 # inconsistent.)

593 df1, allColumns = _makeSingleIndexDataFrame(include_masked=True)

594

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

596

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

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

599

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

601 for name in tab2_df.columns:

602 col1 = df1[name]

603 col2 = tab2_df[name]

604

605 if col1.hasnans:

606 notNull = col1.notnull()

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

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

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

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

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

612 else:

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

614

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

616 def testWriteMultiIndexDataFrameReadAsAstropyTable(self):

617 df1 = _makeMultiIndexDataFrame()

618

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

620

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

622

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

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

625 # recommended.

626

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

628 def testWriteAstropyTableWithMaskedColsReadAsSingleIndexDataFrame(self):

629 tab1 = _makeSimpleAstropyTable(include_masked=True)

630

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

632

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

634

635 tab1_df = astropy_to_pandas(tab1)

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

637

638 tab2_astropy = pandas_to_astropy(tab2)

639 for col in tab1.dtype.names:

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

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

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

643

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

645 def testWriteSingleIndexDataFrameReadAsArrowTable(self):

646 df1, allColumns = _makeSingleIndexDataFrame()

647

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

649

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

651

652 tab2_df = arrow_to_pandas(tab2)

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

654

655 # Check reading the columns.

656 columns = list(tab2.schema.names)

657 columns2 = self.butler.get(

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

659 )

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

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

662

663 # Override the component using a dataset type.

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

665 "ArrowColumnList"

666 )

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

668

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

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

671 # selects the component.

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

673 "columns"

674 )

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

676

677 # Check reading the schema.

678 schema = tab2.schema

679 schema2 = self.butler.get(

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

681 )

682

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

684 # information be maintained.

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

686 for name in schema.names:

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

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

689

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

691 def testWriteMultiIndexDataFrameReadAsArrowTable(self):

692 df1 = _makeMultiIndexDataFrame()

693

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

695

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

697

698 tab2_df = arrow_to_pandas(tab2)

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

700

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

702 def testWriteSingleIndexDataFrameReadAsNumpyTable(self):

703 df1, allColumns = _makeSingleIndexDataFrame()

704

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

706

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

708

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

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

711

712 # Check reading the columns.

713 columns = list(tab2.dtype.names)

714 columns2 = self.butler.get(

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

716 )

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

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

719

720 # Check reading the schema.

721 schema = ArrowNumpySchema(tab2.dtype)

722 schema2 = self.butler.get(

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

724 )

725

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

727 # will be changed because of pandas indexing.

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

729 for name in schema.schema.names:

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

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

732 # the schema consistently.

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

734 continue

735 else:

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

737

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

739 def testWriteMultiIndexDataFrameReadAsNumpyTable(self):

740 df1 = _makeMultiIndexDataFrame()

741

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

743

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

745

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

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

748 # recommended.

749

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

751 def testWriteSingleIndexDataFrameReadAsNumpyDict(self):

752 df1, allColumns = _makeSingleIndexDataFrame()

753

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

755

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

757

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

759 # The column order is not maintained.

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

761 for col in df1.columns:

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

763

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

765 def testWriteMultiIndexDataFrameReadAsNumpyDict(self):

766 df1 = _makeMultiIndexDataFrame()

767

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

769

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

771

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

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

774 # recommended.

775

776 def testBadDataFrameColumnParquet(self):

777 df1, allColumns = _makeSingleIndexDataFrame()

778

779 # Make a column with mixed type.

780 bad_col1 = [0.0] * len(df1)

781 bad_col1[1] = 0.0 * units.nJy

782 bad_df = df1.copy()

783 bad_df["bad_col1"] = bad_col1

784

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

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

787 with self.assertRaises(RuntimeError):

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

789

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

791 def testWriteReadAstropyTableLossless(self):

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

793

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

795

796 tab2 = self.butler.get(

797 self.datasetType,

798 dataId={},

799 storageClass="ArrowAstropy",

800 parameters={"strip_astropy_meta_yaml": False},

801 )

802

803 # Check that minimal provenance was written by default.

804 expected = {

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

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

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

808 "LSST.BUTLER.N_INPUTS": 0,

809 }

810

811 self.assertEqual(tab2.meta, expected)

812

813 _checkAstropyTableEquality(tab1, tab2)

814

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

816 def testWriteReadAstropyTableProvenance(self):

817 tab1 = _makeSimpleAstropyTable()

818

819 # Create a ref for provenance.

820 astropy_type = DatasetType(

821 "astropy_parquet",

822 dimensions=(),

823 storageClass="ArrowAstropy",

824 universe=self.butler.dimensions,

825 )

826 self.butler.registry.registerDatasetType(astropy_type)

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

828 quantum_id = uuid.uuid4()

829 provenance = DatasetProvenance(quantum_id=quantum_id)

830 provenance.add_input(input_ref)

831

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

833

834 tab2 = self.butler.get(

835 self.datasetType,

836 dataId={},

837 storageClass="ArrowAstropy",

838 parameters={"strip_astropy_meta_yaml": False},

839 )

840

841 expected = {

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

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

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

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

846 "LSST.BUTLER.N_INPUTS": 1,

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

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

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

850 }

851 self.assertEqual(tab2.meta, expected)

852

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

854 # that the previous provenance was stripped.

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

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

857

858 # tab2 will have been updated in place.

859 expected = {

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

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

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

863 "LSST.BUTLER.N_INPUTS": 0,

864 }

865 self.assertEqual(tab2.meta, expected)

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

867 self.assertEqual(prov_ref, put_ref3)

868 self.assertEqual(null_prov, DatasetProvenance())

869

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

871 def testWriteReadNumpyTableLossless(self):

872 tab1 = _makeSimpleNumpyTable(include_multidim=True)

873

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

875

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

877

878 _checkNumpyTableEquality(tab1, tab2)

879

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

881 def testMaskedNumpy(self):

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

883 tab1_np = arrow_to_numpy(tab1)

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

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

886 col = tab1_np["m_f8"]

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

888

889 # Now without a mask.

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

891 tab1_np = arrow_to_numpy(tab1)

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

893

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

895 def testWriteReadArrowTableLossless(self):

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

897

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

899

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

901

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

903 tab1_np = arrow_to_numpy(tab1)

904 tab2_np = arrow_to_numpy(tab2)

905 for col in tab1.column_names:

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

907

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

909 def testWriteReadNumpyDictLossless(self):

910 tab1 = _makeSimpleNumpyTable(include_multidim=True)

911 dict1 = _numpy_to_numpy_dict(tab1)

912

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

914

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

916

917 _checkNumpyDictEquality(dict1, dict2)

918

919

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

921class InMemoryDataFrameDelegateTestCase(ParquetFormatterDataFrameTestCase):

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

923

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

925

926 def testBadDataFrameColumnParquet(self):

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

928 pass

929

930 def testWriteMultiIndexDataFrameReadAsAstropyTable(self):

931 df1 = _makeMultiIndexDataFrame()

932

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

934

935 with self.assertRaises(ValueError):

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

937

938 def testLegacyDataFrame(self):

939 # This test does not work with an inMemoryDatastore.

940 pass

941

942 def testBadInput(self):

943 df1, _ = _makeSingleIndexDataFrame()

944 delegate = ArrowTableDelegate("DataFrame")

945

946 with self.assertRaises(ValueError):

947 delegate.handleParameters(inMemoryDataset="not_a_dataframe")

948

949 with self.assertRaises(AttributeError):

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

951

952 def testStorageClass(self):

953 df1, allColumns = _makeSingleIndexDataFrame()

954

955 factory = StorageClassFactory()

956 factory.addFromConfig(StorageClassConfig())

957

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

959 # Force the name lookup to do name matching.

960 storageClass._pytype = None

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

962

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

964 # Force the name lookup to do name matching.

965 storageClass._pytype = None

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

967

968

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

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

971class ParquetFormatterArrowAstropyTestCase(unittest.TestCase):

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

973

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

975

976 def setUp(self):

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

978 self.root = makeTestTempDir(TESTDIR)

979 config = Config(self.configFile)

980 self.run = "test_run"

981 self.butler = Butler.from_config(

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

983 )

984 self.enterContext(self.butler)

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

986 # inserting dimension data or defining data IDs.

987 self.datasetType = DatasetType(

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

989 )

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

991

992 def tearDown(self):

993 removeTestTempDir(self.root)

994

995 def testAstropyTable(self):

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

997

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

999 # Read the whole Table.

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

1001 _checkAstropyTableEquality(tab1, tab2)

1002 # Read the columns.

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

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

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

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

1007 # Read the rowcount.

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

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

1010 # Read the schema.

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

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

1013 # Read just some columns a few different ways.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

1028 # Passing an unrecognized column should be a ValueError.

1029 with self.assertRaises(ValueError):

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

1031

1032 def testAstropyTableBigEndian(self):

1033 tab1 = _makeSimpleAstropyTable(include_bigendian=True)

1034

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

1036 # Read the whole Table.

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

1038 _checkAstropyTableEquality(tab1, tab2, has_bigendian=True)

1039

1040 def testAstropyTableWithMetadata(self):

1041 tab1 = _makeSimpleAstropyTable(include_multidim=True)

1042

1043 meta = {

1044 "meta_a": 5,

1045 "meta_b": 10.0,

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

1047 "meta_d": True,

1048 "meta_e": "string",

1049 }

1050

1051 tab1.meta.update(meta)

1052

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

1054 # Read the whole Table.

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

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

1057 _checkAstropyTableEquality(tab1, tab2)

1058

1059 def testArrowAstropySchema(self):

1060 tab1 = _makeSimpleAstropyTable()

1061 tab1_arrow = astropy_to_arrow(tab1)

1062 schema = ArrowAstropySchema.from_arrow(tab1_arrow.schema)

1063

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

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

1066 self.assertNotEqual(schema, "not_a_schema")

1067 self.assertEqual(schema, schema)

1068

1069 # Test various inequalities

1070 tab2 = tab1.copy()

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

1072 schema2 = ArrowAstropySchema(tab2)

1073 self.assertNotEqual(schema2, schema)

1074

1075 tab2 = tab1.copy()

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

1077 schema2 = ArrowAstropySchema(tab2)

1078 self.assertNotEqual(schema2, schema)

1079

1080 tab2 = tab1.copy()

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

1082 schema2 = ArrowAstropySchema(tab2)

1083 self.assertNotEqual(schema2, schema)

1084

1085 tab2 = tab1.copy()

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

1087 schema2 = ArrowAstropySchema(tab2)

1088 self.assertNotEqual(schema2, schema)

1089

1090 def testAstropyParquet(self):

1091 tab1 = _makeSimpleAstropyTable()

1092

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

1094 del tab1["dtn"]

1095 del tab1["dtu"]

1096

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

1098 tab1.write(fname)

1099

1100 astropy_type = DatasetType(

1101 "astropy_parquet",

1102 dimensions=(),

1103 storageClass="ArrowAstropy",

1104 universe=self.butler.dimensions,

1105 )

1106 self.butler.registry.registerDatasetType(astropy_type)

1107

1108 data_id = {}

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

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

1111

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

1113

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

1115

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

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

1118 _checkAstropyTableEquality(tab2a, tab2b)

1119

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

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

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

1123 for i, name in enumerate(columns2a):

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

1125

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

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

1128 self.assertEqual(rowcount2a, rowcount2b)

1129

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

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

1132 self.assertEqual(schema2a, schema2b)

1133

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

1135 def testWriteAstropyReadAsArrowTable(self):

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

1137 tab1 = _makeSimpleAstropyTable(include_masked=True)

1138

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

1140

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

1142

1143 tab2_astropy = arrow_to_astropy(tab2)

1144 _checkAstropyTableEquality(tab1, tab2_astropy)

1145

1146 # Check reading the columns.

1147 columns = tab2.schema.names

1148 columns2 = self.butler.get(

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

1150 )

1151 self.assertEqual(columns2, columns)

1152

1153 # Check reading the schema.

1154 schema = tab2.schema

1155 schema2 = self.butler.get(

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

1157 )

1158

1159 self.assertEqual(schema, schema2)

1160

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

1162 def testWriteAstropyReadAsDataFrame(self):

1163 tab1 = _makeSimpleAstropyTable()

1164

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

1166

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

1168

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

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

1171

1172 tab1_df = tab1.to_pandas()

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

1174

1175 # Check reading the columns.

1176 columns = tab2.columns

1177 columns2 = self.butler.get(

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

1179 )

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

1181

1182 # Check reading the schema.

1183 schema = DataFrameSchema(tab2)

1184 schema2 = self.butler.get(

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

1186 )

1187

1188 self.assertEqual(schema2, schema)

1189

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

1191 def testWriteAstropyWithMaskedColsReadAsDataFrame(self):

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

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

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

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

1196 tab1 = _makeSimpleAstropyTable(include_masked=True)

1197

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

1199

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

1201

1202 tab1_df = astropy_to_pandas(tab1)

1203

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

1205 for name in tab2.columns:

1206 col1 = tab1_df[name]

1207 col2 = tab2[name]

1208

1209 if col1.hasnans:

1210 notNull = col1.notnull()

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

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

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

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

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

1216 else:

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

1218

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

1220 def testWriteSingleIndexDataFrameWithMaskedColsReadAsAstropyTable(self):

1221 df1, allColumns = _makeSingleIndexDataFrame(include_masked=True)

1222

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

1224

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

1226

1227 df1_tab = pandas_to_astropy(df1)

1228

1229 _checkAstropyTableEquality(df1_tab, tab2)

1230

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

1232 def testWriteAstropyReadAsNumpyTable(self):

1233 tab1 = _makeSimpleAstropyTable()

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

1235

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

1237

1238 # This is tricky because it loses the units.

1239 tab2_astropy = atable.Table(tab2)

1240

1241 _checkAstropyTableEquality(tab1, tab2_astropy, skip_units=True)

1242

1243 # Check reading the columns.

1244 columns = list(tab2.dtype.names)

1245 columns2 = self.butler.get(

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

1247 )

1248 self.assertEqual(columns2, columns)

1249

1250 # Check reading the schema.

1251 schema = ArrowNumpySchema(tab2.dtype)

1252 schema2 = self.butler.get(

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

1254 )

1255

1256 self.assertEqual(schema2, schema)

1257

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

1259 def testWriteAstropyReadAsNumpyDict(self):

1260 tab1 = _makeSimpleAstropyTable()

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

1262

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

1264

1265 # This is tricky because it loses the units.

1266 tab2_astropy = atable.Table(tab2)

1267

1268 _checkAstropyTableEquality(tab1, tab2_astropy, skip_units=True)

1269

1270 def testBadAstropyColumnParquet(self):

1271 tab1 = _makeSimpleAstropyTable()

1272

1273 # Make a column with mixed type.

1274 bad_col1 = [0.0] * len(tab1)

1275 bad_col1[1] = 0.0 * units.nJy

1276 bad_tab = tab1.copy()

1277 bad_tab["bad_col1"] = bad_col1

1278

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

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

1281 with self.assertRaises(RuntimeError):

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

1283

1284 # Make a column with ragged size.

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

1286 bad_col2[1] = [0, 0]

1287 bad_tab = tab1.copy()

1288 bad_tab["bad_col2"] = bad_col2

1289

1290 with self.assertRaises(RuntimeError):

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

1292

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

1294 def testWriteAstropyTableWithPandasIndexHint(self, testStrip=True):

1295 tab1 = _makeSimpleAstropyTable()

1296

1297 add_pandas_index_to_astropy(tab1, "index")

1298

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

1300

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

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

1303

1304 self.assertIn(ASTROPY_PANDAS_INDEX_KEY, tab2.meta)

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

1306

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

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

1309

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

1311

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

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

1314 _ = self.butler.get(

1315 self.datasetType,

1316 dataId={},

1317 storageClass="DataFrame",

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

1319 )

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

1321

1322 if testStrip:

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

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

1325

1326 self.assertNotIn(ASTROPY_PANDAS_INDEX_KEY, tab5.meta)

1327

1328 with self.assertRaises(ValueError):

1329 add_pandas_index_to_astropy(tab1, "not_a_column")

1330

1331

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

1333class InMemoryArrowAstropyDelegateTestCase(ParquetFormatterArrowAstropyTestCase):

1334 """Tests for InMemoryDatastore, using ArrowTableDelegate with

1335 AstropyTable.

1336 """

1337

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

1339

1340 def testAstropyParquet(self):

1341 # This test does not work with an inMemoryDatastore.

1342 pass

1343

1344 def testBadAstropyColumnParquet(self):

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

1346 pass

1347

1348 def testBadInput(self):

1349 tab1 = _makeSimpleAstropyTable()

1350 delegate = ArrowTableDelegate("ArrowAstropy")

1351

1352 with self.assertRaises(ValueError):

1353 delegate.handleParameters(inMemoryDataset="not_an_astropy_table")

1354

1355 with self.assertRaises(NotImplementedError):

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

1357

1358 with self.assertRaises(AttributeError):

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

1360

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

1362 def testWriteAstropyTableWithPandasIndexHint(self):

1363 super().testWriteAstropyTableWithPandasIndexHint(testStrip=False)

1364

1365

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

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

1368class ParquetFormatterArrowNumpyTestCase(unittest.TestCase):

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

1370

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

1372

1373 def setUp(self):

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

1375 self.root = makeTestTempDir(TESTDIR)

1376 config = Config(self.configFile)

1377 self.butler = Butler.from_config(

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

1379 )

1380 self.enterContext(self.butler)

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

1382 # inserting dimension data or defining data IDs.

1383 self.datasetType = DatasetType(

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

1385 )

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

1387

1388 def tearDown(self):

1389 removeTestTempDir(self.root)

1390

1391 def testNumpyTable(self):

1392 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1393

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

1395 # Read the whole Table.

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

1397 _checkNumpyTableEquality(tab1, tab2)

1398 # Read the columns.

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

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

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

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

1403 # Read the rowcount.

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

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

1406 # Read the schema.

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

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

1409 # Read just some columns a few different ways.

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

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

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

1413 _checkNumpyTableEquality(

1414 tab1[

1415 [

1416 "a",

1417 ]

1418 ],

1419 tab4,

1420 )

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

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

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

1424 _checkNumpyTableEquality(

1425 tab1[

1426 [

1427 "ddd",

1428 ]

1429 ],

1430 tab6,

1431 )

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

1433 _checkNumpyTableEquality(

1434 tab1[

1435 [

1436 "a",

1437 ]

1438 ],

1439 tab7,

1440 )

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

1442 _checkNumpyTableEquality(

1443 tab1[

1444 [

1445 "ddd",

1446 "dtn",

1447 "dtu",

1448 "a",

1449 ]

1450 ],

1451 tab8,

1452 )

1453 # Passing an unrecognized column should be a ValueError.

1454 with self.assertRaises(ValueError):

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

1456

1457 def testNumpyTableBigEndian(self):

1458 tab1 = _makeSimpleNumpyTable(include_bigendian=True)

1459

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

1461 # Read the whole Table.

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

1463 _checkNumpyTableEquality(tab1, tab2, has_bigendian=True)

1464

1465 def testArrowNumpySchema(self):

1466 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1467 tab1_arrow = numpy_to_arrow(tab1)

1468 schema = ArrowNumpySchema.from_arrow(tab1_arrow.schema)

1469

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

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

1472 self.assertNotEqual(schema, "not_a_schema")

1473 self.assertEqual(schema, schema)

1474

1475 # Test inequality

1476 tab2 = tab1.copy()

1477 names = list(tab2.dtype.names)

1478 names[0] = "index2"

1479 tab2.dtype.names = names

1480 schema2 = ArrowNumpySchema(tab2.dtype)

1481 self.assertNotEqual(schema2, schema)

1482

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

1484 def testNumpyDictConversions(self):

1485 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1486

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

1488 tab1_arrow = numpy_to_arrow(tab1)

1489 tab1_dict = arrow_to_numpy_dict(tab1_arrow)

1490 tab1_dict_arrow = numpy_dict_to_arrow(tab1_dict)

1491

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

1493 self.assertEqual(tab1_arrow, tab1_dict_arrow)

1494

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

1496 def testWriteNumpyTableReadAsArrowTable(self):

1497 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1498

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

1500

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

1502

1503 tab2_numpy = arrow_to_numpy(tab2)

1504

1505 _checkNumpyTableEquality(tab1, tab2_numpy)

1506

1507 # Check reading the columns.

1508 columns = tab2.schema.names

1509 columns2 = self.butler.get(

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

1511 )

1512 self.assertEqual(columns2, columns)

1513

1514 # Check reading the schema.

1515 schema = tab2.schema

1516 schema2 = self.butler.get(

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

1518 )

1519 self.assertEqual(schema2, schema)

1520

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

1522 def testWriteNumpyTableReadAsDataFrame(self):

1523 tab1 = _makeSimpleNumpyTable()

1524

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

1526

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

1528

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

1530 # and changes the datatype of the string column.

1531

1532 tab1_df = pd.DataFrame(tab1)

1533

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

1535

1536 # Check reading the columns.

1537 columns = tab2.columns

1538 columns2 = self.butler.get(

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

1540 )

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

1542

1543 # Check reading the schema.

1544 schema = DataFrameSchema(tab2)

1545 schema2 = self.butler.get(

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

1547 )

1548

1549 self.assertEqual(schema2, schema)

1550

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

1552 def testWriteNumpyTableReadAsAstropyTable(self):

1553 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1554

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

1556

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

1558 tab2_numpy = tab2.as_array()

1559

1560 _checkNumpyTableEquality(tab1, tab2_numpy)

1561

1562 # Check reading the columns.

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

1564 columns2 = self.butler.get(

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

1566 )

1567 self.assertEqual(columns2, columns)

1568

1569 # Check reading the schema.

1570 schema = ArrowAstropySchema(tab2)

1571 schema2 = self.butler.get(

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

1573 )

1574

1575 self.assertEqual(schema2, schema)

1576

1577 def testWriteNumpyTableReadAsNumpyDict(self):

1578 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1579

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

1581

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

1583 tab2_numpy = _numpy_dict_to_numpy(tab2)

1584

1585 _checkNumpyTableEquality(tab1, tab2_numpy)

1586

1587 def testBadNumpyColumnParquet(self):

1588 tab1 = _makeSimpleAstropyTable()

1589

1590 # Make a column with mixed type.

1591 bad_col1 = [0.0] * len(tab1)

1592 bad_col1[1] = 0.0 * units.nJy

1593 bad_tab = tab1.copy()

1594 bad_tab["bad_col1"] = bad_col1

1595

1596 bad_tab_np = bad_tab.as_array()

1597

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

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

1600 with self.assertRaises(RuntimeError):

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

1602

1603 # Make a column with ragged size.

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

1605 bad_col2[1] = [0, 0]

1606 bad_tab = tab1.copy()

1607 bad_tab["bad_col2"] = bad_col2

1608

1609 bad_tab_np = bad_tab.as_array()

1610

1611 with self.assertRaises(RuntimeError):

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

1613

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

1615 def testWriteReadAstropyTableLossless(self):

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

1617

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

1619

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

1621

1622 _checkAstropyTableEquality(tab1, tab2)

1623

1624

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

1626class InMemoryArrowNumpyDelegateTestCase(ParquetFormatterArrowNumpyTestCase):

1627 """Tests for InMemoryDatastore, using ArrowTableDelegate with

1628 Numpy table.

1629 """

1630

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

1632

1633 def testBadNumpyColumnParquet(self):

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

1635 pass

1636

1637 def testBadInput(self):

1638 tab1 = _makeSimpleNumpyTable()

1639 delegate = ArrowTableDelegate("ArrowNumpy")

1640

1641 with self.assertRaises(ValueError):

1642 delegate.handleParameters(inMemoryDataset="not_a_numpy_table")

1643

1644 with self.assertRaises(NotImplementedError):

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

1646

1647 with self.assertRaises(AttributeError):

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

1649

1650 def testStorageClass(self):

1651 tab1 = _makeSimpleNumpyTable()

1652

1653 factory = StorageClassFactory()

1654 factory.addFromConfig(StorageClassConfig())

1655

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

1657 # Force the name lookup to do name matching.

1658 storageClass._pytype = None

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

1660

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

1662 # Force the name lookup to do name matching.

1663 storageClass._pytype = None

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

1665

1666

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

1668class ParquetFormatterArrowTableTestCase(unittest.TestCase):

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

1670

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

1672

1673 def setUp(self):

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

1675 self.root = makeTestTempDir(TESTDIR)

1676 config = Config(self.configFile)

1677 self.butler = Butler.from_config(

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

1679 )

1680 self.enterContext(self.butler)

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

1682 # inserting dimension data or defining data IDs.

1683 self.datasetType = DatasetType(

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

1685 )

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

1687

1688 def tearDown(self):

1689 removeTestTempDir(self.root)

1690

1691 def testArrowTable(self):

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

1693

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

1695 # Read the whole Table.

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

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

1698 # comparisons.

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

1700 tab1_np = arrow_to_numpy(tab1)

1701 tab2_np = arrow_to_numpy(tab2)

1702 for col in tab1.column_names:

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

1704 # Read the columns.

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

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

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

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

1709 # Read the rowcount.

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

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

1712 # Read the schema.

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

1714 self.assertEqual(schema, tab1.schema)

1715 # Read just some columns a few different ways.

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

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

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

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

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

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

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

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

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

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

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

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

1728 # Passing an unrecognized column should be a ValueError.

1729 with self.assertRaises(ValueError):

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

1731

1732 def testEmptyArrowTable(self):

1733 data = _makeSimpleNumpyTable()

1734 type_list = _numpy_dtype_to_arrow_types(data.dtype)

1735

1736 schema = pa.schema(type_list)

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

1738

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

1740

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

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

1743 self.assertEqual(tab2, tab1)

1744

1745 tab1_numpy = arrow_to_numpy(tab1)

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

1747 tab1_numpy_arrow = numpy_to_arrow(tab1_numpy)

1748 self.assertEqual(tab1_numpy_arrow, tab1)

1749

1750 tab1_pandas = arrow_to_pandas(tab1)

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

1752 tab1_pandas_arrow = pandas_to_arrow(tab1_pandas)

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

1754 # through empty pandas dataframes.

1755 self.assertEqual(

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

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

1758 )

1759

1760 tab1_astropy = arrow_to_astropy(tab1)

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

1762 tab1_astropy_arrow = astropy_to_arrow(tab1_astropy)

1763 self.assertEqual(tab1_astropy_arrow, tab1)

1764

1765 def testEmptyArrowTableMultidim(self):

1766 data = _makeSimpleNumpyTable(include_multidim=True)

1767 type_list = _numpy_dtype_to_arrow_types(data.dtype)

1768

1769 md = {}

1770 for name in data.dtype.names:

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

1772

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

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

1775

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

1777

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

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

1780 self.assertEqual(tab2, tab1)

1781

1782 tab1_numpy = arrow_to_numpy(tab1)

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

1784 tab1_numpy_arrow = numpy_to_arrow(tab1_numpy)

1785 self.assertEqual(tab1_numpy_arrow, tab1)

1786

1787 tab1_astropy = arrow_to_astropy(tab1)

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

1789 tab1_astropy_arrow = astropy_to_arrow(tab1_astropy)

1790 self.assertEqual(tab1_astropy_arrow, tab1)

1791

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

1793 def testWriteArrowTableReadAsSingleIndexDataFrame(self):

1794 df1, allColumns = _makeSingleIndexDataFrame()

1795

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

1797

1798 # Read back out as a dataframe.

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

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

1801

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

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

1804 df3 = arrow_to_pandas(tab3)

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

1806

1807 # Check reading the columns.

1808 columns = df2.reset_index().columns

1809 columns2 = self.butler.get(

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

1811 )

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

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

1814

1815 # Check reading the schema.

1816 schema = DataFrameSchema(df1)

1817 schema2 = self.butler.get(

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

1819 )

1820 self.assertEqual(schema2, schema)

1821

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

1823 def testWriteArrowTableReadAsMultiIndexDataFrame(self):

1824 df1 = _makeMultiIndexDataFrame()

1825

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

1827

1828 # Read back out as a dataframe.

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

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

1831

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

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

1834 df3 = arrow_to_pandas(atab3)

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

1836

1837 # Check reading the columns.

1838 columns = df2.columns

1839 columns2 = self.butler.get(

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

1841 )

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

1843

1844 # Check reading the schema.

1845 schema = DataFrameSchema(df1)

1846 schema2 = self.butler.get(

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

1848 )

1849 self.assertEqual(schema2, schema)

1850

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

1852 def testWriteArrowTableReadAsAstropyTable(self):

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

1854

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

1856

1857 # Read back out as an astropy table.

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

1859 _checkAstropyTableEquality(tab1, tab2)

1860

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

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

1863 tab3 = arrow_to_astropy(atab3)

1864 _checkAstropyTableEquality(tab1, tab3)

1865

1866 # Check reading the columns.

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

1868 columns2 = self.butler.get(

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

1870 )

1871 self.assertEqual(columns2, columns)

1872

1873 # Check reading the schema.

1874 schema = ArrowAstropySchema(tab1)

1875 schema2 = self.butler.get(

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

1877 )

1878 self.assertEqual(schema2, schema)

1879

1880 # Check the schema conversions and units.

1881 arrow_schema = schema.to_arrow_schema()

1882 for name in arrow_schema.names:

1883 field_metadata = arrow_schema.field(name).metadata

1884 if (

1885 b"description" in field_metadata

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

1887 ):

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

1889 else:

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

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

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

1893

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

1895 def testWriteArrowTableReadAsNumpyTable(self):

1896 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1897

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

1899

1900 # Read back out as a numpy table.

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

1902 _checkNumpyTableEquality(tab1, tab2)

1903

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

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

1906 tab3 = arrow_to_numpy(atab3)

1907 _checkNumpyTableEquality(tab1, tab3)

1908

1909 # Check reading the columns.

1910 columns = list(tab2.dtype.names)

1911 columns2 = self.butler.get(

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

1913 )

1914 self.assertEqual(columns2, columns)

1915

1916 # Check reading the schema.

1917 schema = ArrowNumpySchema(tab1.dtype)

1918 schema2 = self.butler.get(

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

1920 )

1921 self.assertEqual(schema2, schema)

1922

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

1924 def testWriteArrowTableReadAsNumpyDict(self):

1925 tab1 = _makeSimpleNumpyTable(include_multidim=True)

1926

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

1928

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

1930 tab2_numpy = _numpy_dict_to_numpy(tab2)

1931 _checkNumpyTableEquality(tab1, tab2_numpy)

1932

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

1934 def testWriteReadAstropyTableLossless(self):

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

1936

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

1938

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

1940

1941 _checkAstropyTableEquality(tab1, tab2)

1942

1943

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

1945class InMemoryArrowTableDelegateTestCase(ParquetFormatterArrowTableTestCase):

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

1947

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

1949

1950 def testBadInput(self):

1951 tab1 = _makeSimpleArrowTable()

1952 delegate = ArrowTableDelegate("ArrowTable")

1953

1954 with self.assertRaises(ValueError):

1955 delegate.handleParameters(inMemoryDataset="not_an_arrow_table")

1956

1957 with self.assertRaises(NotImplementedError):

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

1959

1960 with self.assertRaises(AttributeError):

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

1962

1963 def testStorageClass(self):

1964 tab1 = _makeSimpleArrowTable()

1965

1966 factory = StorageClassFactory()

1967 factory.addFromConfig(StorageClassConfig())

1968

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

1970 # Force the name lookup to do name matching.

1971 storageClass._pytype = None

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

1973

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

1975 # Force the name lookup to do name matching.

1976 storageClass._pytype = None

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

1978

1979

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

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

1982class ParquetFormatterArrowNumpyDictTestCase(unittest.TestCase):

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

1984

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

1986

1987 def setUp(self):

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

1989 self.root = makeTestTempDir(TESTDIR)

1990 config = Config(self.configFile)

1991 self.butler = Butler.from_config(

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

1993 )

1994 self.enterContext(self.butler)

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

1996 # inserting dimension data or defining data IDs.

1997 self.datasetType = DatasetType(

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

1999 )

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

2001

2002 def tearDown(self):

2003 removeTestTempDir(self.root)

2004

2005 def testNumpyDict(self):

2006 tab1 = _makeSimpleNumpyTable(include_multidim=True)

2007 dict1 = _numpy_to_numpy_dict(tab1)

2008

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

2010 # Read the whole table.

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

2012 _checkNumpyDictEquality(dict1, dict2)

2013 # Read the columns.

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

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

2016 for name in dict1:

2017 self.assertIn(name, columns2)

2018 # Read the rowcount.

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

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

2021 # Read the schema.

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

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

2024 # Read just some columns a few different ways.

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

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

2027 _checkNumpyDictEquality(subdict, tab3)

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

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

2030 _checkNumpyDictEquality(subdict, tab4)

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

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

2033 _checkNumpyDictEquality(subdict, tab5)

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

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

2036 _checkNumpyDictEquality(subdict, tab6)

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

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

2039 _checkNumpyDictEquality(subdict, tab7)

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

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

2042 _checkNumpyDictEquality(subdict, tab8)

2043 # Passing an unrecognized column should be a ValueError.

2044 with self.assertRaises(ValueError):

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

2046

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

2048 def testWriteNumpyDictReadAsArrowTable(self):

2049 tab1 = _makeSimpleNumpyTable(include_multidim=True)

2050 dict1 = _numpy_to_numpy_dict(tab1)

2051

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

2053

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

2055

2056 tab2_dict = arrow_to_numpy_dict(tab2)

2057

2058 _checkNumpyDictEquality(dict1, tab2_dict)

2059

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

2061 def testWriteNumpyDictReadAsDataFrame(self):

2062 tab1 = _makeSimpleNumpyTable()

2063 dict1 = _numpy_to_numpy_dict(tab1)

2064

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

2066

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

2068

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

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

2071 # changes the datatype of the string column.

2072 tab1_df = pd.DataFrame(tab1)

2073

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

2075 for col in tab1_df.columns:

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

2077

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

2079 def testWriteNumpyDictReadAsAstropyTable(self):

2080 tab1 = _makeSimpleNumpyTable(include_multidim=True)

2081 dict1 = _numpy_to_numpy_dict(tab1)

2082

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

2084

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

2086 tab2_dict = _astropy_to_numpy_dict(tab2)

2087

2088 _checkNumpyDictEquality(dict1, tab2_dict)

2089

2090 def testWriteNumpyDictReadAsNumpyTable(self):

2091 tab1 = _makeSimpleNumpyTable(include_multidim=True)

2092 dict1 = _numpy_to_numpy_dict(tab1)

2093

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

2095

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

2097 tab2_dict = _numpy_to_numpy_dict(tab2)

2098

2099 _checkNumpyDictEquality(dict1, tab2_dict)

2100

2101 def testWriteNumpyDictBad(self):

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

2103 with self.assertRaises(RuntimeError):

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

2105

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

2107 with self.assertRaises(RuntimeError):

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

2109

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

2111 with self.assertRaises(RuntimeError):

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

2113

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

2115 with self.assertRaises(RuntimeError):

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

2117

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

2119 def testWriteReadAstropyTableLossless(self):

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

2121

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

2123

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

2125

2126 _checkAstropyTableEquality(tab1, tab2)

2127

2128

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

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

2131class InMemoryNumpyDictDelegateTestCase(ParquetFormatterArrowNumpyDictTestCase):

2132 """Tests for InMemoryDatastore, using ArrowTableDelegate with

2133 Numpy dict.

2134 """

2135

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

2137

2138 def testWriteNumpyDictBad(self):

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

2140 pass

2141

2142

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

2144class ParquetFormatterArrowSchemaTestCase(unittest.TestCase):

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

2146

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

2148

2149 def setUp(self):

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

2151 self.root = makeTestTempDir(TESTDIR)

2152 config = Config(self.configFile)

2153 self.butler = Butler.from_config(

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

2155 )

2156 self.enterContext(self.butler)

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

2158 # inserting dimension data or defining data IDs.

2159 self.datasetType = DatasetType(

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

2161 )

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

2163

2164 def tearDown(self):

2165 removeTestTempDir(self.root)

2166

2167 def _makeTestSchema(self):

2168 schema = pa.schema(

2169 [

2170 pa.field(

2171 "int32",

2172 pa.int32(),

2173 nullable=False,

2174 metadata={

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

2176 "unit": "",

2177 },

2178 ),

2179 pa.field(

2180 "int64",

2181 pa.int64(),

2182 nullable=False,

2183 metadata={

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

2185 "unit": "",

2186 },

2187 ),

2188 pa.field(

2189 "uint64",

2190 pa.uint64(),

2191 nullable=False,

2192 metadata={

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

2194 "unit": "",

2195 },

2196 ),

2197 pa.field(

2198 "float32",

2199 pa.float32(),

2200 nullable=False,

2201 metadata={

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

2203 "unit": "count",

2204 },

2205 ),

2206 pa.field(

2207 "float64",

2208 pa.float64(),

2209 nullable=False,

2210 metadata={

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

2212 "unit": "nJy",

2213 },

2214 ),

2215 pa.field(

2216 "fixed_size_list",

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

2218 nullable=False,

2219 metadata={

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

2221 "unit": "nJy",

2222 },

2223 ),

2224 pa.field(

2225 "variable_size_list",

2226 pa.list_(pa.float64()),

2227 nullable=False,

2228 metadata={

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

2230 "unit": "nJy",

2231 },

2232 ),

2233 # One of these fields will have no description.

2234 pa.field(

2235 "string",

2236 pa.string(),

2237 nullable=False,

2238 metadata={

2239 "unit": "",

2240 },

2241 ),

2242 # One of these fields will have no metadata.

2243 pa.field(

2244 "binary",

2245 pa.binary(),

2246 nullable=False,

2247 ),

2248 ]

2249 )

2250

2251 return schema

2252

2253 def testArrowSchema(self):

2254 schema1 = self._makeTestSchema()

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

2256

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

2258 self.assertEqual(schema2, schema1)

2259

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

2261 def testWriteArrowSchemaReadAsDataFrameSchema(self):

2262 schema1 = self._makeTestSchema()

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

2264

2265 df_schema1 = DataFrameSchema.from_arrow(schema1)

2266

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

2268 self.assertEqual(df_schema2, df_schema1)

2269

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

2271 def testWriteArrowSchemaReadAsArrowAstropySchema(self):

2272 schema1 = self._makeTestSchema()

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

2274

2275 ap_schema1 = ArrowAstropySchema.from_arrow(schema1)

2276

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

2278 self.assertEqual(ap_schema2, ap_schema1)

2279

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

2281 for name in schema1.names:

2282 field_metadata = schema1.field(name).metadata

2283 if field_metadata is None:

2284 continue

2285 if (

2286 b"description" in field_metadata

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

2288 ):

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

2290 else:

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

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

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

2294

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

2296 def testWriteArrowSchemaReadAsArrowNumpySchema(self):

2297 schema1 = self._makeTestSchema()

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

2299

2300 np_schema1 = ArrowNumpySchema.from_arrow(schema1)

2301

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

2303 self.assertEqual(np_schema2, np_schema1)

2304

2305

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

2307class InMemoryArrowSchemaDelegateTestCase(ParquetFormatterArrowSchemaTestCase):

2308 """Tests for InMemoryDatastore and ArrowSchema."""

2309

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

2311

2312

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

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

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

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

2317class ParquetFormatterArrowTableS3TestCase(unittest.TestCase):

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

2319

2320 # Code is adapted from test_butler.py

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

2322 fullConfigKey = None

2323 validationCanFail = True

2324

2325 bucketName = "anybucketname"

2326

2327 root = "butlerRoot/"

2328

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

2330

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

2332

2333 registryStr = "/gen3.sqlite3"

2334

2335 mock_aws = mock_aws()

2336

2337 def setUp(self):

2338 self.root = makeTestTempDir(TESTDIR)

2339

2340 config = Config(self.configFile)

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

2342 self.bucketName = uri.netloc

2343

2344 # Enable S3 mocking of tests.

2345 self.enterContext(clean_test_environment_for_s3())

2346 self.mock_aws.start()

2347

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

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

2350

2351 # need local folder to store registry database

2352 self.reg_dir = makeTestTempDir(TESTDIR)

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

2354

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

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

2357 s3 = boto3.resource("s3")

2358 s3.create_bucket(Bucket=self.bucketName)

2359

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

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

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

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

2364

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

2366 self.enterContext(self.butler)

2367

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

2369 # inserting dimension data or defining data IDs.

2370 self.datasetType = DatasetType(

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

2372 )

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

2374

2375 def tearDown(self):

2376 s3 = boto3.resource("s3")

2377 bucket = s3.Bucket(self.bucketName)

2378 try:

2379 bucket.objects.all().delete()

2380 except botocore.exceptions.ClientError as e:

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

2382 # the key was not reachable - pass

2383 pass

2384 else:

2385 raise

2386

2387 bucket = s3.Bucket(self.bucketName)

2388 bucket.delete()

2389

2390 # Stop the S3 mock.

2391 self.mock_aws.stop()

2392

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

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

2395

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

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

2398

2399 def testArrowTableS3(self):

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

2401

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

2403

2404 # Read the whole Table.

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

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

2407 # comparisons.

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

2409 tab1_np = arrow_to_numpy(tab1)

2410 tab2_np = arrow_to_numpy(tab2)

2411 for col in tab1.column_names:

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

2413 # Read the columns.

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

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

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

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

2418 # Read the rowcount.

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

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

2421 # Read the schema.

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

2423 self.assertEqual(schema, tab1.schema)

2424 # Read just some columns a few different ways.

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

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

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

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

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

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

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

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

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

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

2435 # Passing an unrecognized column should be a ValueError.

2436 with self.assertRaises(ValueError):

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

2438

2439

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

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

2442class ComputeRowGroupSizeTestCase(unittest.TestCase):

2443 """Tests for compute_row_group_size."""

2444

2445 def testRowGroupSizeNoMetadata(self):

2446 numpyTable = _makeSimpleNumpyTable(include_multidim=True)

2447

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

2449 # that adds metadata.

2450 type_list = _numpy_dtype_to_arrow_types(numpyTable.dtype)

2451 schema = pa.schema(type_list)

2452 arrays = _numpy_style_arrays_to_arrow_arrays(

2453 numpyTable.dtype,

2454 len(numpyTable),

2455 numpyTable,

2456 schema,

2457 )

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

2459

2460 row_group_size = compute_row_group_size(arrowTable.schema)

2461

2462 self.assertGreater(row_group_size, 1_000_000)

2463 self.assertLess(row_group_size, 2_000_000)

2464

2465 def testRowGroupSizeWithMetadata(self):

2466 numpyTable = _makeSimpleNumpyTable(include_multidim=True)

2467

2468 arrowTable = numpy_to_arrow(numpyTable)

2469

2470 row_group_size = compute_row_group_size(arrowTable.schema)

2471

2472 self.assertGreater(row_group_size, 1_000_000)

2473 self.assertLess(row_group_size, 2_000_000)

2474

2475 def testRowGroupSizeTinyTable(self):

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

2477

2478 arrowTable = numpy_to_arrow(numpyTable)

2479

2480 row_group_size = compute_row_group_size(arrowTable.schema)

2481

2482 self.assertGreater(row_group_size, 1_000_000)

2483

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

2485 def testRowGroupSizeDataFrameWithLists(self):

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

2487 arrowTable = pandas_to_arrow(df)

2488 row_group_size = compute_row_group_size(arrowTable.schema)

2489

2490 self.assertGreater(row_group_size, 1_000_000)

2491

2492

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

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

2495

2496 Parameters

2497 ----------

2498 table1 : `astropy.table.Table`

2499 table2 : `astropy.table.Table`

2500 skip_units : `bool`

2501 has_bigendian : `bool`

2502 """

2503 if not has_bigendian:

2504 assert table1.dtype == table2.dtype

2505 else:

2506 for name in table1.dtype.names:

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

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

2509

2510 # Strip provenance before comparison.

2511 DatasetProvenance.strip_provenance_from_flat_dict(table1.meta)

2512 DatasetProvenance.strip_provenance_from_flat_dict(table2.meta)

2513 assert table1.meta == table2.meta

2514 if not skip_units:

2515 for name in table1.columns:

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

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

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

2519

2520 for name in table1.columns:

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

2522 has_masked = False

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

2524 c1 = table1[name].filled()

2525 has_masked = True

2526 else:

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

2528 if has_masked:

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

2530 c2 = table2[name].filled()

2531 else:

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

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

2534 np.testing.assert_array_equal(c1, c2)

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

2536 if has_masked:

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

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

2539

2540

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

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

2543

2544 Parameters

2545 ----------

2546 table1 : `numpy.ndarray`

2547 table2 : `numpy.ndarray`

2548 has_bigendian : `bool`

2549 """

2550 assert table1.dtype.names == table2.dtype.names

2551 for name in table1.dtype.names:

2552 if not has_bigendian:

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

2554 else:

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

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

2557 assert np.all(table1 == table2)

2558

2559

2560def _checkNumpyDictEquality(dict1, dict2):

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

2562

2563 Parameters

2564 ----------

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

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

2567 """

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

2569 for name in dict1:

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

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

2572

2573

2574if __name__ == "__main__":

2575 unittest.main()