""" Tests that work on both the Python and C engines but do not have a specific classification into the other test modules. """ from io import StringIO import numpy as np import pytest from pandas._libs import parsers as libparsers from pandas.errors import DtypeWarning from pandas import ( DataFrame, concat, ) import pandas._testing as tm pytestmark = pytest.mark.usefixtures("pyarrow_skip") @pytest.mark.parametrize("index_col", [0, "index"]) def test_read_chunksize_with_index(all_parsers, index_col): parser = all_parsers data = """index,A,B,C,D foo,2,3,4,5 bar,7,8,9,10 baz,12,13,14,15 qux,12,13,14,15 foo2,12,13,14,15 bar2,12,13,14,15 """ expected = DataFrame( [ ["foo", 2, 3, 4, 5], ["bar", 7, 8, 9, 10], ["baz", 12, 13, 14, 15], ["qux", 12, 13, 14, 15], ["foo2", 12, 13, 14, 15], ["bar2", 12, 13, 14, 15], ], columns=["index", "A", "B", "C", "D"], ) expected = expected.set_index("index") with parser.read_csv(StringIO(data), index_col=0, chunksize=2) as reader: chunks = list(reader) tm.assert_frame_equal(chunks[0], expected[:2]) tm.assert_frame_equal(chunks[1], expected[2:4]) tm.assert_frame_equal(chunks[2], expected[4:]) @pytest.mark.parametrize("chunksize", [1.3, "foo", 0]) def test_read_chunksize_bad(all_parsers, chunksize): data = """index,A,B,C,D foo,2,3,4,5 bar,7,8,9,10 baz,12,13,14,15 qux,12,13,14,15 foo2,12,13,14,15 bar2,12,13,14,15 """ parser = all_parsers msg = r"'chunksize' must be an integer >=1" with pytest.raises(ValueError, match=msg): with parser.read_csv(StringIO(data), chunksize=chunksize) as _: pass @pytest.mark.parametrize("chunksize", [2, 8]) def test_read_chunksize_and_nrows(all_parsers, chunksize): # see gh-15755 data = """index,A,B,C,D foo,2,3,4,5 bar,7,8,9,10 baz,12,13,14,15 qux,12,13,14,15 foo2,12,13,14,15 bar2,12,13,14,15 """ parser = all_parsers kwargs = {"index_col": 0, "nrows": 5} expected = parser.read_csv(StringIO(data), **kwargs) with parser.read_csv(StringIO(data), chunksize=chunksize, **kwargs) as reader: tm.assert_frame_equal(concat(reader), expected) def test_read_chunksize_and_nrows_changing_size(all_parsers): data = """index,A,B,C,D foo,2,3,4,5 bar,7,8,9,10 baz,12,13,14,15 qux,12,13,14,15 foo2,12,13,14,15 bar2,12,13,14,15 """ parser = all_parsers kwargs = {"index_col": 0, "nrows": 5} expected = parser.read_csv(StringIO(data), **kwargs) with parser.read_csv(StringIO(data), chunksize=8, **kwargs) as reader: tm.assert_frame_equal(reader.get_chunk(size=2), expected.iloc[:2]) tm.assert_frame_equal(reader.get_chunk(size=4), expected.iloc[2:5]) with pytest.raises(StopIteration, match=""): reader.get_chunk(size=3) def test_get_chunk_passed_chunksize(all_parsers): parser = all_parsers data = """A,B,C 1,2,3 4,5,6 7,8,9 1,2,3""" with parser.read_csv(StringIO(data), chunksize=2) as reader: result = reader.get_chunk() expected = DataFrame([[1, 2, 3], [4, 5, 6]], columns=["A", "B", "C"]) tm.assert_frame_equal(result, expected) @pytest.mark.parametrize("kwargs", [{}, {"index_col": 0}]) def test_read_chunksize_compat(all_parsers, kwargs): # see gh-12185 data = """index,A,B,C,D foo,2,3,4,5 bar,7,8,9,10 baz,12,13,14,15 qux,12,13,14,15 foo2,12,13,14,15 bar2,12,13,14,15 """ parser = all_parsers result = parser.read_csv(StringIO(data), **kwargs) with parser.read_csv(StringIO(data), chunksize=2, **kwargs) as reader: tm.assert_frame_equal(concat(reader), result) def test_read_chunksize_jagged_names(all_parsers): # see gh-23509 parser = all_parsers data = "\n".join(["0"] * 7 + [",".join(["0"] * 10)]) expected = DataFrame([[0] + [np.nan] * 9] * 7 + [[0] * 10]) with parser.read_csv(StringIO(data), names=range(10), chunksize=4) as reader: result = concat(reader) tm.assert_frame_equal(result, expected) def test_chunk_begins_with_newline_whitespace(all_parsers): # see gh-10022 parser = all_parsers data = "\n hello\nworld\n" result = parser.read_csv(StringIO(data), header=None) expected = DataFrame([" hello", "world"]) tm.assert_frame_equal(result, expected) @pytest.mark.slow def test_chunks_have_consistent_numerical_type(all_parsers, monkeypatch): # mainly an issue with the C parser heuristic = 2**3 parser = all_parsers integers = [str(i) for i in range(heuristic - 1)] data = "a\n" + "\n".join(integers + ["1.0", "2.0"] + integers) # Coercions should work without warnings. with tm.assert_produces_warning(None): with monkeypatch.context() as m: m.setattr(libparsers, "DEFAULT_BUFFER_HEURISTIC", heuristic) result = parser.read_csv(StringIO(data)) assert type(result.a[0]) is np.float64 assert result.a.dtype == float def test_warn_if_chunks_have_mismatched_type(all_parsers): warning_type = None parser = all_parsers size = 10000 # see gh-3866: if chunks are different types and can't # be coerced using numerical types, then issue warning. if parser.engine == "c" and parser.low_memory: warning_type = DtypeWarning # Use larger size to hit warning path size = 499999 integers = [str(i) for i in range(size)] data = "a\n" + "\n".join(integers + ["a", "b"] + integers) buf = StringIO(data) df = parser.read_csv_check_warnings( warning_type, r"Columns \(0\) have mixed types. " "Specify dtype option on import or set low_memory=False.", buf, ) assert df.a.dtype == object @pytest.mark.parametrize("iterator", [True, False]) def test_empty_with_nrows_chunksize(all_parsers, iterator): # see gh-9535 parser = all_parsers expected = DataFrame(columns=["foo", "bar"]) nrows = 10 data = StringIO("foo,bar\n") if iterator: with parser.read_csv(data, chunksize=nrows) as reader: result = next(iter(reader)) else: result = parser.read_csv(data, nrows=nrows) tm.assert_frame_equal(result, expected) def test_read_csv_memory_growth_chunksize(all_parsers): # see gh-24805 # # Let's just make sure that we don't crash # as we iteratively process all chunks. parser = all_parsers with tm.ensure_clean() as path: with open(path, "w", encoding="utf-8") as f: for i in range(1000): f.write(str(i) + "\n") with parser.read_csv(path, chunksize=20) as result: for _ in result: pass def test_chunksize_with_usecols_second_block_shorter(all_parsers): # GH#21211 parser = all_parsers data = """1,2,3,4 5,6,7,8 9,10,11 """ result_chunks = parser.read_csv( StringIO(data), names=["a", "b"], chunksize=2, usecols=[0, 1], header=None, ) expected_frames = [ DataFrame({"a": [1, 5], "b": [2, 6]}), DataFrame({"a": [9], "b": [10]}, index=[2]), ] for i, result in enumerate(result_chunks): tm.assert_frame_equal(result, expected_frames[i]) def test_chunksize_second_block_shorter(all_parsers): # GH#21211 parser = all_parsers data = """a,b,c,d 1,2,3,4 5,6,7,8 9,10,11 """ result_chunks = parser.read_csv(StringIO(data), chunksize=2) expected_frames = [ DataFrame({"a": [1, 5], "b": [2, 6], "c": [3, 7], "d": [4, 8]}), DataFrame({"a": [9], "b": [10], "c": [11], "d": [np.nan]}, index=[2]), ] for i, result in enumerate(result_chunks): tm.assert_frame_equal(result, expected_frames[i])