Skip to content

Commit

Permalink
Add tests for uint8, uint16, uint32 and uint64. #76
Browse files Browse the repository at this point in the history
  • Loading branch information
igorborgest committed May 8, 2020
1 parent 1a37722 commit 7d190b3
Showing 1 changed file with 35 additions and 0 deletions.
35 changes: 35 additions & 0 deletions testing/test_awswrangler/test_data_lake.py
Expand Up @@ -1396,3 +1396,38 @@ def test_unsigned_parquet(bucket, database):
df["c0"] = df.c0.astype("uint64")
with pytest.raises(wr.exceptions.UnsupportedType):
wr.s3.to_parquet(df=df, path=path, dataset=True, database=database, table=table, mode="overwrite")

wr.s3.delete_objects(path=path)
wr.catalog.delete_table_if_exists(database=database, table=table)


def test_parquet_uint64(bucket):
path = f"s3://{bucket}/test_parquet_uint64/"
wr.s3.delete_objects(path=path)
df = pd.DataFrame(
{
"c0": [0, 0, (2 ** 8) - 1],
"c1": [0, 0, (2 ** 16) - 1],
"c2": [0, 0, (2 ** 32) - 1],
"c3": [0, 0, (2 ** 64) - 1],
"c4": [0, 1, 2],
}
)
print(df)
df["c0"] = df.c0.astype("uint8")
df["c1"] = df.c1.astype("uint16")
df["c2"] = df.c2.astype("uint32")
df["c3"] = df.c3.astype("uint64")
paths = wr.s3.to_parquet(df=df, path=path, dataset=True, mode="overwrite", partition_cols=["c4"])["paths"]
wr.s3.wait_objects_exist(paths=paths, use_threads=False)
df = wr.s3.read_parquet(path=path, dataset=True)
print(df)
print(df.dtypes)
assert len(df.index) == 3
assert len(df.columns) == 5
assert df.c0.max() == (2 ** 8) - 1
assert df.c1.max() == (2 ** 16) - 1
assert df.c2.max() == (2 ** 32) - 1
assert df.c3.max() == (2 ** 64) - 1
assert df.c4.astype("uint8").sum() == 3
wr.s3.delete_objects(path=path)

0 comments on commit 7d190b3

Please sign in to comment.