Coverage for tests/test_parquet.py: 99%

1# This file is part of daf_butler. 

2# 

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. 

8# 

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

27 

28"""Tests for ParquetFormatter. 

29 

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

31""" 

32 

33import datetime 

34import os 

35import unittest 

36import uuid 

37 

38try: 

39 import pyarrow as pa 

40except ImportError: 

41 pa = None 

42try: 

43 import astropy.table as atable 

44 from astropy import units 

45except ImportError: 

46 atable = None 

47try: 

48 import numpy as np 

49except ImportError: 

50 np = None 

51try: 

52 import pandas as pd 

53except ImportError: 

54 pd = None 

55 

56try: 

57 import fsspec 

58except ImportError: 

59 fsspec = None 

60 

61 

62from lsst.daf.butler import ( 

63 Butler, 

64 Config, 

65 DatasetProvenance, 

66 DatasetRef, 

67 DatasetType, 

68 FileDataset, 

69 StorageClassConfig, 

70 StorageClassFactory, 

71) 

72 

73try: 

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

75except ImportError: 

76 pa = None 

77 

78try: 

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

80 ASTROPY_PANDAS_INDEX_KEY, 

81 ArrowAstropySchema, 

82 ArrowNumpySchema, 

83 DataFrameSchema, 

84 ParquetFormatter, 

85 _append_numpy_multidim_metadata, 

86 _astropy_to_numpy_dict, 

87 _numpy_dict_to_numpy, 

88 _numpy_dtype_to_arrow_types, 

89 _numpy_style_arrays_to_arrow_arrays, 

90 _numpy_to_numpy_dict, 

91 add_pandas_index_to_astropy, 

92 arrow_to_astropy, 

93 arrow_to_numpy, 

94 arrow_to_numpy_dict, 

95 arrow_to_pandas, 

96 astropy_to_arrow, 

97 astropy_to_pandas, 

98 compute_row_group_size, 

99 numpy_dict_to_arrow, 

100 numpy_to_arrow, 

101 pandas_to_arrow, 

102 pandas_to_astropy, 

103 ) 

104except ImportError: 

105 pa = None 

106 pd = None 

107 atable = None 

108 np = None 

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

110from lsst.resources.tests import make_remote_test_uri 

111 

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

113 

114 

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

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

117 

118 Parameters 

119 ---------- 

120 include_multidim : `bool` 

121 Include multi-dimensional columns. 

122 include_bigendian : `bool` 

123 Include big-endian columns. 

124 

125 Returns 

126 ------- 

127 numpyTable : `numpy.ndarray` 

128 """ 

129 nrow = 5 

130 

131 dtype = [ 

132 ("index", "i4"), 

133 ("a", "f8"), 

134 ("b", "f8"), 

135 ("c", "f8"), 

136 ("ddd", "f8"), 

137 ("f", "i8"), 

138 ("strcol", "U10"), 

139 ("bytecol", "S10"), 

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

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

142 ] 

143 

144 if include_multidim: 

145 dtype.extend( 

146 [ 

147 ("d1", "f4", (5,)), 

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

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

150 ] 

151 ) 

152 

153 if include_bigendian: 

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

155 

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

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

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

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

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

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

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

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

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

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

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

167 

168 if include_multidim: 

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

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

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

172 

173 if include_bigendian: 

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

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

176 

177 return data 

178 

179 

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

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

182 

183 Parameters 

184 ---------- 

185 include_masked : `bool` 

186 Include masked columns. 

187 include_lists : `bool` 

188 Include list columns. 

189 

190 Returns 

191 ------- 

192 dataFrame : `~pandas.DataFrame` 

193 The test dataframe. 

194 allColumns : `list` [`str`] 

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

196 """ 

197 data = _makeSimpleNumpyTable() 

198 df = pd.DataFrame(data) 

199 df = df.set_index("index") 

200 

201 if include_masked: 

202 nrow = len(df) 

203 

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

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

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

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

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

209 

210 if include_lists: 

211 nrow = len(df) 

212 

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

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

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

216 

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

218 

219 return df, allColumns 

220 

221 

222def _makeMultiIndexDataFrame(): 

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

224 

225 Returns 

226 ------- 

227 dataFrame : `~pandas.DataFrame` 

228 The test dataframe. 

229 """ 

230 columns = pd.MultiIndex.from_tuples( 

231 [ 

232 ("g", "a"), 

233 ("g", "b"), 

234 ("g", "c"), 

235 ("r", "a"), 

236 ("r", "b"), 

237 ("r", "c"), 

238 ], 

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

240 ) 

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

242 

243 return df 

244 

245 

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

247 """Make an astropy table for testing. 

248 

249 Parameters 

250 ---------- 

251 include_multidim : `bool` 

252 Include multi-dimensional columns. 

253 include_masked : `bool` 

254 Include masked columns. 

255 include_bigendian : `bool` 

256 Include big-endian columns. 

257 

258 Returns 

259 ------- 

260 astropyTable : `astropy.table.Table` 

261 The test table. 

262 """ 

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

264 # Add a couple of units. 

265 table = atable.Table(data) 

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

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

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

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

270 

271 # Add some masked columns. 

272 if include_masked: 

273 nrow = len(table) 

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

275 mask[1] = True 

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

277 # to be able test after serialization. 

278 

279 # Masked 64-bit integer. 

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

281 arr[mask] = -1 

282 arr[0] = 1649900760361600113 

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

284 # Masked 32-bit float. 

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

286 arr[mask] = np.nan 

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

288 # Unmasked 32-bit float with NaNs. 

289 table["um_f4"] = arr 

290 # Masked 64-bit float. 

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

292 arr[mask] = np.nan 

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

294 # Unmasked 64-bit float with NaNs. 

295 table["um_f8"] = arr 

296 # Masked boolean. 

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

298 arr[mask] = True 

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

300 # Masked unsigned 32-bit unsigned int. 

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

302 arr[mask] = 0 

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

304 # Masked string. 

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

306 # Masked bytes. 

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

308 

309 return table 

310 

311 

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

313 """Make an arrow table for testing. 

314 

315 Parameters 

316 ---------- 

317 include_multidim : `bool` 

318 Include multi-dimensional columns. 

319 include_masked : `bool` 

320 Include masked columns. 

321 

322 Returns 

323 ------- 

324 arrowTable : `pyarrow.Table` 

325 The test table. 

326 """ 

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

328 return astropy_to_arrow(data) 

329 

330 

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

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

333class ParquetFormatterDataFrameTestCase(unittest.TestCase): 

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

335 

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

337 

338 def setUp(self): 

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

340 self.root = makeTestTempDir(TESTDIR) 

341 config = Config(self.configFile) 

342 self.run = "test_run" 

343 self.butler = Butler.from_config( 

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

345 ) 

346 self.enterContext(self.butler) 

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

348 # inserting dimension data or defining data IDs. 

349 self.datasetType = DatasetType( 

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

351 ) 

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

353 

354 def tearDown(self): 

355 removeTestTempDir(self.root) 

356 

357 def testSingleIndexDataFrame(self): 

358 df1, allColumns = _makeSingleIndexDataFrame(include_masked=True) 

359 

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

361 # Read the whole DataFrame. 

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

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

364 # Read just the column descriptions. 

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

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

367 # Read the rowcount. 

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

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

370 # Read the schema. 

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

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

373 # Read just some columns a few different ways. 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

388 # Passing an unrecognized column should be a ValueError. 

389 with self.assertRaises(ValueError): 

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

391 

392 def testSingleIndexDataFrameWithLists(self): 

393 df1, allColumns = _makeSingleIndexDataFrame(include_lists=True) 

394 

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

396 # Read the whole DataFrame. 

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

398 

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

400 # from lists to arrays. 

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

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

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

404 

405 def testMultiIndexDataFrame(self): 

406 df1 = _makeMultiIndexDataFrame() 

407 

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

409 # Read the whole DataFrame. 

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

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

412 # Read just the column descriptions. 

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

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

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

416 # Read the rowcount. 

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

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

419 # Read the schema. 

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

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

422 # Read just some columns a few different ways. 

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

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

425 df4 = self.butler.get( 

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

427 ) 

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

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

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

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

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

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

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

435 # Passing an unrecognized column should be a ValueError. 

436 with self.assertRaises(ValueError): 

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

438 

439 def testSingleIndexDataFrameEmptyString(self): 

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

441 df1, _ = _makeSingleIndexDataFrame() 

442 

443 # Set one of the strings to None 

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

445 

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

447 # Read the whole DataFrame. 

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

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

450 

451 def testSingleIndexDataFrameAllEmptyStrings(self): 

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

453 column. 

454 """ 

455 df1, _ = _makeSingleIndexDataFrame() 

456 

457 # Set all of the strings to None 

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

459 

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

461 # Read the whole DataFrame. 

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

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

464 

465 def testLegacyDataFrame(self): 

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

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

468 functionality. 

469 """ 

470 df1, allColumns = _makeSingleIndexDataFrame() 

471 

472 if isinstance(df1.index, pd.RangeIndex): 472 ↛ 478line 472 didn't jump to line 478 because the condition on line 472 was never true

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

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

475 # Unfortunately, parquet files serialized directly with 

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

477 # they are sequential integers. 

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

479 

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

481 df1.to_parquet(fname) 

482 

483 legacy_type = DatasetType( 

484 "legacy_dataframe", 

485 dimensions=(), 

486 storageClass="DataFrame", 

487 universe=self.butler.dimensions, 

488 ) 

489 self.butler.registry.registerDatasetType(legacy_type) 

490 

491 data_id = {} 

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

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

494 

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

496 

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

498 

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

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

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

502 

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

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

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

506 

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

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

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

510 

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

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

513 self.assertEqual(rowcount2a, rowcount2b) 

514 

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

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

517 self.assertEqual(schema2a, schema2b) 

518 

519 def testDataFrameSchema(self): 

520 tab1 = _makeSimpleArrowTable() 

521 

522 schema = DataFrameSchema.from_arrow(tab1.schema) 

523 

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

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

526 self.assertNotEqual(schema, "not_a_schema") 

527 self.assertEqual(schema, schema) 

528 

529 tab2 = _makeMultiIndexDataFrame() 

530 schema2 = DataFrameSchema(tab2) 

531 

532 self.assertNotEqual(schema, schema2) 

533 

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

535 def testWriteSingleIndexDataFrameReadAsAstropyTable(self): 

536 df1, allColumns = _makeSingleIndexDataFrame() 

537 

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

539 

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

541 

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

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

544 

545 # Check reading the columns. 

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

547 columns2 = self.butler.get( 

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

549 ) 

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

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

552 

553 # Check reading the schema. 

554 schema = ArrowAstropySchema(tab2) 

555 schema2 = self.butler.get( 

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

557 ) 

558 

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

560 # will be changed because of pandas indexing. 

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

562 for name in schema.schema.columns: 

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

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

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

566 

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

568 def testWriteSingleIndexDataFrameWithMaskedColsReadAsAstropyTable(self): 

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

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

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

572 # inconsistent.) 

573 df1, allColumns = _makeSingleIndexDataFrame(include_masked=True) 

574 

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

576 

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

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

579 

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

581 for name in tab2_df.columns: 

582 col1 = df1[name] 

583 col2 = tab2_df[name] 

584 

585 if col1.hasnans: 

586 notNull = col1.notnull() 

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

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

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

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

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

592 else: 

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

594 

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

596 def testWriteMultiIndexDataFrameReadAsAstropyTable(self): 

597 df1 = _makeMultiIndexDataFrame() 

598 

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

600 

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

602 

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

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

605 # recommended. 

606 

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

608 def testWriteAstropyTableWithMaskedColsReadAsSingleIndexDataFrame(self): 

609 tab1 = _makeSimpleAstropyTable(include_masked=True) 

610 

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

612 

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

614 

615 tab1_df = astropy_to_pandas(tab1) 

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

617 

618 tab2_astropy = pandas_to_astropy(tab2) 

619 for col in tab1.dtype.names: 

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

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

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

623 

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

625 def testWriteSingleIndexDataFrameReadAsArrowTable(self): 

626 df1, allColumns = _makeSingleIndexDataFrame() 

627 

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

629 

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

631 

632 tab2_df = arrow_to_pandas(tab2) 

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

634 

635 # Check reading the columns. 

636 columns = list(tab2.schema.names) 

637 columns2 = self.butler.get( 

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

639 ) 

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

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

642 

643 # Override the component using a dataset type. 

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

645 "ArrowColumnList" 

646 ) 

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

648 

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

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

651 # selects the component. 

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

653 "columns" 

654 ) 

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

656 

657 # Check reading the schema. 

658 schema = tab2.schema 

659 schema2 = self.butler.get( 

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

661 ) 

662 

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

664 # information be maintained. 

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

666 for name in schema.names: 

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

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

669 

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

671 def testWriteMultiIndexDataFrameReadAsArrowTable(self): 

672 df1 = _makeMultiIndexDataFrame() 

673 

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

675 

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

677 

678 tab2_df = arrow_to_pandas(tab2) 

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

680 

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

682 def testWriteSingleIndexDataFrameReadAsNumpyTable(self): 

683 df1, allColumns = _makeSingleIndexDataFrame() 

684 

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

686 

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

688 

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

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

691 

692 # Check reading the columns. 

693 columns = list(tab2.dtype.names) 

694 columns2 = self.butler.get( 

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

696 )