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)