def test_to_arrow_empty_dataframe(self):
fd, filename = tempfile.mkstemp()
os.close(fd)
# Remove the file. Then we'll test that ProcessResult.to_arrow() does
# not write it (because the result is an error)
os.unlink(filename)
try:
result = ProcessResult.coerce("bad, bad error").to_arrow(Path(filename))
self.assertEqual(
result,
atypes.RenderResult(
atypes.ArrowTable(None, atypes.TableMetadata(0, [])),
[
atypes.RenderError(
atypes.I18nMessage.TODO_i18n("bad, bad error"), []
)
],
{},
),
)
with self.assertRaises(FileNotFoundError):
open(filename)
finally:
try:
os.unlink(filename)
except FileNotFoundError:
pass
def test_to_arrow_normal_dataframe(self):
fd, filename = tempfile.mkstemp()
os.close(fd)
# Remove the file. Then we'll test that ProcessResult.to_arrow() does
# not write it (because the result is an error)
os.unlink(filename)
try:
process_result = ProcessResult.coerce(pd.DataFrame({"A": [1, 2]}))
result = process_result.to_arrow(Path(filename))
self.assertEqual(
result,
atypes.RenderResult(
atypes.ArrowTable(
Path(filename),
pyarrow.table({"A": [1, 2]}),
atypes.TableMetadata(
2,
[
atypes.Column(
"A",
ColumnType.Number(
# Whatever .format
# ProcessResult.coerce() gave
process_result.columns[0].type.format),
)
],
),
),
[],
{},
),
)
finally:
os.unlink(filename)
def test_to_arrow(self):
self.assertEqual(
TableShape(
3,
[
Column("A", ColumnType.NUMBER("{:,d}")),
Column("B", ColumnType.TEXT()),
],
).to_arrow(),
atypes.TableMetadata(
3,
[
atypes.Column("A", atypes.ColumnType.Number("{:,d}")),
atypes.Column("B", atypes.ColumnType.Text()),
],
),
)
def test_table_metadata_to_thrift(self):
self.assertEqual(
types.TableMetadata(
4,
[
types.Column("A", types.ColumnType.Text()),
types.Column("B", types.ColumnType.Text()),
],
).to_thrift(),
ttypes.TableMetadata(
4,
[
ttypes.Column(
"A", ttypes.ColumnType(text_type=ttypes.ColumnTypeText())
),
ttypes.Column(
"B", ttypes.ColumnType(text_type=ttypes.ColumnTypeText())
),
],
),
)