@@ -248,13 +248,8 @@ def test_build_idx_grib_mapping(self):
248
248
f"{ mapping_fname } .idx_grib_mapping.parquet" ,
249
249
)
250
250
# Build the mapping from idx to cfgrib metadata and assert it matches the fixture
251
- expected = pd .read_parquet (test_path )
252
- pd .testing .assert_frame_equal (
253
- mapping ,
254
- expected .assign (
255
- step = lambda x : x .step .astype ("timedelta64[ns]" )
256
- ),
257
- )
251
+ expected = pd .read_parquet (test_path , engine = "fastparquet" )
252
+ pd .testing .assert_frame_equal (mapping , expected )
258
253
259
254
# parse the idx files for 20231104 and compare the mapped result to the direct indexed result
260
255
test_name = fnames ["20231104" ]
@@ -289,11 +284,7 @@ def test_build_idx_grib_mapping(self):
289
284
"20231104" ,
290
285
f"{ test_name } .kindex.parquet" ,
291
286
)
292
- expected = pd .read_parquet (kindex_test_path )
293
-
294
- expected = expected .assign (
295
- step = lambda x : x .step .astype ("timedelta64[ns]" )
296
- )
287
+ expected = pd .read_parquet (kindex_test_path , engine = "fastparquet" )
297
288
298
289
expected = expected .set_index (
299
290
["varname" , "typeOfLevel" , "stepType" , "step" , "level" ]
@@ -393,13 +384,8 @@ def test_kerchunk_indexing(self):
393
384
TEST_DATE ,
394
385
f"{ fname } .kindex.parquet" ,
395
386
)
396
- expected = pd .read_parquet (test_path )
397
- pd .testing .assert_frame_equal (
398
- kindex ,
399
- expected .assign (
400
- step = lambda x : x .step .astype ("timedelta64[ns]" )
401
- ),
402
- )
387
+ expected = pd .read_parquet (test_path , engine = "fastparquet" )
388
+ pd .testing .assert_frame_equal (kindex , expected )
403
389
404
390
@unittest .skip ("TODO" )
405
391
def test_extract_dataset_chunk_index (self ):
@@ -444,10 +430,8 @@ def _test_index_extraction(self, sample_prefix: str):
444
430
test_path = os .path .join (
445
431
THIS_DIR , "grib_idx_fixtures" , sample_prefix , "kerchunk_index.parquet"
446
432
)
447
- expected = pd .read_parquet (test_path )
448
- pd .testing .assert_frame_equal (
449
- k_index , expected .assign (step = lambda x : x .step .astype ("timedelta64[ns]" ))
450
- )
433
+ expected = pd .read_parquet (test_path , engine = "fastparquet" )
434
+ pd .testing .assert_frame_equal (k_index , expected )
451
435
452
436
def test_strip_datavar_chunks (self ):
453
437
for sample_prefix , pre , post in [
@@ -583,7 +567,8 @@ def _reinflate_grib_store(
583
567
kind = pd .read_parquet (
584
568
os .path .join (
585
569
THIS_DIR , "grib_idx_fixtures" , dataset , "test_reinflate.parquet"
586
- )
570
+ ),
571
+ engine = "fastparquet" ,
587
572
)
588
573
589
574
zstore = reinflate_grib_store (
@@ -666,9 +651,6 @@ def _reinflate_grib_store(
666
651
with fsspec .open (test_path , "r" ) as f :
667
652
expected_keys = ujson .loads (f .read ())
668
653
669
- if key_set != expected_keys :
670
- print ("hello" )
671
-
672
654
self .assertListEqual (key_set , expected_keys )
673
655
674
656
def test_reinflate_grib_store (self ):
0 commit comments