def _check_roundtrip(con, df):
writer = DataFrameWriter(con, df)
path = writer.write_temp_csv()
table = writer.delimited_table(path)
df2 = table.execute()
tm.assert_frame_equal(df2, df)
def test_writer_cleanup_context_manager(con, hdfs, alltypes_df):
with DataFrameWriter(con, alltypes_df) as writer:
path = writer.write_temp_csv()
assert hdfs.exists(path)
# the path shouldn't exist after the context manager exits
assert not hdfs.exists(path)
def write_dataframe(
self,
df: pd.DataFrame,
path: str,
format: Literal['csv'] = 'csv',
) -> Any:
"""Write a pandas DataFrame to indicated file path.
Parameters
----------
df
Pandas DataFrame
path
Absolute file path
format
File format
"""
writer = DataFrameWriter(self, df)
return writer.write_csv(path)
def test_writer_cleanup_deletes_hdfs_dir(con, hdfs, alltypes_df):
writer = DataFrameWriter(con, alltypes_df)
path = writer.write_temp_csv()
assert hdfs.exists(path)
writer.cleanup()
assert not hdfs.exists(path)
# noop
writer.cleanup()
assert not hdfs.exists(path)
def _setup_insert(self, obj):
if isinstance(obj, pd.DataFrame):
with DataFrameWriter(self, obj) as writer:
yield writer.delimited_table(writer.write_temp_csv())
else:
yield obj