Skip to content

Commit dba847e

Browse files
committed
Added a test for non-default indexes
1 parent 0710646 commit dba847e

File tree

1 file changed

+83
-27
lines changed

1 file changed

+83
-27
lines changed

pandas/tests/io/test_parquet.py

Lines changed: 83 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -247,33 +247,6 @@ def test_columns_dtypes_invalid(self, engine):
247247
datetime.datetime(2011, 1, 1, 1, 1)]
248248
self.check_error_on_write(df, engine, ValueError)
249249

250-
def test_write_with_index(self, engine):
251-
252-
df = pd.DataFrame({'A': [1, 2, 3]})
253-
self.check_round_trip(df, engine, write_kwargs={'compression': None})
254-
255-
# non-default index
256-
for index in [[2, 3, 4],
257-
pd.date_range('20130101', periods=3),
258-
list('abc'),
259-
[1, 3, 4],
260-
pd.MultiIndex.from_tuples([('a', 1), ('a', 2),
261-
('b', 1)]),
262-
]:
263-
264-
df.index = index
265-
self.check_error_on_write(df, engine, ValueError)
266-
267-
# index with meta-data
268-
df.index = [0, 1, 2]
269-
df.index.name = 'foo'
270-
self.check_error_on_write(df, engine, ValueError)
271-
272-
# column multi-index
273-
df.index = [0, 1, 2]
274-
df.columns = pd.MultiIndex.from_tuples([('a', 1), ('a', 2), ('b', 1)]),
275-
self.check_error_on_write(df, engine, ValueError)
276-
277250
@pytest.mark.parametrize('compression', [None, 'gzip', 'snappy', 'brotli'])
278251
def test_compression(self, engine, compression):
279252

@@ -360,6 +333,61 @@ def test_categorical_unsupported(self, pa_lt_070):
360333
df = pd.DataFrame({'a': pd.Categorical(list('abc'))})
361334
self.check_error_on_write(df, pa, NotImplementedError)
362335

336+
def test_write_with_index(self, pa_lt_070):
337+
engine = pa_lt_070
338+
df = pd.DataFrame({'A': [1, 2, 3]})
339+
self.check_round_trip(df, engine, write_kwargs={'compression': None})
340+
341+
# non-default index
342+
for index in [[2, 3, 4],
343+
pd.date_range('20130101', periods=3),
344+
list('abc'),
345+
[1, 3, 4],
346+
pd.MultiIndex.from_tuples([('a', 1), ('a', 2),
347+
('b', 1)]),
348+
]:
349+
350+
df.index = index
351+
self.check_error_on_write(df, engine, ValueError)
352+
353+
# index with meta-data
354+
df.index = [0, 1, 2]
355+
df.index.name = 'foo'
356+
self.check_error_on_write(df, engine, ValueError)
357+
358+
# column multi-index
359+
df.index = [0, 1, 2]
360+
df.columns = pd.MultiIndex.from_tuples([('a', 1), ('a', 2), ('b', 1)]),
361+
self.check_error_on_write(df, engine, ValueError)
362+
363+
def test_write_with_non_default_index(self, pa_ge_070):
364+
engine = pa_ge_070
365+
366+
df = pd.DataFrame({'A': [1, 2, 3]})
367+
self.check_round_trip(df, engine, write_kwargs={'compression': None})
368+
369+
# non-default index
370+
for index in [[2, 3, 4],
371+
pd.date_range('20130101', periods=3),
372+
list('abc'),
373+
[1, 3, 4],
374+
pd.MultiIndex.from_tuples([('a', 1), ('a', 2),
375+
('b', 1)]),
376+
]:
377+
378+
df.index = index
379+
self.check_round_trip(df, engine)
380+
381+
# index with meta-data
382+
df.index = [0, 1, 2]
383+
df.index.name = 'foo'
384+
self.check_round_trip(df, engine)
385+
386+
# column multi-index
387+
df.index = [0, 1, 2]
388+
df.columns = pd.MultiIndex.from_tuples([('a', 1), ('a', 2), ('b', 1)]),
389+
self.check_error_on_write(df, engine, ValueError)
390+
363391

364392
class TestParquetFastParquet(Base):
365393

@@ -385,6 +413,34 @@ def test_basic(self, fp):
385413

386414
self.check_round_trip(df, fp, write_kwargs={'compression': None})
387415

416+
def test_write_with_index(self, fp):
417+
engine = fp
418+
419+
df = pd.DataFrame({'A': [1, 2, 3]})
420+
self.check_round_trip(df, engine, write_kwargs={'compression': None})
421+
422+
# non-default index
423+
for index in [[2, 3, 4],
424+
pd.date_range('20130101', periods=3),
425+
list('abc'),
426+
[1, 3, 4],
427+
pd.MultiIndex.from_tuples([('a', 1), ('a', 2),
428+
('b', 1)]),
429+
]:
430+
431+
df.index = index
432+
self.check_error_on_write(df, engine, ValueError)
433+
434+
# index with meta-data
435+
df.index = [0, 1, 2]
436+
df.index.name = 'foo'
437+
self.check_error_on_write(df, engine, ValueError)
438+
439+
# column multi-index
440+
df.index = [0, 1, 2]
441+
df.columns = pd.MultiIndex.from_tuples([('a', 1), ('a', 2), ('b', 1)]),
442+
self.check_error_on_write(df, engine, ValueError)
443+
388444
@pytest.mark.skip(reason="not supported")
389445
def test_duplicate_columns(self, fp):
390446

0 commit comments

Comments
 (0)