Add columns parameters to Pandas.to_csv()

Author: igorborgest

awswrangler/data_types.py

@@ -370,8 +370,8 @@ def extract_pyarrow_schema_from_pandas(dataframe: pd.DataFrame,
     :param indexes_position: "right" or "left"
     :return: Pyarrow schema (e.g. [("col name": "bigint"), ("col2 name": "int")]
     """
-    cols = []
-    cols_dtypes = {}
+    cols: List[str] = []
+    cols_dtypes: Dict[str, str] = {}
     if indexes_position not in ("right", "left"):
         raise ValueError(f"indexes_position must be \"right\" or \"left\"")
 
@@ -384,10 +384,10 @@ def extract_pyarrow_schema_from_pandas(dataframe: pd.DataFrame,
             cols.append(name)
 
     # Filling cols_dtypes and indexes
-    indexes = []
+    indexes: List[str] = []
     for field in pa.Schema.from_pandas(df=dataframe[cols], preserve_index=preserve_index):
         name = str(field.name)
-        dtype = field.type
+        dtype = str(field.type)
         cols_dtypes[name] = dtype
         if name not in dataframe.columns:
             indexes.append(name)

awswrangler/glue.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Dict, Optional, Any, Iterator, List, Union
+from typing import TYPE_CHECKING, Dict, Optional, Any, Iterator, List, Union, Tuple
 from math import ceil
 from itertools import islice
 import re
@@ -55,16 +55,16 @@ def get_table_python_types(self, database: str, table: str) -> Dict[str, Optiona
     def metadata_to_glue(self,
                          dataframe,
                          path: str,
-                         objects_paths,
-                         file_format,
-                         database=None,
-                         table=None,
-                         partition_cols=None,
-                         preserve_index=True,
+                         objects_paths: List[str],
+                         file_format: str,
+                         database: str,
+                         table: Optional[str],
+                         partition_cols: Optional[List[str]] = None,
+                         preserve_index: bool = True,
                          mode: str = "append",
-                         compression=None,
-                         cast_columns=None,
-                         extra_args: Optional[Dict[str, Optional[Union[str, int]]]] = None,
+                         compression: Optional[str] = None,
+                         cast_columns: Optional[Dict[str, str]] = None,
+                         extra_args: Optional[Dict[str, Optional[Union[str, int, List[str]]]]] = None,
                          description: Optional[str] = None,
                          parameters: Optional[Dict[str, str]] = None,
                          columns_comments: Optional[Dict[str, str]] = None) -> None:
@@ -88,6 +88,8 @@ def metadata_to_glue(self,
         :return: None
         """
         indexes_position = "left" if file_format == "csv" else "right"
+        schema: List[Tuple[str, str]]
+        partition_cols_schema: List[Tuple[str, str]]
         schema, partition_cols_schema = Glue._build_schema(dataframe=dataframe,
                                                            partition_cols=partition_cols,
                                                            preserve_index=preserve_index,
@@ -138,14 +140,14 @@ def does_table_exists(self, database, table):
             return False
 
     def create_table(self,
-                     database,
-                     table,
-                     schema,
-                     path,
-                     file_format,
-                     compression,
-                     partition_cols_schema=None,
-                     extra_args=None,
+                     database: str,
+                     table: str,
+                     schema: List[Tuple[str, str]],
+                     path: str,
+                     file_format: str,
+                     compression: Optional[str],
+                     partition_cols_schema: List[Tuple[str, str]],
+                     extra_args: Optional[Dict[str, Union[str, int, List[str], None]]] = None,
                      description: Optional[str] = None,
                      parameters: Optional[Dict[str, str]] = None,
                      columns_comments: Optional[Dict[str, str]] = None) -> None:
@@ -166,13 +168,17 @@ def create_table(self,
         :return: None
         """
         if file_format == "parquet":
-            table_input = Glue.parquet_table_definition(table, partition_cols_schema, schema, path, compression)
+            table_input: Dict[str, Any] = Glue.parquet_table_definition(table=table,
+                                                                        partition_cols_schema=partition_cols_schema,
+                                                                        schema=schema,
+                                                                        path=path,
+                                                                        compression=compression)
         elif file_format == "csv":
-            table_input = Glue.csv_table_definition(table,
-                                                    partition_cols_schema,
-                                                    schema,
-                                                    path,
-                                                    compression,
+            table_input = Glue.csv_table_definition(table=table,
+                                                    partition_cols_schema=partition_cols_schema,
+                                                    schema=schema,
+                                                    path=path,
+                                                    compression=compression,
                                                     extra_args=extra_args)
         else:
             raise UnsupportedFileFormat(file_format)
@@ -223,19 +229,23 @@ def get_connection_details(self, name):
         return self._client_glue.get_connection(Name=name, HidePassword=False)["Connection"]
 
     @staticmethod
-    def _build_schema(dataframe, partition_cols, preserve_index, indexes_position, cast_columns=None):
+    def _build_schema(
+            dataframe,
+            partition_cols: Optional[List[str]],
+            preserve_index: bool,
+            indexes_position: str,
+            cast_columns: Optional[Dict[str, str]] = None) -> Tuple[List[Tuple[str, str]], List[Tuple[str, str]]]:
         if cast_columns is None:
             cast_columns = {}
         logger.debug(f"dataframe.dtypes:\n{dataframe.dtypes}")
-        if not partition_cols:
+        if partition_cols is None:
             partition_cols = []
 
-        pyarrow_schema = data_types.extract_pyarrow_schema_from_pandas(dataframe=dataframe,
-                                                                       preserve_index=preserve_index,
-                                                                       indexes_position=indexes_position)
+        pyarrow_schema: List[Tuple[str, str]] = data_types.extract_pyarrow_schema_from_pandas(
+            dataframe=dataframe, preserve_index=preserve_index, indexes_position=indexes_position)
 
-        schema_built = []
-        partition_cols_types = {}
+        schema_built: List[Tuple[str, str]] = []
+        partition_cols_types: Dict[str, str] = {}
         for name, dtype in pyarrow_schema:
             if (cast_columns is not None) and (name in cast_columns.keys()):
                 if name in partition_cols:
@@ -256,7 +266,7 @@ def _build_schema(dataframe, partition_cols, preserve_index, indexes_position, c
                 else:
                     schema_built.append((name, athena_type))
 
-        partition_cols_schema_built = [(name, partition_cols_types[name]) for name in partition_cols]
+        partition_cols_schema_built: List = [(name, partition_cols_types[name]) for name in partition_cols]
 
         logger.debug(f"schema_built:\n{schema_built}")
         logger.debug(f"partition_cols_schema_built:\n{partition_cols_schema_built}")
@@ -269,12 +279,12 @@ def parse_table_name(path):
         return path.rpartition("/")[2]
 
     @staticmethod
-    def csv_table_definition(table,
-                             partition_cols_schema,
-                             schema,
-                             path,
-                             compression,
-                             extra_args: Optional[Dict[str, Optional[Union[str, int]]]] = None):
+    def csv_table_definition(table: str,
+                             partition_cols_schema: List[Tuple[str, str]],
+                             schema: List[Tuple[str, str]],
+                             path: str,
+                             compression: Optional[str],
+                             extra_args: Optional[Dict[str, Optional[Union[str, int, List[str]]]]] = None):
         if extra_args is None:
             extra_args = {"sep": ","}
         if partition_cols_schema is None:
@@ -301,6 +311,9 @@ def csv_table_definition(table,
             refined_schema = [(name, dtype) if dtype in dtypes_allowed else (name, "string") for name, dtype in schema]
         else:
             raise InvalidSerDe(f"{serde} in not in the valid SerDe list.")
+        if "columns" in extra_args:
+            refined_schema = [(name, dtype) for name, dtype in refined_schema
+                              if name in extra_args["columns"]]  # type: ignore
         return {
             "Name": table,
             "PartitionKeys": [{
@@ -378,7 +391,8 @@ def csv_partition_definition(partition, compression, extra_args=None):
         }
 
     @staticmethod
-    def parquet_table_definition(table, partition_cols_schema, schema, path, compression):
+    def parquet_table_definition(table: str, partition_cols_schema: List[Tuple[str, str]],
+                                 schema: List[Tuple[str, str]], path: str, compression: Optional[str]):
         if not partition_cols_schema:
             partition_cols_schema = []
         compressed = False if compression is None else True

awswrangler/pandas.py

@@ -696,6 +696,7 @@ def to_csv(self,
                path: str,
                sep: Optional[str] = None,
                na_rep: Optional[str] = None,
+               columns: Optional[List[str]] = None,
                quoting: Optional[int] = None,
                escapechar: Optional[str] = None,
                serde: Optional[str] = "OpenCSVSerDe",
@@ -718,6 +719,7 @@ def to_csv(self,
         :param path: AWS S3 path (E.g. s3://bucket-name/folder_name/
         :param sep: Same as pandas.to_csv()
         :param na_rep: Same as pandas.to_csv()
+        :param columns: Same as pandas.to_csv()
         :param quoting: Same as pandas.to_csv()
         :param escapechar: Same as pandas.to_csv()
         :param serde: SerDe library name (e.g. OpenCSVSerDe, LazySimpleSerDe) (For Athena/Glue Catalog only)
@@ -738,9 +740,10 @@ def to_csv(self,
             raise InvalidSerDe(f"{serde} in not in the valid SerDe list ({Pandas.VALID_CSV_SERDES})")
         if (database is not None) and (serde is None):
             raise InvalidParameters(f"It is not possible write to a Glue Database without a SerDe.")
-        extra_args: Dict[str, Optional[Union[str, int]]] = {
+        extra_args: Dict[str, Optional[Union[str, int, List[str]]]] = {
             "sep": sep,
             "na_rep": na_rep,
+            "columns": columns,
             "serde": serde,
             "escapechar": escapechar,
             "quoting": quoting
@@ -822,14 +825,14 @@ def to_s3(self,
               file_format: str,
               database: Optional[str] = None,
               table: Optional[str] = None,
-              partition_cols=None,
-              preserve_index=True,
+              partition_cols: Optional[List[str]] = None,
+              preserve_index: bool = True,
               mode: str = "append",
-              compression=None,
-              procs_cpu_bound=None,
-              procs_io_bound=None,
-              cast_columns=None,
-              extra_args: Optional[Dict[str, Optional[Union[str, int]]]] = None,
+              compression: Optional[str] = None,
+              procs_cpu_bound: Optional[int] = None,
+              procs_io_bound: Optional[int] = None,
+              cast_columns: Optional[Dict[str, str]] = None,
+              extra_args: Optional[Dict[str, Optional[Union[str, int, List[str]]]]] = None,
               inplace: bool = True,
               description: Optional[str] = None,
               parameters: Optional[Dict[str, str]] = None,
@@ -866,6 +869,8 @@ def to_s3(self,
         logger.debug(f"cast_columns: {cast_columns}")
         partition_cols = [Athena.normalize_column_name(x) for x in partition_cols]
         logger.debug(f"partition_cols: {partition_cols}")
+        if extra_args is not None and "columns" in extra_args:
+            extra_args["columns"] = [Athena.normalize_column_name(x) for x in extra_args["columns"]]
         dataframe = Pandas.drop_duplicated_columns(dataframe=dataframe, inplace=inplace)
         if compression is not None:
             compression = compression.lower()
@@ -1112,6 +1117,9 @@ def write_csv_dataframe(dataframe, path, preserve_index, compression, fs, extra_
         sep = extra_args.get("sep")
         if sep is not None:
             csv_extra_args["sep"] = sep
+        columns = extra_args.get("columns")
+        if columns is not None:
+            csv_extra_args["columns"] = columns
 
         serde = extra_args.get("serde")
         if serde is None:
@@ -1519,7 +1527,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") and field.name != "__index_level_0__"]
+        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}")
@@ -1612,6 +1623,7 @@ def to_aurora(self,
                   temp_s3_path: Optional[str] = None,
                   preserve_index: bool = False,
                   mode: str = "append",
+                  columns: Optional[List[str]] = None,
                   procs_cpu_bound: Optional[int] = None,
                   procs_io_bound: Optional[int] = None,
                   inplace=True) -> None:
@@ -1626,6 +1638,7 @@ def to_aurora(self,
         :param temp_s3_path: S3 path to write temporary files (E.g. s3://BUCKET_NAME/ANY_NAME/)
         :param preserve_index: Should we preserve the Dataframe index?
         :param mode: append or overwrite
+        :param columns: List of columns to load
         :param procs_cpu_bound: Number of cores used for CPU bound tasks
         :param procs_io_bound: Number of cores used for I/O bound tasks
         :param inplace: True is cheapest (CPU and Memory) but False leaves your DataFrame intact
@@ -1654,6 +1667,7 @@ def to_aurora(self,
                                            serde=None,
                                            sep=",",
                                            na_rep=na_rep,
+                                           columns=columns,
                                            quoting=csv.QUOTE_MINIMAL,
                                            escapechar="\"",
                                            preserve_index=preserve_index,

testing/test_awswrangler/test_pandas.py

@@ -2185,6 +2185,7 @@ def test_to_parquet_categorical_partitions(bucket):
     x['Year'] = x['Year'].astype('category')
     wr.pandas.to_parquet(x[x.Year == 1990], path=path, partition_cols=["Year"])
     y = wr.pandas.read_parquet(path=path)
+    wr.s3.delete_objects(path=path)
     assert len(x[x.Year == 1990].index) == len(y.index)
 
 
@@ -2197,5 +2198,32 @@ def test_range_index(bucket, database):
     print(x)
     wr.pandas.to_parquet(dataframe=x, path=path, database=database)
     df = wr.pandas.read_parquet(path=path)
+    wr.s3.delete_objects(path=path)
     assert len(x.columns) == len(df.columns)
     assert len(x.index) == len(df.index)
+
+
+def test_to_csv_columns(bucket, database):
+    path = f"s3://{bucket}/test_to_csv_columns"
+    wr.s3.delete_objects(path=path)
+    df = pd.DataFrame({
+        "A": [1, 2, 3],
+        "B": [4, 5, 6],
+        "C": ["foo", "boo", "bar"]
+    })
+    wr.s3.delete_objects(path=path)
+    wr.pandas.to_csv(
+        dataframe=df,
+        database=database,
+        path=path,
+        columns=["A", "B"],
+        mode="overwrite",
+        preserve_index=False,
+        procs_cpu_bound=1,
+        inplace=False
+    )
+    sleep(10)
+    df2 = wr.pandas.read_sql_athena(database=database, sql="SELECT * FROM test_to_csv_columns")
+    wr.s3.delete_objects(path=path)
+    assert len(df.columns) == len(df2.columns) + 1
+    assert len(df.index) == len(df2.index)