(enhancement): Enable missing unit tests and Redshift, Athena, LF load tests (#1736)

* (enhancement): Enable missing unit tests and Redshift, Athena, LF load tests

Author: jaidisido

CONTRIBUTING.md

@@ -245,113 +245,6 @@ You can choose from three different environments to test your fixes/changes, bas
 
 ``./test_infra/scripts/delete-stack.sh databases``
 
-## Ray Load Tests Environment 
-**DISCLAIMER**: Make sure you know what you are doing. These steps will charge some services on your AWS account and require a minimum security skill to keep your environment safe.
-
-* Pick up a Linux or MacOS.
-* Install Python 3.7, 3.8 or 3.9 with [poetry](https://github.com/python-poetry/poetry) for package management
-* Fork the AWS SDK for pandas repository and clone that into your development environment
-
-* Then run the command bellow to install all dependencies:
-
-``poetry install``
-
-* Go to the ``test_infra`` directory
-
-``cd test_infra``
-
-* Install CDK dependencies:
-
-``poetry install``
-
-* [OPTIONAL] Set AWS_DEFAULT_REGION to define the region the Ray Test environment will deploy into. You may want to choose a region which you don't currently use:
-
-``export AWS_DEFAULT_REGION=ap-northeast-1``
-
-* Go to the ``scripts`` directory
-
-``cd scripts``
-
-* Deploy the `ray` CDK stack.
-
-``./deploy-stack.sh ray``
-
-* Configure Ray Cluster 
-
-``vi ray-cluster-config.yaml`` 
-
-```
-# Update the following file to match your enviroment
-# The following is an example
-cluster_name: ray-cluster
-
-min_workers: 2
-max_workers: 2
-
-provider:
-    type: aws
-    region: us-east-1 # change region as required
-    availability_zone: us-east-1a,us-east-1b,us-east-1c # change azs as required
-    security_group:
-        GroupName: ray_client_security_group
-    cache_stopped_nodes: False
-
-available_node_types:
-  ray.head.default:
-    node_config:
-      InstanceType: r5n.2xlarge # change instance type as required
-      IamInstanceProfile:
-        Arn: arn:aws:iam::{UPDATE YOUR ACCOUNT ID HERE}:instance-profile/ray-cluster-instance-profile
-      ImageId: ami-0ea510fcb67686b48 # latest ray images -> https://github.com/amzn/amazon-ray#amazon-ray-images
-      SubnetId: {replace with subnet within above AZs}
-
-  ray.worker.default:
-      min_workers: 2
-      max_workers: 2
-      node_config:
-        InstanceType: r5n.2xlarge
-        IamInstanceProfile:
-          Arn: arn:aws:iam::{UPDATE YOUR ACCOUNT ID HERE}:instance-profile/ray-cluster-instance-profile
-        ImageId: ami-0ea510fcb67686b48 # latest ray images -> https://github.com/amzn/amazon-ray#amazon-ray-images
-        SubnetId: {replace with subnet within above AZs}
-
-setup_commands:
-- pip install "awswrangler[modin, ray]==3.0.0rc1"
-- pip install pytest
-
-```
-
-* Create Ray Cluster 
-``ray up -y ray-cluster-config.yaml``
-
-* Push Load Tests to Ray Cluster
-``ray rsync-up ray-cluster-config.yaml tests/load /home/ubuntu/``
-
-* Submit Pytest Run to Ray Cluster
-```
-echo '''
-import os
-
-import pytest
-
-args = "-v load/"
-
-if not os.getenv("AWS_DEFAULT_REGION"):
-    os.environ["AWS_DEFAULT_REGION"] = "us-east-1" # Set your region as necessary
-
-result = pytest.main(args.split(" "))
-
-print(f"result: {result}")
-''' > handler.py
-ray submit ray-cluster-config.yaml handler.py
-```
-
-* Teardown Cluster 
-``ray down -y ray-cluster-config.yaml``
-
-[More on launching Ray Clusters on AWS](https://docs.ray.io/en/master/cluster/vms/user-guides/launching-clusters/aws.html#)
-
-
 ## Recommended Visual Studio Code Recommended setting
 
 ```json

awswrangler/lakeformation/_utils.py

@@ -44,7 +44,7 @@ def _build_table_objects(
     partitions_values: Dict[str, List[str]],
     use_threads: Union[bool, int],
     boto3_session: Optional[boto3.Session],
-) -> List[Dict[str, Any]]:
+) -> List[List[Dict[str, Any]]]:
     table_objects: List[Dict[str, Any]] = []
     paths_desc: Dict[str, Dict[str, Any]] = describe_objects(
         path=paths, use_threads=use_threads, boto3_session=boto3_session
@@ -58,7 +58,7 @@ def _build_table_objects(
         if partitions_values:
             table_object["PartitionValues"] = partitions_values[f"{path.rsplit('/', 1)[0].rstrip('/')}/"]
         table_objects.append(table_object)
-    return table_objects
+    return _utils.chunkify(table_objects, max_length=100)  # LF write operations is limited to 100 objects per call
 
 
 def _get_table_objects(

awswrangler/redshift.py

@@ -1582,7 +1582,6 @@ def copy(  # pylint: disable=too-many-arguments,too-many-locals
         e.g. s3_additional_kwargs={'ServerSideEncryption': 'aws:kms', 'SSEKMSKeyId': 'YOUR_KMS_KEY_ARN'}
     max_rows_by_file : int
         Max number of rows in each file.
-        Default is None i.e. dont split the files.
         (e.g. 33554432, 268435456)
     precombine_key : str, optional
         When there is a primary_key match during upsert, this column will change the upsert method,

awswrangler/s3/_write_dataset.py

@@ -316,19 +316,22 @@ def _to_dataset(
     _logger.debug("paths: %s", paths)
     _logger.debug("partitions_values: %s", partitions_values)
     if (table_type == "GOVERNED") and (table is not None) and (database is not None):
-        add_objects: List[Dict[str, Any]] = lakeformation._build_table_objects(  # pylint: disable=protected-access
+        list_add_objects: List[
+            List[Dict[str, Any]]
+        ] = lakeformation._build_table_objects(  # pylint: disable=protected-access
             paths, partitions_values, use_threads=use_threads, boto3_session=boto3_session
         )
         try:
-            if add_objects:
-                lakeformation._update_table_objects(  # pylint: disable=protected-access
-                    catalog_id=catalog_id,
-                    database=database,
-                    table=table,
-                    transaction_id=transaction_id,  # type: ignore
-                    add_objects=add_objects,
-                    boto3_session=boto3_session,
-                )
+            if list_add_objects:
+                for add_objects in list_add_objects:
+                    lakeformation._update_table_objects(  # pylint: disable=protected-access
+                        catalog_id=catalog_id,
+                        database=database,
+                        table=table,
+                        transaction_id=transaction_id,  # type: ignore
+                        add_objects=add_objects,
+                        boto3_session=boto3_session,
+                    )
         except Exception as ex:
             _logger.error(ex)
             raise

tests/conftest.py

@@ -349,3 +349,10 @@ def random_glue_database():
     database_name = get_time_str_with_random_suffix()
     yield database_name
     wr.catalog.delete_database(database_name)
+
+
+@pytest.fixture(scope="function")
+def redshift_con():
+    con = wr.redshift.connect("aws-sdk-pandas-redshift")
+    yield con
+    con.close()