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._repo_template_cache import make_repo_for_test 

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

111from lsst.resources.tests import make_remote_test_uri 

112 

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

114 

115 

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

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

118 

119 Parameters 

120 ---------- 

121 include_multidim : `bool` 

122 Include multi-dimensional columns. 

123 include_bigendian : `bool` 

124 Include big-endian columns. 

125 

126 Returns 

127 ------- 

128 numpyTable : `numpy.ndarray` 

129 """ 

130 nrow = 5 

131 

132 dtype = [ 

133 ("index", "i4"), 

134 ("a", "f8"), 

135 ("b", "f8"), 

136 ("c", "f8"), 

137 ("ddd", "f8"), 

138 ("f", "i8"), 

139 ("strcol", "U10"), 

140 ("bytecol", "S10"), 

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

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

143 ] 

144 

145 if include_multidim: 

146 dtype.extend( 

147 [ 

148 ("d1", "f4", (5,)), 

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

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

151 ] 

152 ) 

153 

154 if include_bigendian: 

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

156 

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

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

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

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

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

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

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

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

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

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

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

168 

169 if include_multidim: 

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

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

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

173 

174 if include_bigendian: 

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

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

177 

178 return data 

179 

180 

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

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

183 

184 Parameters 

185 ---------- 

186 include_masked : `bool` 

187 Include masked columns. 

188 include_lists : `bool` 

189 Include list columns. 

190 

191 Returns 

192 ------- 

193 dataFrame : `~pandas.DataFrame` 

194 The test dataframe. 

195 allColumns : `list` [`str`] 

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

197 """ 

198 data = _makeSimpleNumpyTable() 

199 df = pd.DataFrame(data) 

200 df = df.set_index("index") 

201 

202 if include_masked: 

203 nrow = len(df) 

204 

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

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

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

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

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

210 

211 if include_lists: 

212 nrow = len(df) 

213 

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

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

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

217 

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

219 

220 return df, allColumns 

221 

222 

223def _makeMultiIndexDataFrame(): 

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

225 

226 Returns 

227 ------- 

228 dataFrame : `~pandas.DataFrame` 

229 The test dataframe. 

230 """ 

231 columns = pd.MultiIndex.from_tuples( 

232 [ 

233 ("g", "a"), 

234 ("g", "b"), 

235 ("g", "c"), 

236 ("r", "a"), 

237 ("r", "b"), 

238 ("r", "c"), 

239 ], 

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

241 ) 

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

243 

244 return df 

245 

246 

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

248 """Make an astropy table for testing. 

249 

250 Parameters 

251 ---------- 

252 include_multidim : `bool` 

253 Include multi-dimensional columns. 

254 include_masked : `bool` 

255 Include masked columns. 

256 include_bigendian : `bool` 

257 Include big-endian columns. 

258 

259 Returns 

260 ------- 

261 astropyTable : `astropy.table.Table` 

262 The test table. 

263 """ 

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

265 # Add a couple of units. 

266 table = atable.Table(data) 

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

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

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

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

271 

272 # Add some masked columns. 

273 if include_masked: 

274 nrow = len(table) 

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

276 mask[1] = True 

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

278 # to be able test after serialization. 

279 

280 # Masked 64-bit integer. 

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

282 arr[mask] = -1 

283 arr[0] = 1649900760361600113 

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

285 # Masked 32-bit float. 

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

287 arr[mask] = np.nan 

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

289 # Unmasked 32-bit float with NaNs. 

290 table["um_f4"] = arr 

291 # Masked 64-bit float. 

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

293 arr[mask] = np.nan 

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

295 # Unmasked 64-bit float with NaNs. 

296 table["um_f8"] = arr 

297 # Masked boolean. 

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

299 arr[mask] = True 

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

301 # Masked unsigned 32-bit unsigned int. 

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

303 arr[mask] = 0 

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

305 # Masked string. 

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

307 # Masked bytes. 

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

309 

310 return table 

311 

312 

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

314 """Make an arrow table for testing. 

315 

316 Parameters 

317 ---------- 

318 include_multidim : `bool` 

319 Include multi-dimensional columns. 

320 include_masked : `bool` 

321 Include masked columns. 

322 

323 Returns 

324 ------- 

325 arrowTable : `pyarrow.Table` 

326 The test table. 

327 """ 

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

329 return astropy_to_arrow(data) 

330 

331 

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

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

334class ParquetFormatterDataFrameTestCase(unittest.TestCase): 

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

336 

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

338 

339 def setUp(self): 

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

341 self.root = makeTestTempDir(TESTDIR) 

342 config = Config(self.configFile) 

343 self.run = "test_run" 

344 self.butler = Butler.from_config( 

345 make_repo_for_test(self.root, config=config), writeable=True, run=self.run 

346 ) 

347 self.enterContext(self.butler) 

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

349 # inserting dimension data or defining data IDs. 

350 self.datasetType = DatasetType( 

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

352 ) 

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

354 

355 def tearDown(self): 

356 removeTestTempDir(self.root) 

357 

358 def testSingleIndexDataFrame(self): 

359 df1, allColumns = _makeSingleIndexDataFrame(include_masked=True) 

360 

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

362 # Read the whole DataFrame. 

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

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

365 # Read just the column descriptions. 

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

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

368 # Read the rowcount. 

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

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

371 # Read the schema. 

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

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

374 # Read just some columns a few different ways. 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

389 # Passing an unrecognized column should be a ValueError. 

390 with self.assertRaises(ValueError): 

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

392 

393 def testSingleIndexDataFrameWithLists(self): 

394 df1, allColumns = _makeSingleIndexDataFrame(include_lists=True) 

395 

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

397 # Read the whole DataFrame. 

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

399 

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

401 # from lists to arrays. 

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

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

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

405 

406 def testMultiIndexDataFrame(self): 

407 df1 = _makeMultiIndexDataFrame() 

408 

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

410 # Read the whole DataFrame. 

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

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

413 # Read just the column descriptions. 

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

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

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

417 # Read the rowcount. 

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

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

420 # Read the schema. 

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

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

423 # Read just some columns a few different ways. 

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

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

426 df4 = self.butler.get( 

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

428 ) 

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

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

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

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

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

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

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

436 # Passing an unrecognized column should be a ValueError. 

437 with self.assertRaises(ValueError): 

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

439 

440 def testSingleIndexDataFrameEmptyString(self): 

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

442 df1, _ = _makeSingleIndexDataFrame() 

443 

444 # Set one of the strings to None 

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

446 

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

448 # Read the whole DataFrame. 

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

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

451 

452 def testSingleIndexDataFrameAllEmptyStrings(self): 

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

454 column. 

455 """ 

456 df1, _ = _makeSingleIndexDataFrame() 

457 

458 # Set all of the strings to None 

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

460 

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

462 # Read the whole DataFrame. 

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

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

465 

466 def testLegacyDataFrame(self): 

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

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

469 functionality. 

470 """ 

471 df1, allColumns = _makeSingleIndexDataFrame() 

472 

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

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

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

476 # Unfortunately, parquet files serialized directly with 

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

478 # they are sequential integers. 

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

480 

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

482 df1.to_parquet(fname) 

483 

484 legacy_type = DatasetType( 

485 "legacy_dataframe", 

486 dimensions=(), 

487 storageClass="DataFrame", 

488 universe=self.butler.dimensions, 

489 ) 

490 self.butler.registry.registerDatasetType(legacy_type) 

491 

492 data_id = {} 

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

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

495 

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

497 

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

499 

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

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

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

503 

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

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

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

507 

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

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

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

511 

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

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

514 self.assertEqual(rowcount2a, rowcount2b) 

515 

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

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

518 self.assertEqual(schema2a, schema2b) 

519 

520 def testDataFrameSchema(self): 

521 tab1 = _makeSimpleArrowTable() 

522 

523 schema = DataFrameSchema.from_arrow(tab1.schema) 

524 

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

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

527 self.assertNotEqual(schema, "not_a_schema") 

528 self.assertEqual(schema, schema) 

529 

530 tab2 = _makeMultiIndexDataFrame() 

531 schema2 = DataFrameSchema(tab2) 

532 

533 self.assertNotEqual(schema, schema2) 

534 

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

536 def testWriteSingleIndexDataFrameReadAsAstropyTable(self): 

537 df1, allColumns = _makeSingleIndexDataFrame() 

538 

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

540 

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

542 

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

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

545 

546 # Check reading the columns. 

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

548 columns2 = self.butler.get( 

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

550 ) 

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

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

553 

554 # Check reading the schema. 

555 schema = ArrowAstropySchema(tab2) 

556 schema2 = self.butler.get( 

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

558 ) 

559 

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

561 # will be changed because of pandas indexing. 

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

563 for name in schema.schema.columns: 

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

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

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

567 

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

569 def testWriteSingleIndexDataFrameWithMaskedColsReadAsAstropyTable(self): 

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

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

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

573 # inconsistent.) 

574 df1, allColumns = _makeSingleIndexDataFrame(include_masked=True) 

575 

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

577 

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

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

580 

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

582 for name in tab2_df.columns: 

583 col1 = df1[name] 

584 col2 = tab2_df[name] 

585 

586 if col1.hasnans: 

587 notNull = col1.notnull() 

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

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

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

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

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

593 else: 

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

595 

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

597 def testWriteMultiIndexDataFrameReadAsAstropyTable(self): 

598 df1 = _makeMultiIndexDataFrame() 

599 

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

601 

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

603 

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

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

606 # recommended. 

607 

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

609 def testWriteAstropyTableWithMaskedColsReadAsSingleIndexDataFrame(self): 

610 tab1 = _makeSimpleAstropyTable(include_masked=True) 

611 

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

613 

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

615 

616 tab1_df = astropy_to_pandas(tab1) 

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

618 

619 tab2_astropy = pandas_to_astropy(tab2) 

620 for col in tab1.dtype.names: 

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

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

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

624 

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

626 def testWriteSingleIndexDataFrameReadAsArrowTable(self): 

627 df1, allColumns = _makeSingleIndexDataFrame() 

628 

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

630 

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

632 

633 tab2_df = arrow_to_pandas(tab2) 

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

635 

636 # Check reading the columns. 

637 columns = list(tab2.schema.names) 

638 columns2 = self.butler.get( 

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

640 ) 

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

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

643 

644 # Override the component using a dataset type. 

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

646 "ArrowColumnList" 

647 ) 

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

649 

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

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

652 # selects the component. 

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

654 "columns" 

655 ) 

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

657 

658 # Check reading the schema. 

659 schema = tab2.schema 

660 schema2 = self.butler.get( 

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

662 ) 

663 

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

665 # information be maintained. 

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

667 for name in schema.names: 

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

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

670 

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

672 def testWriteMultiIndexDataFrameReadAsArrowTable(self): 

673 df1 = _makeMultiIndexDataFrame() 

674 

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

676 

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

678 

679 tab2_df = arrow_to_pandas(tab2) 

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

681 

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

683 def testWriteSingleIndexDataFrameReadAsNumpyTable(self): 

684 df1, allColumns = _makeSingleIndexDataFrame() 

685 

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

687 

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

689 

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

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

692 

693 # Check reading the columns. 

694 columns = list(tab2.dtype.names) 

695 columns2 = self.butler.get( 

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