Improving spark-to-redshift stress test

Author: igorborgest

awswrangler/redshift.py

@@ -137,6 +137,7 @@ def write_load_manifest(self, manifest_path: str, objects_paths: List[str], proc
         client_s3 = self._session.boto3_session.client(service_name="s3", config=self._session.botocore_config)
         bucket: str
         bucket, path = manifest_path.replace("s3://", "").split("/", 1)
+        logger.info(f"payload: {payload}")
         client_s3.put_object(Body=payload, Bucket=bucket, Key=path)
         return manifest
 

testing/test_awswrangler/test_redshift.py

@@ -287,9 +287,10 @@ def test_stress_to_redshift_spark_big(session, bucket, redshift_parameters):
             "B": list(range(1_000_000)),
             "C": list(range(1_000_000))
         }))
+    dataframe.cache()
 
     for i in range(10):
-        print(i)
+        print(f"Run number: {i}")
         con = Redshift.generate_connection(
             database="test",
             host=redshift_parameters.get("RedshiftAddress"),
@@ -299,15 +300,16 @@ def test_stress_to_redshift_spark_big(session, bucket, redshift_parameters):
         )
         session.spark.to_redshift(
             dataframe=dataframe,
-            path=f"s3://{bucket}/redshift-load/",
+            path=f"s3://{bucket}/redshift-load-{i}/",
             connection=con,
             schema="public",
             table="test",
             iam_role=redshift_parameters.get("RedshiftRole"),
             mode="overwrite",
-            min_num_partitions=4,
+            min_num_partitions=16,
         )
         con.close()
+        dataframe.unpersist()
 
 
 @pytest.mark.parametrize(