def _read_text_chunked( paths: List[str], chunksize: int, parser_func: Callable[..., pd.DataFrame], path_root: Optional[str], boto3_session: boto3.Session, pandas_kwargs: Dict[str, Any], s3_additional_kwargs: Optional[Dict[str, str]], dataset: bool, ) -> Iterator[pd.DataFrame]: for path in paths: _logger.debug("path: %s", path) fs: s3fs.S3FileSystem = _utils.get_fs( s3fs_block_size=8_388_608, session=boto3_session, s3_additional_kwargs=s3_additional_kwargs, # 8 MB (8 * 2**20) ) mode, encoding, newline = _get_read_details( path=path, pandas_kwargs=pandas_kwargs) with _utils.open_file(fs=fs, path=path, mode=mode, encoding=encoding, newline=newline) as f: reader: pandas.io.parsers.TextFileReader = parser_func( f, chunksize=chunksize, **pandas_kwargs) for df in reader: yield _apply_partitions(df=df, dataset=dataset, path=path, path_root=path_root)
def _read_parquet_row_group( row_group: int, path: str, columns: Optional[List[str]], categories: Optional[List[str]], boto3_primitives: _utils.Boto3PrimitivesType, s3_additional_kwargs: Optional[Dict[str, str]], ) -> pa.Table: boto3_session: boto3.Session = _utils.boto3_from_primitives( primitives=boto3_primitives) fs: s3fs.S3FileSystem = _utils.get_fs( s3fs_block_size=134_217_728, session=boto3_session, s3_additional_kwargs=s3_additional_kwargs, # 128 MB (128 * 2**20) ) with _utils.open_file(fs=fs, path=path, mode="rb") as f: pq_file: pyarrow.parquet.ParquetFile = pyarrow.parquet.ParquetFile( source=f, read_dictionary=categories) num_row_groups: int = pq_file.num_row_groups _logger.debug("Reading Row Group %s/%s [multi-threaded]", row_group + 1, num_row_groups) return pq_file.read_row_group(i=row_group, columns=columns, use_threads=False, use_pandas_metadata=False)
def _read_text_file( path: str, parser_func: Callable[..., pd.DataFrame], path_root: Optional[str], boto3_session: Union[boto3.Session, Dict[str, Optional[str]]], pandas_kwargs: Dict[str, Any], s3_additional_kwargs: Optional[Dict[str, str]], dataset: bool, ) -> pd.DataFrame: fs: s3fs.S3FileSystem = _utils.get_fs( s3fs_block_size=134_217_728, session=boto3_session, s3_additional_kwargs=s3_additional_kwargs, # 128 MB (128 * 2**20) ) mode, encoding, newline = _get_read_details(path=path, pandas_kwargs=pandas_kwargs) with _utils.open_file(fs=fs, path=path, mode=mode, encoding=encoding, newline=newline) as f: df: pd.DataFrame = parser_func(f, **pandas_kwargs) return _apply_partitions(df=df, dataset=dataset, path=path, path_root=path_root)
def _read_parquet_init( path: Union[str, List[str]], filters: Optional[Union[List[Tuple], List[List[Tuple]]]] = None, categories: List[str] = None, validate_schema: bool = True, dataset: bool = False, use_threads: bool = True, boto3_session: Optional[boto3.Session] = None, s3_additional_kwargs: Optional[Dict[str, str]] = None, ) -> pyarrow.parquet.ParquetDataset: """Encapsulate all initialization before the use of the pyarrow.parquet.ParquetDataset.""" session: boto3.Session = _utils.ensure_session(session=boto3_session) if dataset is False: path_or_paths: Union[str, List[str]] = path2list(path=path, boto3_session=session) elif isinstance(path, str): path_or_paths = path[:-1] if path.endswith("/") else path else: path_or_paths = path _logger.debug("path_or_paths: %s", path_or_paths) fs: s3fs.S3FileSystem = _utils.get_fs( session=session, s3_additional_kwargs=s3_additional_kwargs) cpus: int = _utils.ensure_cpu_count(use_threads=use_threads) data: pyarrow.parquet.ParquetDataset = pyarrow.parquet.ParquetDataset( path_or_paths=path_or_paths, filesystem=fs, metadata_nthreads=cpus, filters=filters, read_dictionary=categories, validate_schema=validate_schema, split_row_groups=False, use_legacy_dataset=True, ) return data
def _read_text_full( parser_func: Callable, path_root: str, path: str, boto3_session: Union[boto3.Session, Dict[str, Optional[str]]], pandas_args: Dict[str, Any], s3_additional_kwargs: Optional[Dict[str, str]] = None, dataset: bool = False, ) -> pd.DataFrame: fs: s3fs.S3FileSystem = _utils.get_fs( session=boto3_session, s3_additional_kwargs=s3_additional_kwargs) if pandas_args.get("compression", "infer") == "infer": pandas_args["compression"] = infer_compression(path, compression="infer") mode: str = "r" if pandas_args.get("compression") is None else "rb" encoding: Optional[str] = pandas_args.get("encoding", None) newline: Optional[str] = pandas_args.get("lineterminator", None) with fs.open(path=path, mode=mode, encoding=encoding, newline=newline) as f: df: pd.DataFrame = parser_func(f, **pandas_args) if dataset is True: partitions: Dict[str, Any] = _utils.extract_partitions_from_path( path_root=path_root, path=path) for column_name, value in partitions.items(): df[column_name] = value return df
def _read_text_chunksize( parser_func: Callable, path_root: str, paths: List[str], boto3_session: boto3.Session, chunksize: int, pandas_args: Dict[str, Any], s3_additional_kwargs: Optional[Dict[str, str]] = None, dataset: bool = False, ) -> Iterator[pd.DataFrame]: fs: s3fs.S3FileSystem = _utils.get_fs( session=boto3_session, s3_additional_kwargs=s3_additional_kwargs) for path in paths: _logger.debug("path: %s", path) partitions: Dict[str, Any] = {} if dataset is True: partitions = _utils.extract_partitions_from_path( path_root=path_root, path=path) if pandas_args.get("compression", "infer") == "infer": pandas_args["compression"] = infer_compression(path, compression="infer") mode: str = "r" if pandas_args.get("compression") is None else "rb" with fs.open(path, mode) as f: reader: pandas.io.parsers.TextFileReader = parser_func( f, chunksize=chunksize, **pandas_args) for df in reader: if dataset is True: for column_name, value in partitions.items(): df[column_name] = value yield df
def _to_text( file_format: str, df: pd.DataFrame, boto3_session: Optional[boto3.Session], s3_additional_kwargs: Optional[Dict[str, str]], path: Optional[str] = None, path_root: Optional[str] = None, **pandas_kwargs, ) -> str: if df.empty is True: raise exceptions.EmptyDataFrame() if path is None and path_root is not None: file_path: str = f"{path_root}{uuid.uuid4().hex}.{file_format}" elif path is not None and path_root is None: file_path = path else: raise RuntimeError("path and path_root received at the same time.") fs: s3fs.S3FileSystem = _utils.get_fs( s3fs_block_size=33_554_432, session=boto3_session, s3_additional_kwargs=s3_additional_kwargs, # 32 MB (32 * 2**20) ) encoding: Optional[str] = pandas_kwargs.get("encoding", None) newline: Optional[str] = pandas_kwargs.get("line_terminator", None) with _utils.open_file(fs=fs, path=file_path, mode="w", encoding=encoding, newline=newline) as f: _logger.debug("pandas_kwargs: %s", pandas_kwargs) if file_format == "csv": df.to_csv(f, **pandas_kwargs) elif file_format == "json": df.to_json(f, **pandas_kwargs) return file_path
def _get_fs( boto3_session: Optional[boto3.Session], s3_additional_kwargs: Optional[Dict[str, str]]) -> s3fs.S3FileSystem: return _utils.get_fs( s3fs_block_size=33_554_432, # 32 MB (32 * 2**20) session=boto3_session, s3_additional_kwargs=s3_additional_kwargs, )
def _to_parquet_file( df: pd.DataFrame, schema: pa.Schema, index: bool, compression: Optional[str], compression_ext: str, cpus: int, dtype: Dict[str, str], boto3_session: Optional[boto3.Session], s3_additional_kwargs: Optional[Dict[str, str]], path: Optional[str] = None, path_root: Optional[str] = None, ) -> str: if path is None and path_root is not None: file_path: str = f"{path_root}{uuid.uuid4().hex}{compression_ext}.parquet" elif path is not None and path_root is None: file_path = path else: raise RuntimeError("path and path_root received at the same time.") _logger.debug("file_path: %s", file_path) table: pa.Table = pyarrow.Table.from_pandas(df=df, schema=schema, nthreads=cpus, preserve_index=index, safe=True) for col_name, col_type in dtype.items(): if col_name in table.column_names: col_index = table.column_names.index(col_name) pyarrow_dtype = _data_types.athena2pyarrow(col_type) field = pa.field(name=col_name, type=pyarrow_dtype) table = table.set_column( col_index, field, table.column(col_name).cast(pyarrow_dtype)) _logger.debug("Casting column %s (%s) to %s (%s)", col_name, col_index, col_type, pyarrow_dtype) fs: s3fs.S3FileSystem = _utils.get_fs( s3fs_block_size=33_554_432, session=boto3_session, s3_additional_kwargs=s3_additional_kwargs, # 32 MB (32 * 2**20) ) with pyarrow.parquet.ParquetWriter( where=file_path, write_statistics=True, use_dictionary=True, filesystem=fs, coerce_timestamps="ms", compression=compression, flavor="spark", schema=table.schema, ) as writer: writer.write_table(table) return file_path
def _read_parquet_metadata_file( path: str, boto3_session: boto3.Session, s3_additional_kwargs: Optional[Dict[str, str]], ) -> Dict[str, str]: fs: s3fs.S3FileSystem = _utils.get_fs( s3fs_block_size=4_194_304, session=boto3_session, s3_additional_kwargs=s3_additional_kwargs # 4 MB (4 * 2**20) ) with _utils.open_file(fs=fs, path=path, mode="rb") as f: pq_file: pyarrow.parquet.ParquetFile = pyarrow.parquet.ParquetFile( source=f) return _data_types.athena_types_from_pyarrow_schema( schema=pq_file.schema.to_arrow_schema(), partitions=None)[0]
def _count_row_groups( path: str, categories: Optional[List[str]], boto3_session: boto3.Session, s3_additional_kwargs: Optional[Dict[str, str]], ) -> int: fs: s3fs.S3FileSystem = _utils.get_fs( s3fs_block_size=4_194_304, session=boto3_session, s3_additional_kwargs=s3_additional_kwargs # 4 MB (4 * 2**20) ) with _utils.open_file(fs=fs, path=path, mode="rb") as f: pq_file: pyarrow.parquet.ParquetFile = pyarrow.parquet.ParquetFile( source=f, read_dictionary=categories) return pq_file.num_row_groups
def _read_parquet_file( path: str, columns: Optional[List[str]], categories: Optional[List[str]], boto3_session: boto3.Session, s3_additional_kwargs: Optional[Dict[str, str]], ) -> pa.Table: fs: s3fs.S3FileSystem = _utils.get_fs( s3fs_block_size=134_217_728, session=boto3_session, s3_additional_kwargs=s3_additional_kwargs, # 128 MB (128 * 2**20) ) with _utils.open_file(fs=fs, path=path, mode="rb") as f: pq_file: pyarrow.parquet.ParquetFile = pyarrow.parquet.ParquetFile( source=f, read_dictionary=categories) return pq_file.read(columns=columns, use_threads=False, use_pandas_metadata=False)
def _to_text( file_format: str, df: pd.DataFrame, path: str, fs: Optional[s3fs.S3FileSystem] = None, boto3_session: Optional[boto3.Session] = None, s3_additional_kwargs: Optional[Dict[str, str]] = None, **pandas_kwargs, ) -> None: if df.empty is True: # pragma: no cover raise exceptions.EmptyDataFrame() if fs is None: fs = _utils.get_fs(session=boto3_session, s3_additional_kwargs=s3_additional_kwargs) encoding: Optional[str] = pandas_kwargs.get("encoding", None) newline: Optional[str] = pandas_kwargs.get("line_terminator", None) with fs.open(path=path, mode="w", encoding=encoding, newline=newline) as f: if file_format == "csv": df.to_csv(f, **pandas_kwargs) elif file_format == "json": df.to_json(f, **pandas_kwargs)
def to_parquet( # pylint: disable=too-many-arguments,too-many-locals df: pd.DataFrame, path: str, index: bool = False, compression: Optional[str] = "snappy", use_threads: bool = True, boto3_session: Optional[boto3.Session] = None, s3_additional_kwargs: Optional[Dict[str, str]] = None, sanitize_columns: bool = False, dataset: bool = False, partition_cols: Optional[List[str]] = None, mode: Optional[str] = None, catalog_versioning: bool = False, database: Optional[str] = None, table: Optional[str] = None, dtype: Optional[Dict[str, str]] = None, description: Optional[str] = None, parameters: Optional[Dict[str, str]] = None, columns_comments: Optional[Dict[str, str]] = None, regular_partitions: bool = True, projection_enabled: bool = False, projection_types: Optional[Dict[str, str]] = None, projection_ranges: Optional[Dict[str, str]] = None, projection_values: Optional[Dict[str, str]] = None, projection_intervals: Optional[Dict[str, str]] = None, projection_digits: Optional[Dict[str, str]] = None, ) -> Dict[str, Union[List[str], Dict[str, List[str]]]]: """Write Parquet file or dataset on Amazon S3. The concept of Dataset goes beyond the simple idea of files and enable more complex features like partitioning, casting and catalog integration (Amazon Athena/AWS Glue Catalog). Note ---- If `dataset=True` The table name and all column names will be automatically sanitized using `wr.catalog.sanitize_table_name` and `wr.catalog.sanitize_column_name`. Please, pass `sanitize_columns=True` to force the same behaviour for `dataset=False`. Note ---- On `append` mode, the `parameters` will be upsert on an existing table. Note ---- In case of `use_threads=True` the number of threads that will be spawned will be get from os.cpu_count(). Parameters ---------- df: pandas.DataFrame Pandas DataFrame https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.html path : str S3 path (for file e.g. ``s3://bucket/prefix/filename.parquet``) (for dataset e.g. ``s3://bucket/prefix``). index : bool True to store the DataFrame index in file, otherwise False to ignore it. compression: str, optional Compression style (``None``, ``snappy``, ``gzip``). use_threads : bool True to enable concurrent requests, False to disable multiple threads. If enabled os.cpu_count() will be used as the max number of threads. boto3_session : boto3.Session(), optional Boto3 Session. The default boto3 session will be used if boto3_session receive None. s3_additional_kwargs: Forward to s3fs, useful for server side encryption https://s3fs.readthedocs.io/en/latest/#serverside-encryption sanitize_columns : bool True to sanitize columns names or False to keep it as is. True value is forced if `dataset=True`. dataset : bool If True store a parquet dataset instead of a single file. If True, enable all follow arguments: partition_cols, mode, database, table, description, parameters, columns_comments, . partition_cols: List[str], optional List of column names that will be used to create partitions. Only takes effect if dataset=True. mode: str, optional ``append`` (Default), ``overwrite``, ``overwrite_partitions``. Only takes effect if dataset=True. catalog_versioning : bool If True and `mode="overwrite"`, creates an archived version of the table catalog before updating it. database : str, optional Glue/Athena catalog: Database name. table : str, optional Glue/Athena catalog: Table name. dtype : Dict[str, str], optional Dictionary of columns names and Athena/Glue types to be casted. Useful when you have columns with undetermined or mixed data types. (e.g. {'col name': 'bigint', 'col2 name': 'int'}) description : str, optional Glue/Athena catalog: Table description parameters : Dict[str, str], optional Glue/Athena catalog: Key/value pairs to tag the table. columns_comments : Dict[str, str], optional Glue/Athena catalog: Columns names and the related comments (e.g. {'col0': 'Column 0.', 'col1': 'Column 1.', 'col2': 'Partition.'}). regular_partitions : bool Create regular partitions (Non projected partitions) on Glue Catalog. Disable when you will work only with Partition Projection. Keep enabled even when working with projections is useful to keep Redshift Spectrum working with the regular partitions. projection_enabled : bool Enable Partition Projection on Athena (https://docs.aws.amazon.com/athena/latest/ug/partition-projection.html) projection_types : Optional[Dict[str, str]] Dictionary of partitions names and Athena projections types. Valid types: "enum", "integer", "date", "injected" https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html (e.g. {'col_name': 'enum', 'col2_name': 'integer'}) projection_ranges: Optional[Dict[str, str]] Dictionary of partitions names and Athena projections ranges. https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html (e.g. {'col_name': '0,10', 'col2_name': '-1,8675309'}) projection_values: Optional[Dict[str, str]] Dictionary of partitions names and Athena projections values. https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html (e.g. {'col_name': 'A,B,Unknown', 'col2_name': 'foo,boo,bar'}) projection_intervals: Optional[Dict[str, str]] Dictionary of partitions names and Athena projections intervals. https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html (e.g. {'col_name': '1', 'col2_name': '5'}) projection_digits: Optional[Dict[str, str]] Dictionary of partitions names and Athena projections digits. https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html (e.g. {'col_name': '1', 'col2_name': '2'}) Returns ------- Dict[str, Union[List[str], Dict[str, List[str]]]] Dictionary with: 'paths': List of all stored files paths on S3. 'partitions_values': Dictionary of partitions added with keys as S3 path locations and values as a list of partitions values as str. Examples -------- Writing single file >>> import awswrangler as wr >>> import pandas as pd >>> wr.s3.to_parquet( ... df=pd.DataFrame({'col': [1, 2, 3]}), ... path='s3://bucket/prefix/my_file.parquet', ... ) { 'paths': ['s3://bucket/prefix/my_file.parquet'], 'partitions_values': {} } Writing single file encrypted with a KMS key >>> import awswrangler as wr >>> import pandas as pd >>> wr.s3.to_parquet( ... df=pd.DataFrame({'col': [1, 2, 3]}), ... path='s3://bucket/prefix/my_file.parquet', ... s3_additional_kwargs={ ... 'ServerSideEncryption': 'aws:kms', ... 'SSEKMSKeyId': 'YOUR_KMY_KEY_ARN' ... } ... ) { 'paths': ['s3://bucket/prefix/my_file.parquet'], 'partitions_values': {} } Writing partitioned dataset >>> import awswrangler as wr >>> import pandas as pd >>> wr.s3.to_parquet( ... df=pd.DataFrame({ ... 'col': [1, 2, 3], ... 'col2': ['A', 'A', 'B'] ... }), ... path='s3://bucket/prefix', ... dataset=True, ... partition_cols=['col2'] ... ) { 'paths': ['s3://.../col2=A/x.parquet', 's3://.../col2=B/y.parquet'], 'partitions_values: { 's3://.../col2=A/': ['A'], 's3://.../col2=B/': ['B'] } } Writing dataset to S3 with metadata on Athena/Glue Catalog. >>> import awswrangler as wr >>> import pandas as pd >>> wr.s3.to_parquet( ... df=pd.DataFrame({ ... 'col': [1, 2, 3], ... 'col2': ['A', 'A', 'B'] ... }), ... path='s3://bucket/prefix', ... dataset=True, ... partition_cols=['col2'], ... database='default', # Athena/Glue database ... table='my_table' # Athena/Glue table ... ) { 'paths': ['s3://.../col2=A/x.parquet', 's3://.../col2=B/y.parquet'], 'partitions_values: { 's3://.../col2=A/': ['A'], 's3://.../col2=B/': ['B'] } } Writing dataset casting empty column data type >>> import awswrangler as wr >>> import pandas as pd >>> wr.s3.to_parquet( ... df=pd.DataFrame({ ... 'col': [1, 2, 3], ... 'col2': ['A', 'A', 'B'], ... 'col3': [None, None, None] ... }), ... path='s3://bucket/prefix', ... dataset=True, ... database='default', # Athena/Glue database ... table='my_table' # Athena/Glue table ... dtype={'col3': 'date'} ... ) { 'paths': ['s3://.../x.parquet'], 'partitions_values: {} } """ if (database is None) ^ (table is None): raise exceptions.InvalidArgumentCombination( "Please pass database and table arguments to be able to store the metadata into the Athena/Glue Catalog." ) if df.empty is True: raise exceptions.EmptyDataFrame() partition_cols = partition_cols if partition_cols else [] dtype = dtype if dtype else {} partitions_values: Dict[str, List[str]] = {} # Sanitize table to respect Athena's standards if (sanitize_columns is True) or (dataset is True): df = catalog.sanitize_dataframe_columns_names(df=df) partition_cols = [catalog.sanitize_column_name(p) for p in partition_cols] dtype = {catalog.sanitize_column_name(k): v.lower() for k, v in dtype.items()} catalog.drop_duplicated_columns(df=df) session: boto3.Session = _utils.ensure_session(session=boto3_session) cpus: int = _utils.ensure_cpu_count(use_threads=use_threads) fs: s3fs.S3FileSystem = _utils.get_fs(session=session, s3_additional_kwargs=s3_additional_kwargs) compression_ext: Optional[str] = _COMPRESSION_2_EXT.get(compression, None) if compression_ext is None: raise exceptions.InvalidCompression(f"{compression} is invalid, please use None, snappy or gzip.") if dataset is False: if path.endswith("/"): # pragma: no cover raise exceptions.InvalidArgumentValue( "If <dataset=False>, the argument <path> should be a object path, not a directory." ) if partition_cols: raise exceptions.InvalidArgumentCombination("Please, pass dataset=True to be able to use partition_cols.") if mode is not None: raise exceptions.InvalidArgumentCombination("Please pass dataset=True to be able to use mode.") if any(arg is not None for arg in (database, table, description, parameters)): raise exceptions.InvalidArgumentCombination( "Please pass dataset=True to be able to use any one of these " "arguments: database, table, description, parameters, " "columns_comments." ) df = _data_types.cast_pandas_with_athena_types(df=df, dtype=dtype) schema: pa.Schema = _data_types.pyarrow_schema_from_pandas( df=df, index=index, ignore_cols=partition_cols, dtype=dtype ) _logger.debug("schema: \n%s", schema) paths = [ _to_parquet_file( df=df, path=path, schema=schema, index=index, compression=compression, cpus=cpus, fs=fs, dtype=dtype ) ] else: mode = "append" if mode is None else mode if ( (mode in ("append", "overwrite_partitions")) and (database is not None) and (table is not None) ): # Fetching Catalog Types catalog_types: Optional[Dict[str, str]] = catalog.get_table_types( database=database, table=table, boto3_session=session ) if catalog_types is not None: for k, v in catalog_types.items(): dtype[k] = v paths, partitions_values = _to_parquet_dataset( df=df, path=path, index=index, compression=compression, compression_ext=compression_ext, cpus=cpus, fs=fs, use_threads=use_threads, partition_cols=partition_cols, dtype=dtype, mode=mode, boto3_session=session, ) if (database is not None) and (table is not None): columns_types, partitions_types = _data_types.athena_types_from_pandas_partitioned( df=df, index=index, partition_cols=partition_cols, dtype=dtype ) catalog.create_parquet_table( database=database, table=table, path=path, columns_types=columns_types, partitions_types=partitions_types, compression=compression, description=description, parameters=parameters, columns_comments=columns_comments, boto3_session=session, mode=mode, catalog_versioning=catalog_versioning, projection_enabled=projection_enabled, projection_types=projection_types, projection_ranges=projection_ranges, projection_values=projection_values, projection_intervals=projection_intervals, projection_digits=projection_digits, ) if partitions_values and (regular_partitions is True): _logger.debug("partitions_values:\n%s", partitions_values) catalog.add_parquet_partitions( database=database, table=table, partitions_values=partitions_values, compression=compression, boto3_session=session, ) return {"paths": paths, "partitions_values": partitions_values}
def to_csv( # pylint: disable=too-many-arguments,too-many-locals df: pd.DataFrame, path: str, sep: str = ",", index: bool = True, columns: Optional[List[str]] = None, use_threads: bool = True, boto3_session: Optional[boto3.Session] = None, s3_additional_kwargs: Optional[Dict[str, str]] = None, sanitize_columns: bool = False, dataset: bool = False, partition_cols: Optional[List[str]] = None, mode: Optional[str] = None, catalog_versioning: bool = False, database: Optional[str] = None, table: Optional[str] = None, dtype: Optional[Dict[str, str]] = None, description: Optional[str] = None, parameters: Optional[Dict[str, str]] = None, columns_comments: Optional[Dict[str, str]] = None, regular_partitions: bool = True, projection_enabled: bool = False, projection_types: Optional[Dict[str, str]] = None, projection_ranges: Optional[Dict[str, str]] = None, projection_values: Optional[Dict[str, str]] = None, projection_intervals: Optional[Dict[str, str]] = None, projection_digits: Optional[Dict[str, str]] = None, **pandas_kwargs, ) -> Dict[str, Union[List[str], Dict[str, List[str]]]]: """Write CSV file or dataset on Amazon S3. The concept of Dataset goes beyond the simple idea of files and enable more complex features like partitioning, casting and catalog integration (Amazon Athena/AWS Glue Catalog). Note ---- If `dataset=True` The table name and all column names will be automatically sanitized using `wr.catalog.sanitize_table_name` and `wr.catalog.sanitize_column_name`. Please, pass `sanitize_columns=True` to force the same behaviour for `dataset=False`. Note ---- If `dataset=True`, `pandas_kwargs` will be ignored due restrictive quoting, date_format, escapechar, encoding, etc required by Athena/Glue Catalog. Note ---- By now Pandas does not support in-memory CSV compression. https://github.com/pandas-dev/pandas/issues/22555 So the `compression` will not be supported on Wrangler too. Note ---- On `append` mode, the `parameters` will be upsert on an existing table. Note ---- In case of `use_threads=True` the number of threads that will be spawned will be get from os.cpu_count(). Parameters ---------- df: pandas.DataFrame Pandas DataFrame https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.html path : str Amazon S3 path (e.g. s3://bucket/filename.csv). sep : str String of length 1. Field delimiter for the output file. index : bool Write row names (index). columns : List[str], optional Columns to write. use_threads : bool True to enable concurrent requests, False to disable multiple threads. If enabled os.cpu_count() will be used as the max number of threads. boto3_session : boto3.Session(), optional Boto3 Session. The default boto3 Session will be used if boto3_session receive None. s3_additional_kwargs: Forward to s3fs, useful for server side encryption https://s3fs.readthedocs.io/en/latest/#serverside-encryption sanitize_columns : bool True to sanitize columns names or False to keep it as is. True value is forced if `dataset=True`. dataset : bool If True store a parquet dataset instead of a single file. If True, enable all follow arguments: partition_cols, mode, database, table, description, parameters, columns_comments, . partition_cols: List[str], optional List of column names that will be used to create partitions. Only takes effect if dataset=True. mode : str, optional ``append`` (Default), ``overwrite``, ``overwrite_partitions``. Only takes effect if dataset=True. catalog_versioning : bool If True and `mode="overwrite"`, creates an archived version of the table catalog before updating it. database : str, optional Glue/Athena catalog: Database name. table : str, optional Glue/Athena catalog: Table name. dtype : Dict[str, str], optional Dictionary of columns names and Athena/Glue types to be casted. Useful when you have columns with undetermined or mixed data types. (e.g. {'col name': 'bigint', 'col2 name': 'int'}) description : str, optional Glue/Athena catalog: Table description parameters : Dict[str, str], optional Glue/Athena catalog: Key/value pairs to tag the table. columns_comments : Dict[str, str], optional Glue/Athena catalog: Columns names and the related comments (e.g. {'col0': 'Column 0.', 'col1': 'Column 1.', 'col2': 'Partition.'}). regular_partitions : bool Create regular partitions (Non projected partitions) on Glue Catalog. Disable when you will work only with Partition Projection. Keep enabled even when working with projections is useful to keep Redshift Spectrum working with the regular partitions. projection_enabled : bool Enable Partition Projection on Athena (https://docs.aws.amazon.com/athena/latest/ug/partition-projection.html) projection_types : Optional[Dict[str, str]] Dictionary of partitions names and Athena projections types. Valid types: "enum", "integer", "date", "injected" https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html (e.g. {'col_name': 'enum', 'col2_name': 'integer'}) projection_ranges: Optional[Dict[str, str]] Dictionary of partitions names and Athena projections ranges. https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html (e.g. {'col_name': '0,10', 'col2_name': '-1,8675309'}) projection_values: Optional[Dict[str, str]] Dictionary of partitions names and Athena projections values. https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html (e.g. {'col_name': 'A,B,Unknown', 'col2_name': 'foo,boo,bar'}) projection_intervals: Optional[Dict[str, str]] Dictionary of partitions names and Athena projections intervals. https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html (e.g. {'col_name': '1', 'col2_name': '5'}) projection_digits: Optional[Dict[str, str]] Dictionary of partitions names and Athena projections digits. https://docs.aws.amazon.com/athena/latest/ug/partition-projection-supported-types.html (e.g. {'col_name': '1', 'col2_name': '2'}) pandas_kwargs : keyword arguments forwarded to pandas.DataFrame.to_csv() https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.to_csv.html Returns ------- None None. Examples -------- Writing single file >>> import awswrangler as wr >>> import pandas as pd >>> wr.s3.to_csv( ... df=pd.DataFrame({'col': [1, 2, 3]}), ... path='s3://bucket/prefix/my_file.csv', ... ) { 'paths': ['s3://bucket/prefix/my_file.csv'], 'partitions_values': {} } Writing single file encrypted with a KMS key >>> import awswrangler as wr >>> import pandas as pd >>> wr.s3.to_csv( ... df=pd.DataFrame({'col': [1, 2, 3]}), ... path='s3://bucket/prefix/my_file.csv', ... s3_additional_kwargs={ ... 'ServerSideEncryption': 'aws:kms', ... 'SSEKMSKeyId': 'YOUR_KMY_KEY_ARN' ... } ... ) { 'paths': ['s3://bucket/prefix/my_file.csv'], 'partitions_values': {} } Writing partitioned dataset >>> import awswrangler as wr >>> import pandas as pd >>> wr.s3.to_csv( ... df=pd.DataFrame({ ... 'col': [1, 2, 3], ... 'col2': ['A', 'A', 'B'] ... }), ... path='s3://bucket/prefix', ... dataset=True, ... partition_cols=['col2'] ... ) { 'paths': ['s3://.../col2=A/x.csv', 's3://.../col2=B/y.csv'], 'partitions_values: { 's3://.../col2=A/': ['A'], 's3://.../col2=B/': ['B'] } } Writing dataset to S3 with metadata on Athena/Glue Catalog. >>> import awswrangler as wr >>> import pandas as pd >>> wr.s3.to_csv( ... df=pd.DataFrame({ ... 'col': [1, 2, 3], ... 'col2': ['A', 'A', 'B'] ... }), ... path='s3://bucket/prefix', ... dataset=True, ... partition_cols=['col2'], ... database='default', # Athena/Glue database ... table='my_table' # Athena/Glue table ... ) { 'paths': ['s3://.../col2=A/x.csv', 's3://.../col2=B/y.csv'], 'partitions_values: { 's3://.../col2=A/': ['A'], 's3://.../col2=B/': ['B'] } } Writing dataset casting empty column data type >>> import awswrangler as wr >>> import pandas as pd >>> wr.s3.to_csv( ... df=pd.DataFrame({ ... 'col': [1, 2, 3], ... 'col2': ['A', 'A', 'B'], ... 'col3': [None, None, None] ... }), ... path='s3://bucket/prefix', ... dataset=True, ... database='default', # Athena/Glue database ... table='my_table' # Athena/Glue table ... dtype={'col3': 'date'} ... ) { 'paths': ['s3://.../x.csv'], 'partitions_values: {} } """ if (database is None) ^ (table is None): raise exceptions.InvalidArgumentCombination( "Please pass database and table arguments to be able to store the metadata into the Athena/Glue Catalog." ) if df.empty is True: raise exceptions.EmptyDataFrame() partition_cols = partition_cols if partition_cols else [] dtype = dtype if dtype else {} partitions_values: Dict[str, List[str]] = {} # Sanitize table to respect Athena's standards if (sanitize_columns is True) or (dataset is True): df = catalog.sanitize_dataframe_columns_names(df=df) partition_cols = [catalog.sanitize_column_name(p) for p in partition_cols] dtype = {catalog.sanitize_column_name(k): v.lower() for k, v in dtype.items()} catalog.drop_duplicated_columns(df=df) session: boto3.Session = _utils.ensure_session(session=boto3_session) fs: s3fs.S3FileSystem = _utils.get_fs(session=session, s3_additional_kwargs=s3_additional_kwargs) if dataset is False: if partition_cols: raise exceptions.InvalidArgumentCombination("Please, pass dataset=True to be able to use partition_cols.") if mode is not None: raise exceptions.InvalidArgumentCombination("Please pass dataset=True to be able to use mode.") if columns_comments: raise exceptions.InvalidArgumentCombination("Please pass dataset=True to be able to use columns_comments.") if any(arg is not None for arg in (database, table, description, parameters)): raise exceptions.InvalidArgumentCombination( "Please pass dataset=True to be able to use any one of these " "arguments: database, table, description, parameters, " "columns_comments." ) pandas_kwargs["sep"] = sep pandas_kwargs["index"] = index pandas_kwargs["columns"] = columns _to_text(file_format="csv", df=df, path=path, fs=fs, **pandas_kwargs) paths = [path] else: mode = "append" if mode is None else mode if columns: df = df[columns] if ( (mode in ("append", "overwrite_partitions")) and (database is not None) and (table is not None) ): # Fetching Catalog Types catalog_types: Optional[Dict[str, str]] = catalog.get_table_types( database=database, table=table, boto3_session=session ) if catalog_types is not None: for k, v in catalog_types.items(): dtype[k] = v paths, partitions_values = _to_csv_dataset( df=df, path=path, index=index, sep=sep, fs=fs, use_threads=use_threads, partition_cols=partition_cols, dtype=dtype, mode=mode, boto3_session=session, ) if (database is not None) and (table is not None): columns_types, partitions_types = _data_types.athena_types_from_pandas_partitioned( df=df, index=index, partition_cols=partition_cols, dtype=dtype, index_left=True ) catalog.create_csv_table( database=database, table=table, path=path, columns_types=columns_types, partitions_types=partitions_types, description=description, parameters=parameters, columns_comments=columns_comments, boto3_session=session, mode=mode, catalog_versioning=catalog_versioning, sep=sep, projection_enabled=projection_enabled, projection_types=projection_types, projection_ranges=projection_ranges, projection_values=projection_values, projection_intervals=projection_intervals, projection_digits=projection_digits, ) if partitions_values and (regular_partitions is True): _logger.debug("partitions_values:\n%s", partitions_values) catalog.add_csv_partitions( database=database, table=table, partitions_values=partitions_values, boto3_session=session, sep=sep ) return {"paths": paths, "partitions_values": partitions_values}
def _read_parquet_chunked( paths: List[str], chunked: Union[bool, int], columns: Optional[List[str]], categories: Optional[List[str]], validate_schema: bool, safe: bool, boto3_session: boto3.Session, dataset: bool, path_root: Optional[str], s3_additional_kwargs: Optional[Dict[str, str]], use_threads: bool, ) -> Iterator[pd.DataFrame]: next_slice: Optional[pd.DataFrame] = None fs: s3fs.S3FileSystem = _utils.get_fs( s3fs_block_size=8_388_608, session=boto3_session, s3_additional_kwargs=s3_additional_kwargs # 8 MB (8 * 2**20) ) last_schema: Optional[Dict[str, str]] = None last_path: str = "" for path in paths: with _utils.open_file(fs=fs, path=path, mode="rb") as f: pq_file: pyarrow.parquet.ParquetFile = pyarrow.parquet.ParquetFile( source=f, read_dictionary=categories) schema: Dict[str, str] = _data_types.athena_types_from_pyarrow_schema( schema=pq_file.schema.to_arrow_schema(), partitions=None)[0] if validate_schema is True and last_schema is not None: if schema != last_schema: raise exceptions.InvalidSchemaConvergence( f"Was detect at least 2 different schemas:\n" f" - {last_path} -> {last_schema}\n" f" - {path} -> {schema}") last_schema = schema last_path = path num_row_groups: int = pq_file.num_row_groups _logger.debug("num_row_groups: %s", num_row_groups) for i in range(num_row_groups): _logger.debug("Reading Row Group %s...", i) df: pd.DataFrame = _arrowtable2df( table=pq_file.read_row_group(i=i, columns=columns, use_threads=use_threads, use_pandas_metadata=False), categories=categories, safe=safe, use_threads=use_threads, dataset=dataset, path=path, path_root=path_root, ) if chunked is True: yield df elif isinstance(chunked, int) and chunked > 0: if next_slice is not None: df = pd.concat(objs=[next_slice, df], ignore_index=True, sort=False, copy=False) while len(df.index) >= chunked: yield df.iloc[:chunked] df = df.iloc[chunked:] if df.empty: next_slice = None else: next_slice = df else: raise exceptions.InvalidArgument(f"chunked: {chunked}") if next_slice is not None: yield next_slice