def test_date(in_type, pd_date_type, out_type): test_data_path = "tests/data/date_type.csv" test_str_dates = pd.read_csv(test_data_path, dtype=str)["my_date"] test_str_dates = [None if pd.isna(s) else s for s in test_str_dates] schema = pa.schema([("my_date", getattr(pa, in_type)())]) # datetime_object if in_type == "date32" and pd_date_type == "pd_period": with pytest.warns(UserWarning): df = pa_read_csv_to_pandas( test_data_path, schema, expect_full_schema=False, pd_date_type=pd_date_type, ) else: df = pa_read_csv_to_pandas(test_data_path, schema, expect_full_schema=False, pd_date_type=pd_date_type) test_str_dates = pd.read_csv(test_data_path, dtype=str)["my_date"] test_str_dates = [None if pd.isna(s) else s for s in test_str_dates] assert str(df.my_date.dtype) == out_type if out_type == "object": assert isinstance(df.my_date[0], datetime.date) actual_str_dates = pd_datetime_series_to_list(df.my_date, out_type.split("[")[0], date=True) assert test_str_dates == actual_str_dates
def test_string(in_type, pd_old_type, pd_new_type): schema = pa.schema([("string_col", pa.string())]) df_old = pa_read_csv_to_pandas("tests/data/string_type.csv", schema, False, pd_string=False) assert str(df_old.my_string.dtype) == pd_old_type df_new = pa_read_csv_to_pandas("tests/data/string_type.csv", schema, False, pd_string=True) assert str(df_new.my_string.dtype) == pd_new_type
def test_datetime(in_type, pd_timestamp_type, out_type): test_data_path = "tests/data/datetime_type.csv" test_str_dates = pd.read_csv(test_data_path, dtype=str)["my_datetime"] test_str_dates = [None if pd.isna(s) else s for s in test_str_dates] type_dict = { "timestamp[s]": pa.timestamp("s"), "timestamp[ms]": pa.timestamp("ms"), "timestamp[us]": pa.timestamp("us"), "timestamp[ns]": pa.timestamp("ns"), } schema = pa.schema([("my_datetime", type_dict[in_type])]) # datetime_object df = pa_read_csv_to_pandas( test_data_path, schema=schema, expect_full_schema=False, pd_timestamp_type=pd_timestamp_type, ) test_str_dates = pd.read_csv(test_data_path, dtype=str)["my_datetime"] test_str_dates = [None if pd.isna(s) else s for s in test_str_dates] assert str(df.my_datetime.dtype) == out_type if out_type == "object": assert isinstance(df.my_datetime[0], datetime.datetime) actual_str_dates = pd_datetime_series_to_list(df.my_datetime, out_type.split("[")[0], date=False) assert test_str_dates == actual_str_dates
def test_to_parquet(schema, boolean_args, date_args): original = pa_read_csv_to_pandas( "tests/data/all_types.csv", schema, pd_boolean=boolean_args, pd_integer=boolean_args, pd_string=boolean_args, pd_date_type=date_args, pd_timestamp_type=date_args, ) # output as parquet with tempfile.NamedTemporaryFile(suffix=".parquet") as f: tmp_out_file = f.name pd_to_parquet(original, tmp_out_file) # read in as parquet reloaded = pa_read_parquet_to_pandas( tmp_out_file, schema, pd_boolean=boolean_args, pd_integer=boolean_args, pd_string=boolean_args, pd_date_type=date_args, pd_timestamp_type=date_args, ) assert_frame_equal(original, reloaded)
def test_int_csv(in_type, pd_old_type, pd_new_type): """Testing csv mapping from pyarrow to Pandas data types. Args: in_type (str): pyarrow data type read in from the csv. pd_old_type (str): old pandas data type mapping. pd_new_type (str): new pandas data type mapping. """ schema = pa.schema([("int_col", getattr(pa, in_type)())]) test_file = "tests/data/int_type.csv" df_old = pa_read_csv_to_pandas(test_file, schema, False, pd_integer=False) assert str(df_old.my_int.dtype) == pd_old_type df_new = pa_read_csv_to_pandas(test_file, schema, False, pd_integer=True) assert str(df_new.my_int.dtype) == pd_new_type
def test_out_of_bounds_datetime(pd_timestamp_type, expect_error): test_data_path = "tests/data/datetime_type_oob.csv" test_str_dates = pd.read_csv(test_data_path, dtype=str)["my_datetime"] test_str_dates = [None if pd.isna(s) else s for s in test_str_dates] schema = pa.schema([("my_datetime", pa.timestamp("s"))]) out_type_lu = { "datetime_object": "object", "pd_timestamp": "datetime64", "pd_period": "period", } out_type = out_type_lu[pd_timestamp_type] # datetime_object try: df = pa_read_csv_to_pandas( test_data_path, schema=schema, expect_full_schema=False, pd_timestamp_type=pd_timestamp_type, ) except pa.lib.ArrowInvalid: assert expect_error is True else: df = pa_read_csv_to_pandas( test_data_path, schema=schema, expect_full_schema=False, pd_timestamp_type=pd_timestamp_type, ) test_str_dates = pd.read_csv(test_data_path, dtype=str)["my_datetime"] test_str_dates = [None if pd.isna(s) else s for s in test_str_dates] if out_type == "object": assert isinstance(df.my_datetime[0], datetime.datetime) actual_str_dates = pd_datetime_series_to_list(df.my_datetime, out_type.split("[")[0], date=False) assert test_str_dates == actual_str_dates
def test_bool_csv_and_json(): schema = pa.schema([("i", pa.int8()), ("my_bool", pa.bool_()), ("my_nullable_bool", pa.bool_())]) df_csv = pa_read_csv_to_pandas("tests/data/bool_type.csv", schema, pd_boolean=True) df_jsonl = pa_read_json_to_pandas("tests/data/bool_type.jsonl", schema, pd_boolean=True) assert df_csv.equals(df_jsonl)
def test_file_reader_works_with_schema(): csv_schema = pa.schema([("test", pa.string()), ("a_column", pa.string())]) df_csv = pa_read_csv_to_pandas("tests/data/example_data.csv") df_csv_schema = pa_read_csv_to_pandas("tests/data/example_data.csv", csv_schema) assert_frame_equal(df_csv, df_csv_schema) json_schema = pa.schema([("a", pa.int64()), ("b", pa.float64()), ("c", pa.string()), ("d", pa.bool_())]) df_json = pa_read_json_to_pandas("tests/data/example_data.jsonl") df_json_schema = pa_read_json_to_pandas("tests/data/example_data.jsonl", json_schema) assert_frame_equal(df_json, df_json_schema) # Check raises error on both readers missing_schema = pa.schema([("b", pa.float64()), ("c", pa.string()), ("d", pa.bool_())]) with pytest.raises(ValueError): pa_read_json_to_pandas("tests/data/example_data.jsonl", missing_schema) with pytest.raises(ValueError): pa_read_csv_to_pandas("tests/data/example_data.csv", missing_schema)
def test_timestamps_as_strs(): test_data_path = "tests/data/datetime_type.csv" test_str_dates = pd.read_csv(test_data_path, dtype="string")["my_datetime"] schema = pa.schema([("my_datetime", pa.string())]) df = pa_read_csv_to_pandas(test_data_path, schema, expect_full_schema=False) assert df["my_datetime"].to_list() == test_str_dates.to_list() df = pa_read_json_to_pandas(test_data_path.replace(".csv", ".jsonl"), schema, expect_full_schema=False) assert df["my_datetime"].to_list() == test_str_dates.to_list()
def test_pd_to_csv(boolean_args, date_args, schema): original = pa_read_csv_to_pandas( "tests/data/all_types.csv", schema, pd_boolean=boolean_args, pd_integer=boolean_args, pd_string=boolean_args, pd_date_type=date_args, pd_timestamp_type=date_args, ) # Write to StringIO then convert to BytesIO so Arrow can read it output = io.StringIO() pd_to_csv(original, output) as_bytes = io.BytesIO(bytearray(output.getvalue(), "utf-8")) reloaded = pa_read_csv_to_pandas( as_bytes, schema, pd_boolean=boolean_args, pd_integer=boolean_args, pd_string=boolean_args, pd_date_type=date_args, pd_timestamp_type=date_args, ) assert_frame_equal(original, reloaded)
def test_decimal_float(arrow_type, pd_type): type_lu = { "float32": pa.float32(), "float64": pa.float64(), "decimal": pa.decimal128(5, 3), } schema = pa.schema([("i", pa.int8()), ("my_decimal", type_lu[arrow_type])]) df_csv = pa_read_csv_to_pandas("tests/data/decimal_type.csv", schema) df_json = pa_read_json_to_pandas("tests/data/decimal_type.jsonl", schema) assert str(df_csv.my_decimal.dtype) == pd_type assert str(df_json.my_decimal.dtype) == pd_type assert_frame_equal(df_csv, df_json)
def test_csv_options(in_type, pd_old_type, pd_new_type): schema = pa.schema([("string_col", pa.string())]) read_options = csv.ReadOptions(skip_rows=1) parse_options = csv.ParseOptions(quote_char="'", escape_char="\\", delimiter=";", newlines_in_values=True) convert_options = csv.ConvertOptions( include_columns=["i", "my_string", "nonexistent_column"], include_missing_columns=True, null_values=["NULL_STRING"], strings_can_be_null=True, ) df = pa_read_csv_to_pandas( "tests/data/csv_options_test.csv", schema, False, pd_string=False, parse_options=parse_options, convert_options=convert_options, read_options=read_options, ) expected = [ "dsfasd;dsffadsf", "dsfasd;dsffadsf", None, "this text\nhas a line break", "this text, like so, has commas", ] assert df.columns.tolist() == ["i", "my_string", "nonexistent_column"] assert df["nonexistent_column"].isnull().all() assert_series_equal(df["my_string"], Series(expected, name="my_string"))
def test_file_reader_returns_df(): df = pa_read_csv_to_pandas("tests/data/example_data.csv") assert isinstance(df, pd.DataFrame) df = pa_read_json_to_pandas("tests/data/example_data.jsonl") assert isinstance(df, pd.DataFrame)