a Gb@sdZddlmZddlmZddlmZddlZddl Z ddl Z ddl m Z ddl mZe jdZdd Ze jjd d Zd d Ze jjddZdS)zd Tests multithreading behaviour for reading and parsing files for each parser defined in parsers.py ) ExitStack)BytesIO) ThreadPoolN) DataFrameZ pyarrow_skipcCs\ttj|dtdd}d|d<d|d<d|d<tjd|dd |d <tj|d d |d <|S)z Construct a DataFrame for testing. Parameters ---------- num_rows : int The number of rows for our DataFrame. Returns ------- df : DataFrame Zabcde)columnsZfooZbarZbazz20000101 09:00:00s)ZperiodsZfreqdateZint64)Zdtypeint)rnpZrandomZrandlistpdZ date_rangeZarange)num_rowsdfrGlib/python3.9/site-packages/pandas/tests/io/parser/test_multi_thread.py_construct_dataframes rc s|}dd}fddt|D}t\fdd|D}td}||j|}|d}|D]}t||qdWdn1s0YdS)Ni'dcs(g|] }dddtDqS) cSs&g|]}|dd|dd|dqS)d,r.0irrr 3zCtest_multi_thread_string_io_read_csv...)joinrangeencode)r_) max_row_rangerrr2sz8test_multi_thread_string_io_read_csv..csg|]}t|qSr) enter_contextr)rb)stackrrr9rr)rrr!rmapread_csvtmassert_frame_equal) all_parsersparserZ num_filesZ bytes_to_dffilespoolresultsZ first_resultresultr)r r#r$test_multi_thread_string_io_read_csv+s r/c sfdd}fddtD}td}|||}Wdn1sR0Y|dj}|ddD] } || _qrt|} | S) a Generate a DataFrame via multi-thread. Parameters ---------- parser : BaseParser The parser object to use for reading the data. path : str The location of the CSV file to read. num_rows : int The number of rows to read per task. num_tasks : int The number of tasks to use for reading this DataFrame. Returns ------- df : DataFrame csB|\}}|s"jdd|dgdSjddt|d|dgdS)aj Create a reader for part of the CSV. Parameters ---------- arg : tuple A tuple of the following: * start : int The starting row to start for parsing CSV * nrows : int The number of rows to read. Returns ------- df : DataFrame rr ) index_colheadernrows parse_datesN )r0r1Zskiprowsr2r3)r&r )argstartr2)r*pathrrreaderXs  z0_generate_multi_thread_dataframe..readercs g|]}|fqSrrr)r num_tasksrrrzsz4_generate_multi_thread_dataframe..)Z processesNrr4)rrr%rr concat) r*r8rr:r9Ztasksr,r-r1rfinal_dataframer)rr:r*r8r _generate_multi_thread_dataframeDs"  *  r>cCsjd}d}|}d}t|}t|4}||t||||}t||Wdn1s\0YdS)Niz__thread_pool_reader__.csv)rr'Z ensure_cleanZto_csvr>r()r)r:rr* file_namerr8r=rrr)test_multi_thread_path_multipart_read_csvs  rA)__doc__ contextlibriorZmultiprocessing.poolrZnumpyr ZpytestZpandasr rZpandas._testingZ_testingr'ZmarkZ usefixturesZ pytestmarkrZslowr/r>rArrrrs       F