Merge pull request #116 from awslabs/range-index

Removing regular indexes from the compulsory Int64 cast

Author: igorborgest

awswrangler/pandas.py

@@ -1519,9 +1519,10 @@ def _read_parquet_path(session_primitives: "SessionPrimitives",
             fs.invalidate_cache()
             table = pq.read_table(source=path, columns=columns, filters=filters, filesystem=fs, use_threads=use_threads)
         # Check if we lose some integer during the conversion (Happens when has some null value)
-        integers = [field.name for field in table.schema if str(field.type).startswith("int")]
+        integers = [field.name for field in table.schema if str(field.type).startswith("int") and field.name != "__index_level_0__"]
         logger.debug(f"Converting to Pandas: {path}")
         df = table.to_pandas(use_threads=use_threads, integer_object_nulls=True)
+        logger.debug(f"Casting Int64 columns: {path}")
         for c in integers:
             if not str(df[c].dtype).startswith("int"):
                 df[c] = df[c].astype("Int64")

testing/test_awswrangler/test_pandas.py

@@ -2186,3 +2186,16 @@ def test_to_parquet_categorical_partitions(bucket):
     wr.pandas.to_parquet(x[x.Year == 1990], path=path, partition_cols=["Year"])
     y = wr.pandas.read_parquet(path=path)
     assert len(x[x.Year == 1990].index) == len(y.index)
+
+
+def test_range_index(bucket, database):
+    path = f"s3://{bucket}/test_range_index"
+    wr.s3.delete_objects(path=path)
+    d = pd.date_range('1990-01-01', freq='D', periods=10000)
+    vals = pd.np.random.randn(len(d), 4)
+    x = pd.DataFrame(vals, index=d, columns=['A', 'B', 'C', 'D']).reset_index()
+    print(x)
+    wr.pandas.to_parquet(dataframe=x, path=path, database=database)
+    df = wr.pandas.read_parquet(path=path)
+    assert len(x.columns) == len(df.columns)
+    assert len(x.index) == len(df.index)