Repository: aws/aws-sdk-pandas

Issue:
Add Common Attributes to Timestream write
Writing to Timestream with common attributes like dimensions can make a massive cost difference (size of the attributes * number of rows). Right now, you can supply dimension columns, but that requires two things that are inefficient. 

1) adding a repetitive column to a potentially massive dataframe if it does not already exist 
2) writing these repetitive values to each record rather than once per 100-record batch

**Describe the solution you'd like**
One thing that would be great would be to specify, at the dataframe level, a dictionary of common attributes to be passed with all record batches in the frame, assuming they are all consistent.

e.g., if I have a CSV I am reading from that will all be from the same IOT device, I should be able to specify comething like
```
common_attributes = {
            'Dimensions': {‘Name’: ‘device_id’, ‘Value’: ‘12345’},
            'MeasureValueType': 'DOUBLE',
            'MeasureName': measure_name,
        }
````

The most flexible solution would be for a user to supply common_attribute columns when calling `wr.timestream.write`. When processing the dataframe, the library should groupby the specified columns and then send the batches from subgroups with the specified common attributes.
 
**Describe alternatives you've considered**
The solution I’m using on my end right now involves specifying a group id column to the dataframe and looping through using a dictionary specifying the common attributes for that group.


Discussion:
Thanks @ddemarco96, I agree that adding a `common_attributes` parameter to the api is warranted. However, I am not in favour of fetching it via a data frame column. The library focuses on I/O operations and we tend to avoid manipulating data frames with groupbys. The simplest implementation is to enable passing a `common_attributes` dictionary like in the boto3 [call](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/timestream-write/client/write_records.html). The user can always handle the groupby themselves if they wish.

Resolving patch:
diff --git a/awswrangler/timestream.py b/awswrangler/timestream.py
--- a/awswrangler/timestream.py
+++ b/awswrangler/timestream.py
@@ -10,7 +10,7 @@
 import pandas as pd
 from botocore.config import Config
 
-from awswrangler import _data_types, _utils
+from awswrangler import _data_types, _utils, exceptions
 
 _logger: logging.Logger = logging.getLogger(__name__)
 
@@ -27,61 +27,111 @@ def _df2list(df: pd.DataFrame) -> List[List[Any]]:
     return parameters
 
 
+def _format_timestamp(timestamp: Union[int, datetime]) -> str:
+    if isinstance(timestamp, int):
+        return str(round(timestamp / 1_000_000))
+    if isinstance(timestamp, datetime):
+        return str(round(timestamp.timestamp() * 1_000))
+    raise exceptions.InvalidArgumentType("`time_col` must be of type timestamp.")
+
+
 def _format_measure(measure_name: str, measure_value: Any, measure_type: str) -> Dict[str, str]:
     return {
         "Name": measure_name,
-        "Value": str(round(measure_value.timestamp() * 1_000) if measure_type == "TIMESTAMP" else measure_value),
+        "Value": _format_timestamp(measure_value) if measure_type == "TIMESTAMP" else str(measure_value),
         "Type": measure_type,
     }
 
 
+def _sanitize_common_attributes(
+    common_attributes: Optional[Dict[str, Any]],
+    version: int,
+    measure_name: Optional[str],
+) -> Dict[str, Any]:
+    common_attributes = {} if not common_attributes else common_attributes
+    # Values in common_attributes take precedence
+    common_attributes.setdefault("Version", version)
+
+    if "Time" not in common_attributes:
+        # TimeUnit is MILLISECONDS by default for Timestream writes
+        # But if a time_col is supplied (i.e. Time is not in common_attributes)
+        # then TimeUnit must be set to MILLISECONDS explicitly
+        common_attributes["TimeUnit"] = "MILLISECONDS"
+
+    if "MeasureValue" in common_attributes and "MeasureValueType" not in common_attributes:
+        raise exceptions.InvalidArgumentCombination(
+            "MeasureValueType must be supplied alongside MeasureValue in common_attributes."
+        )
+
+    if measure_name:
+        common_attributes.setdefault("MeasureName", measure_name)
+    elif "MeasureName" not in common_attributes:
+        raise exceptions.InvalidArgumentCombination(
+            "MeasureName must be supplied with the `measure_name` argument or in common_attributes."
+        )
+    return common_attributes
+
+
 def _write_batch(
+    timestream_client: boto3.client,
     database: str,
     table: str,
-    cols_names: List[str],
-    measure_cols_names: List[str],
+    common_attributes: Dict[str, Any],
+    cols_names: List[Optional[str]],
+    measure_cols: List[Optional[str]],
     measure_types: List[str],
-    version: int,
+    dimension_cols: List[Optional[str]],
     batch: List[Any],
-    timestream_client: boto3.client,
-    measure_name: Optional[str] = None,
 ) -> List[Dict[str, str]]:
-    try:
-        time_loc = 0
-        measure_cols_loc = 1
-        dimensions_cols_loc = 1 + len(measure_cols_names)
-        records: List[Dict[str, Any]] = []
-        for rec in batch:
-            record: Dict[str, Any] = {
-                "Dimensions": [
-                    {"Name": name, "DimensionValueType": "VARCHAR", "Value": str(value)}
-                    for name, value in zip(cols_names[dimensions_cols_loc:], rec[dimensions_cols_loc:])
-                ],
-                "Time": str(round(rec[time_loc].timestamp() * 1_000)),
-                "TimeUnit": "MILLISECONDS",
-                "Version": version,
-            }
-            if len(measure_cols_names) == 1:
-                measure_value = rec[measure_cols_loc]
-                if pd.isnull(measure_value):
-                    continue
-                record["MeasureName"] = measure_name if measure_name else measure_cols_names[0]
-                record["MeasureValueType"] = measure_types[0]
-                record["MeasureValue"] = str(measure_value)
-            else:
-                record["MeasureName"] = measure_name if measure_name else measure_cols_names[0]
-                record["MeasureValueType"] = "MULTI"
-                record["MeasureValues"] = [
-                    _format_measure(measure_name, measure_value, measure_value_type)
-                    for measure_name, measure_value, measure_value_type in zip(
-                        measure_cols_names, rec[measure_cols_loc:dimensions_cols_loc], measure_types
-                    )
-                    if not pd.isnull(measure_value)
-                ]
-                if len(record["MeasureValues"]) == 0:
-                    continue
+    records: List[Dict[str, Any]] = []
+    scalar = bool(len(measure_cols) == 1 and "MeasureValues" not in common_attributes)
+    time_loc = 0
+    measure_cols_loc = 1 if cols_names[0] else 0
+    dimensions_cols_loc = 1 if len(measure_cols) == 1 else 1 + len(measure_cols)
+    if all(cols_names):
+        # Time and Measures are supplied in the data frame
+        dimensions_cols_loc = 1 + len(measure_cols)
+    elif all(v is None for v in cols_names[:2]):
+        # Time and Measures are supplied in common_attributes
+        dimensions_cols_loc = 0
+
+    for row in batch:
+        record: Dict[str, Any] = {}
+        if "Time" not in common_attributes:
+            record["Time"] = _format_timestamp(row[time_loc])
+        if scalar and "MeasureValue" not in common_attributes:
+            measure_value = row[measure_cols_loc]
+            if pd.isnull(measure_value):
+                continue
+            record["MeasureValue"] = str(measure_value)
+        elif not scalar and "MeasureValues" not in common_attributes:
+            record["MeasureValues"] = [
+                _format_measure(measure_name, measure_value, measure_value_type)  # type: ignore[arg-type]
+                for measure_name, measure_value, measure_value_type in zip(
+                    measure_cols, row[measure_cols_loc:dimensions_cols_loc], measure_types
+                )
+                if not pd.isnull(measure_value)
+            ]
+            if len(record["MeasureValues"]) == 0:
+                continue
+        if "MeasureValueType" not in common_attributes:
+            record["MeasureValueType"] = measure_types[0] if scalar else "MULTI"
+        # Dimensions can be specified in both common_attributes and the data frame
+        dimensions = (
+            [
+                {"Name": name, "DimensionValueType": "VARCHAR", "Value": str(value)}
+                for name, value in zip(dimension_cols, row[dimensions_cols_loc:])
+            ]
+            if all(dimension_cols)
+            else []
+        )
+        if dimensions:
+            record["Dimensions"] = dimensions
+        if record:
             records.append(record)
-        if len(records) > 0:
+
+    try:
+        if records:
             _utils.try_it(
                 f=timestream_client.write_records,
                 ex=(
@@ -91,6 +141,7 @@ def _write_batch(
                 max_num_tries=5,
                 DatabaseName=database,
                 TableName=table,
+                CommonAttributes=common_attributes,
                 Records=records,
             )
     except timestream_client.exceptions.RejectedRecordsException as ex:
@@ -192,12 +243,13 @@ def write(
     df: pd.DataFrame,
     database: str,
     table: str,
-    time_col: str,
-    measure_col: Union[str, List[str]],
-    dimensions_cols: List[str],
+    time_col: Optional[str] = None,
+    measure_col: Union[str, List[Optional[str]], None] = None,
+    dimensions_cols: Optional[List[Optional[str]]] = None,
     version: int = 1,
     num_threads: int = 32,
     measure_name: Optional[str] = None,
+    common_attributes: Optional[Dict[str, Any]] = None,
     boto3_session: Optional[boto3.Session] = None,
 ) -> List[Dict[str, str]]:
     """Store a Pandas DataFrame into a Amazon Timestream table.
@@ -206,6 +258,16 @@ def write(
     this function will not throw a Python exception.
     Instead it will return the rejection information.
 
+    Note
+    ----
+    Values in `common_attributes` take precedence over all other arguments and data frame values.
+    Dimension attributes are merged with attributes in record objects.
+    Example: common_attributes = {"Dimensions": {"Name": "device_id", "Value": "12345"}, "MeasureValueType": "DOUBLE"}.
+
+    Note
+    ----
+    If the `time_col` column is supplied it must be of type timestamp. `TimeUnit` is set to MILLISECONDS by default.
+
     Parameters
     ----------
     df: pandas.DataFrame
@@ -214,11 +276,11 @@ def write(
         Amazon Timestream database name.
     table : str
         Amazon Timestream table name.
-    time_col : str
+    time_col : Optional[str]
         DataFrame column name to be used as time. MUST be a timestamp column.
-    measure_col : Union[str, List[str]]
+    measure_col : Union[str, List[str], None]
         DataFrame column name(s) to be used as measure.
-    dimensions_cols : List[str]
+    dimensions_cols : Optional[List[str]]
         List of DataFrame column names to be used as dimensions.
     version : int
         Version number used for upserts.
@@ -226,6 +288,9 @@ def write(
     measure_name : Optional[str]
         Name that represents the data attribute of the time series.
         Overrides ``measure_col`` if specified.
+    common_attributes : Optional[Dict[str, Any]]
+        Dictionary of attributes that is shared across all records in the request.
+        Using common attributes can optimize the cost of writes by reducing the size of request payloads.
     num_threads : str
         Number of thread to be used for concurrent writing.
     boto3_session : boto3.Session(), optional
@@ -279,30 +344,50 @@ def write(
         session=boto3_session,
         botocore_config=Config(read_timeout=20, max_pool_connections=5000, retries={"max_attempts": 10}),
     )
-    measure_cols_names = measure_col if isinstance(measure_col, list) else [measure_col]
-    _logger.debug("measure_cols_names: %s", measure_cols_names)
 
-    measure_types: List[str] = [
-        _data_types.timestream_type_from_pandas(df[[measure_col_name]]) for measure_col_name in measure_cols_names
+    measure_cols = measure_col if isinstance(measure_col, list) else [measure_col]
+    measure_types = [
+        _data_types.timestream_type_from_pandas(df[[measure_col_name]])
+        for measure_col_name in measure_cols
+        if measure_col_name
     ]
-    _logger.debug("measure_types: %s", measure_types)
-    cols_names: List[str] = [time_col] + measure_cols_names + dimensions_cols
-    _logger.debug("cols_names: %s", cols_names)
-    batches: List[List[Any]] = _utils.chunkify(lst=_df2list(df=df[cols_names]), max_length=100)
-    _logger.debug("len(batches): %s", len(batches))
+    dimensions_cols = dimensions_cols if dimensions_cols else [dimensions_cols]  # type: ignore[list-item]
+    cols_names: List[Optional[str]] = [time_col] + measure_cols + dimensions_cols
+    measure_name = measure_name if measure_name else measure_cols[0]
+    common_attributes = _sanitize_common_attributes(common_attributes, version, measure_name)
+
+    _logger.debug(
+        "common_attributes: %s\n, cols_names: %s\n, measure_types: %s",
+        common_attributes,
+        cols_names,
+        measure_types,
+    )
+
+    # User can supply arguments in one of two ways:
+    # 1. With the `common_attributes` dictionary which takes precedence
+    # 2. With data frame columns
+    # However, the data frame cannot be completely empty.
+    # So if all values in `cols_names` are None, an exception is raised.
+    if any(cols_names):
+        batches: List[List[Any]] = _utils.chunkify(lst=_df2list(df=df[[c for c in cols_names if c]]), max_length=100)
+    else:
+        raise exceptions.InvalidArgumentCombination(
+            "At least one of `time_col`, `measure_col` or `dimensions_cols` must be specified."
+        )
+
     with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor:
         res: List[List[Any]] = list(
             executor.map(
                 _write_batch,
+                itertools.repeat(timestream_client),
                 itertools.repeat(database),
                 itertools.repeat(table),
+                itertools.repeat(common_attributes),
                 itertools.repeat(cols_names),
-                itertools.repeat(measure_cols_names),
+                itertools.repeat(measure_cols),
                 itertools.repeat(measure_types),
-                itertools.repeat(version),
+                itertools.repeat(dimensions_cols),
                 batches,
-                itertools.repeat(timestream_client),
-                itertools.repeat(measure_name),
             )
         )
         return [item for sublist in res for item in sublist]


Test patch:
diff --git a/tests/test_timestream.py b/tests/test_timestream.py
--- a/tests/test_timestream.py
+++ b/tests/test_timestream.py
@@ -1,5 +1,6 @@
 import datetime as dt
 import logging
+import time
 from datetime import datetime
 
 import boto3
@@ -220,6 +221,162 @@ def test_real_csv_load_scenario(timestream_database_and_table):
     assert df["counter"].iloc[0] == 126_000
 
 
+# This test covers every possible combination of common_attributes and data frame inputs
+@pytest.mark.parametrize("record_type", ["SCALAR", "MULTI"])
+@pytest.mark.parametrize(
+    "common_attributes_details",
+    [
+        ({}, (4, 6)),
+        ({"MeasureName": "cpu_util"}, (4, 6)),
+        ({"MeasureValue": "13.1", "MeasureValueType": "DOUBLE"}, (4, 5)),
+        (
+            {
+                "MeasureValues": [
+                    {"Name": "cpu_util", "Value": "12.0", "Type": "DOUBLE"},
+                    {"Name": "mem_util", "Value": "45.6", "Type": "DOUBLE"},
+                ],
+                "MeasureValueType": "MULTI",
+            },
+            (5, 6),
+        ),
+        ({"Time": str(round(time.time())), "TimeUnit": "SECONDS"}, (4, 5)),
+        ({"Time": str(round(time.time()) * 1_000), "MeasureValue": "13.1", "MeasureValueType": "DOUBLE"}, (4, 5)),
+        ({"Dimensions": [{"Name": "region", "Value": "us-east-1"}]}, (4, 5)),
+        ({"Dimensions": [{"Name": "region", "Value": "us-east-1"}, {"Name": "host", "Value": "linux"}]}, (5, 6)),
+        (
+            {
+                "Dimensions": [{"Name": "region", "Value": "us-east-1"}],
+                "MeasureValue": "13.1",
+                "MeasureValueType": "DOUBLE",
+            },
+            (4, 4),
+        ),
+        (
+            {
+                "Dimensions": [{"Name": "region", "Value": "us-east-1"}, {"Name": "host", "Value": "linux"}],
+                "MeasureValues": [
+                    {"Name": "cpu_util", "Value": "12.0", "Type": "DOUBLE"},
+                    {"Name": "mem_util", "Value": "45.6", "Type": "DOUBLE"},
+                ],
+                "MeasureValueType": "MULTI",
+            },
+            (6, 6),
+        ),
+    ],
+)
+def test_common_attributes(timestream_database_and_table, record_type, common_attributes_details):
+    common_attributes, num_cols = common_attributes_details
+    df = pd.DataFrame({"dummy": ["a"] * 3})
+
+    kwargs = {
+        "df": df,
+        "database": timestream_database_and_table,
+        "table": timestream_database_and_table,
+        "common_attributes": common_attributes,
+        "measure_name": "cpu",
+    }
+    if "Time" not in common_attributes:
+        df["time"] = [datetime.now() + dt.timedelta(seconds=(60 * c)) for c in range(3)]
+        kwargs["time_col"] = "time"
+    if all(k not in common_attributes for k in ("MeasureValue", "MeasureValues")):
+        df["cpu_utilization"] = [13.1] * 3
+        measure_cols = ["cpu_utilization"]
+        if record_type == "MULTI":
+            df["mem_utilization"] = [45.2] * 3
+            measure_cols += ["mem_utilization"]
+        kwargs["measure_col"] = measure_cols
+    if "Dimensions" not in common_attributes:
+        df["region"] = ["us-east-1", "us-east-2", "us-west-2"]
+        dimensions_cols = ["region"]
+        if record_type == "MULTI":
+            df["host"] = ["AL2", "Ubuntu", "Debian"]
+            dimensions_cols += ["host"]
+        kwargs["dimensions_cols"] = dimensions_cols
+
+    rejected_records = wr.timestream.write(**kwargs)
+    assert len(rejected_records) == 0
+
+    df = wr.timestream.query(f"""SELECT * FROM "{timestream_database_and_table}"."{timestream_database_and_table}" """)
+    assert df.shape == (3, num_cols[1] if record_type == "MULTI" else num_cols[0])
+
+
+@pytest.mark.parametrize("record_type", ["SCALAR", "MULTI"])
+def test_merge_dimensions(timestream_database_and_table, record_type):
+    df = pd.DataFrame(
+        {
+            "time": [datetime.now()] * 3,
+            "measure": [13.1, 13.2, 13.3],
+            "dim1": ["az1", "az2", "az3"],
+        }
+    )
+
+    kwargs = {
+        "df": df,
+        "database": timestream_database_and_table,
+        "table": timestream_database_and_table,
+        "common_attributes": {
+            "Dimensions": [{"Name": "region", "Value": "us-east-1"}, {"Name": "host", "Value": "linux"}]
+        },
+        "time_col": "time",
+        "measure_col": ["measure"],
+        "dimensions_cols": ["dim1"],
+    }
+    if record_type == "MULTI":
+        df["dim2"] = ["arm", "intel", "arm"]
+        kwargs["dimensions_cols"] = ["dim1", "dim2"]
+
+    rejected_records = wr.timestream.write(**kwargs)
+    assert len(rejected_records) == 0
+
+    df = wr.timestream.query(f"""SELECT * FROM "{timestream_database_and_table}"."{timestream_database_and_table}" """)
+    assert df.shape == (3, 7) if record_type == "MULTI" else (3, 6)
+
+
+def test_exceptions():
+    df = pd.DataFrame(
+        {
+            "time": [datetime.now(), datetime.now(), datetime.now()],
+            "measure1": [1.0, 1.1, 1.2],
+        }
+    )
+
+    # No columns supplied
+    with pytest.raises(wr.exceptions.InvalidArgumentCombination):
+        wr.timestream.write(
+            df=df,
+            database="test",
+            table="test",
+            common_attributes={"MeasureName": "test", "MeasureValueType": "Double"},
+        )
+
+    # Missing MeasureValueType
+    with pytest.raises(wr.exceptions.InvalidArgumentCombination):
+        wr.timestream.write(
+            df=df, database="test", table="test", time_col="time", common_attributes={"MeasureValue": "13.1"}
+        )
+
+    # Missing MeasureName
+    with pytest.raises(wr.exceptions.InvalidArgumentCombination):
+        wr.timestream.write(
+            df=df,
+            database="test",
+            table="test",
+            time_col="time",
+            common_attributes={"MeasureValue": "13.1", "MeasureValueType": "DOUBLE"},
+        )
+
+    # None timestamp time_col
+    with pytest.raises(wr.exceptions.InvalidArgumentType):
+        df["time"] = ["a"] * 3
+        wr.timestream.write(
+            df=df,
+            database="test",
+            table="test",
+            time_col="time",
+            common_attributes={"MeasureName": "test", "MeasureValue": "13.1", "MeasureValueType": "Double"},
+        )
+
+
 def test_multimeasure_scenario(timestream_database_and_table):
     df = pd.DataFrame(
         {
@@ -252,7 +409,7 @@ def test_multimeasure_scenario(timestream_database_and_table):
 
 
 @pytest.mark.parametrize(
-    "test_input,expected",
+    "test_input, expected",
     [
         (
             {