Merge pull request #63 from awslabs/fix-single-row

Fixing single-row-by-partition issue

Author: igorborgest

README.md

@@ -2,7 +2,7 @@
 
 > Utility belt to handle data on AWS.
 
-[![Release](https://img.shields.io/badge/release-0.0.18-brightgreen.svg)](https://pypi.org/project/awswrangler/)
+[![Release](https://img.shields.io/badge/release-0.0.19-brightgreen.svg)](https://pypi.org/project/awswrangler/)
 [![Downloads](https://img.shields.io/pypi/dm/awswrangler.svg)](https://pypi.org/project/awswrangler/)
 [![Python Version](https://img.shields.io/badge/python-3.6%20%7C%203.7-brightgreen.svg)](https://pypi.org/project/awswrangler/)
 [![Documentation Status](https://readthedocs.org/projects/aws-data-wrangler/badge/?version=latest)](https://aws-data-wrangler.readthedocs.io/en/latest/?badge=latest)

awswrangler/__version__.py

@@ -1,4 +1,4 @@
 __title__ = "awswrangler"
 __description__ = "Utility belt to handle data on AWS."
-__version__ = "0.0.18"
+__version__ = "0.0.19"
 __license__ = "Apache License 2.0"

awswrangler/pandas.py

@@ -893,7 +893,7 @@ def write_parquet_dataframe(dataframe, path, preserve_index, compression, fs, ca
         dtypes = copy.deepcopy(dataframe.dtypes.to_dict())
         for name, dtype in dtypes.items():
             if str(dtype) == "Int64":
-                dataframe.loc[:, name] = dataframe[name].astype("float64")
+                dataframe[name] = dataframe[name].astype("float64")
                 casted_in_pandas.append(name)
                 cast_columns[name] = "bigint"
                 logger.debug(f"Casting column {name} Int64 to float64")

testing/test_awswrangler/test_pandas.py

@@ -1099,3 +1099,53 @@ def test_partition_cast(session, bucket, database):
     assert str(df2.dtypes[3]).startswith("bool")
     assert str(df2.dtypes[4]).startswith("datetime")
     session.s3.delete_objects(path=path)
+
+
+@pytest.mark.parametrize("procs", [1, 2, 8])
+def test_partition_single_row(session, bucket, database, procs):
+    data = {
+        "col1": [
+            1,
+            2,
+            3,
+        ],
+        "datecol": [
+            "2019-11-09",
+            "2019-11-09",
+            "2019-11-08",
+        ],
+        "partcol": [
+            "2019-11-09",
+            "2019-11-09",
+            "2019-11-08",
+        ]
+    }
+    df = pd.DataFrame(data)
+    df = df.astype({"datecol": "datetime64", "partcol": "datetime64"})
+    schema = {
+        "col1": "bigint",
+        "datecol": "date",
+        "partcol": "date",
+    }
+    path = f"s3://{bucket}/test/"
+    session.pandas.to_parquet(dataframe=df,
+                              database=database,
+                              path=path,
+                              partition_cols=["datecol"],
+                              mode="overwrite",
+                              cast_columns=schema,
+                              procs_cpu_bound=procs,
+                              preserve_index=False)
+    df2 = None
+    for counter in range(10):
+        df2 = session.pandas.read_sql_athena(sql="select * from test", database=database)
+        assert len(list(df.columns)) == len(list(df2.columns))
+        if len(df.index) == len(df2.index):
+            break
+        sleep(1)
+    print(df2.dtypes)
+    assert len(df.index) == len(df2.index)
+    assert df2.dtypes[0] == "Int64"
+    assert df2.dtypes[1] == "object"
+    assert df2.dtypes[2] == "object"
+    session.s3.delete_objects(path=path)