arthur-common 1.0.1__py3-none-any.whl → 2.1.48__py3-none-any.whl

arthur_common/aggregations/functions/inference_null_count.py

@@ -1,4 +1,4 @@
-from typing import Annotated
+from typing import Annotated, Optional
 from uuid import UUID
 
 from arthur_common.aggregations.aggregator import NumericAggregationFunction
@@ -7,6 +7,7 @@ from arthur_common.models.schema_definitions import (
     DType,
     MetricColumnParameterAnnotation,
     MetricDatasetParameterAnnotation,
+    MetricMultipleColumnParameterAnnotation,
     ScalarType,
     ScopeSchemaTag,
 )
@@ -60,23 +61,62 @@ class InferenceNullCountAggregationFunction(NumericAggregationFunction):
                 description="A column containing nullable values to count.",
             ),
         ],
+        segmentation_cols: Annotated[
+            Optional[list[str]],
+            MetricMultipleColumnParameterAnnotation(
+                source_dataset_parameter_key="dataset",
+                allowed_column_types=[
+                    ScalarType(dtype=DType.INT),
+                    ScalarType(dtype=DType.BOOL),
+                    ScalarType(dtype=DType.STRING),
+                    ScalarType(dtype=DType.UUID),
+                ],
+                tag_hints=[],
+                friendly_name="Segmentation Columns",
+                description="All columns to include as dimensions for segmentation.",
+                optional=True,
+            ),
+        ] = None,
     ) -> list[NumericMetric]:
+        """Executed SQL with no segmentation columns:
+                select time_bucket(INTERVAL '5 minutes', {escaped_timestamp_col}) as ts, \
+                count(*) as count \
+                from {dataset.dataset_table_name} where {escaped_nullable_col} is null \
+                group by ts \
+        """
+        segmentation_cols = [] if not segmentation_cols else segmentation_cols
         escaped_timestamp_col = escape_identifier(timestamp_col)
         escaped_nullable_col = escape_identifier(nullable_col)
-        count_query = f" \
-            select time_bucket(INTERVAL '5 minutes', {escaped_timestamp_col}) as ts, \
-            count(*) as count \
-            from {dataset.dataset_table_name} where {escaped_nullable_col} is null \
-            group by ts \
-        "
+
+        # build query components with segmentation columns
+        escaped_segmentation_cols = [
+            escape_identifier(col) for col in segmentation_cols
+        ]
+        all_select_clause_cols = [
+            f"time_bucket(INTERVAL '5 minutes', {escaped_timestamp_col}) as ts",
+            f"count(*) as count",
+        ] + escaped_segmentation_cols
+        all_group_by_cols = ["ts"] + escaped_segmentation_cols
+
+        # build query
+        count_query = f"""
+            select {", ".join(all_select_clause_cols)}
+            from {dataset.dataset_table_name}
+            where {escaped_nullable_col} is null
+            group by {", ".join(all_group_by_cols)}
+        """
+
         results = ddb_conn.sql(count_query).df()
 
-        series = self.dimensionless_query_results_to_numeric_metrics(
+        series = self.group_query_results_to_numeric_metrics(
             results,
             "count",
+            segmentation_cols,
             "ts",
         )
-        series.dimensions = [Dimension(name="column_name", value=nullable_col)]
+        # preserve dimension that identifies the name of the nullable column used for the aggregation
+        for point in series:
+            point.dimensions.append(Dimension(name="column_name", value=nullable_col))
 
-        metric = self.series_to_metric(self.METRIC_NAME, [series])
+        metric = self.series_to_metric(self.METRIC_NAME, series)
         return [metric]

arthur_common/aggregations/functions/mean_absolute_error.py

@@ -1,4 +1,4 @@
-from typing import Annotated
+from typing import Annotated, Optional
 from uuid import UUID
 
 from arthur_common.aggregations.aggregator import NumericAggregationFunction
@@ -8,6 +8,7 @@ from arthur_common.models.schema_definitions import (
     DType,
     MetricColumnParameterAnnotation,
     MetricDatasetParameterAnnotation,
+    MetricMultipleColumnParameterAnnotation,
     ScalarType,
     ScopeSchemaTag,
 )
@@ -75,36 +76,75 @@ class MeanAbsoluteErrorAggregationFunction(NumericAggregationFunction):
                 description="A column containing float typed ground truth values.",
             ),
         ],
+        segmentation_cols: Annotated[
+            Optional[list[str]],
+            MetricMultipleColumnParameterAnnotation(
+                source_dataset_parameter_key="dataset",
+                allowed_column_types=[
+                    ScalarType(dtype=DType.INT),
+                    ScalarType(dtype=DType.BOOL),
+                    ScalarType(dtype=DType.STRING),
+                    ScalarType(dtype=DType.UUID),
+                ],
+                tag_hints=[],
+                friendly_name="Segmentation Columns",
+                description="All columns to include as dimensions for segmentation.",
+                optional=True,
+            ),
+        ] = None,
     ) -> list[NumericMetric]:
+        """Executed SQL with no segmentation columns:
+                SELECT time_bucket(INTERVAL '5 minutes', {escaped_timestamp_col}) as ts, \
+                SUM(ABS({escaped_prediction_col} - {escaped_ground_truth_col})) as ae, \
+                COUNT(*) as count \
+                FROM {dataset.dataset_table_name} \
+                WHERE {escaped_prediction_col} IS NOT NULL \
+                AND {escaped_ground_truth_col} IS NOT NULL \
+                GROUP BY ts order by ts desc \
+                """
+        segmentation_cols = [] if not segmentation_cols else segmentation_cols
         escaped_timestamp_col = escape_identifier(timestamp_col)
         escaped_prediction_col = escape_identifier(prediction_col)
         escaped_ground_truth_col = escape_identifier(ground_truth_col)
-        count_query = f" \
-            SELECT time_bucket(INTERVAL '5 minutes', {escaped_timestamp_col}) as ts, \
-            SUM(ABS({escaped_prediction_col} - {escaped_ground_truth_col})) as ae, \
-            COUNT(*) as count \
-            FROM {dataset.dataset_table_name} \
-            WHERE {escaped_prediction_col} IS NOT NULL \
-            AND {escaped_ground_truth_col} IS NOT NULL \
-            GROUP BY ts order by ts desc \
-        "
 
-        results = ddb_conn.sql(count_query).df()
-        count_series = self.dimensionless_query_results_to_numeric_metrics(
+        # build query components with segmentation columns
+        escaped_segmentation_cols = [
+            escape_identifier(col) for col in segmentation_cols
+        ]
+        all_select_clause_cols = [
+            f"time_bucket(INTERVAL '5 minutes', {escaped_timestamp_col}) as ts",
+            f"SUM(ABS({escaped_prediction_col} - {escaped_ground_truth_col})) as ae",
+            f"COUNT(*) as count",
+        ] + escaped_segmentation_cols
+        all_group_by_cols = ["ts"] + escaped_segmentation_cols
+
+        # build query
+        mae_query = f"""
+            SELECT {", ".join(all_select_clause_cols)}
+            FROM {dataset.dataset_table_name}
+            WHERE {escaped_prediction_col} IS NOT NULL
+                  AND {escaped_ground_truth_col} IS NOT NULL
+            GROUP BY {", ".join(all_group_by_cols)} order by ts desc
+        """
+
+        results = ddb_conn.sql(mae_query).df()
+        count_series = self.group_query_results_to_numeric_metrics(
             results,
             "count",
+            segmentation_cols,
             "ts",
         )
-        absolute_error_series = self.dimensionless_query_results_to_numeric_metrics(
+        absolute_error_series = self.group_query_results_to_numeric_metrics(
             results,
             "ae",
+            segmentation_cols,
             "ts",
         )
 
-        count_metric = self.series_to_metric("absolute_error_count", [count_series])
+        count_metric = self.series_to_metric("absolute_error_count", count_series)
         absolute_error_metric = self.series_to_metric(
             "absolute_error_sum",
-            [absolute_error_series],
+            absolute_error_series,
         )
 
         return [count_metric, absolute_error_metric]

arthur_common/aggregations/functions/mean_squared_error.py

@@ -1,4 +1,4 @@
-from typing import Annotated
+from typing import Annotated, Optional
 from uuid import UUID
 
 from arthur_common.aggregations.aggregator import NumericAggregationFunction
@@ -8,6 +8,7 @@ from arthur_common.models.schema_definitions import (
     DType,
     MetricColumnParameterAnnotation,
     MetricDatasetParameterAnnotation,
+    MetricMultipleColumnParameterAnnotation,
     ScalarType,
     ScopeSchemaTag,
 )
@@ -75,36 +76,75 @@ class MeanSquaredErrorAggregationFunction(NumericAggregationFunction):
                 description="A column containing float typed ground truth values.",
             ),
         ],
+        segmentation_cols: Annotated[
+            Optional[list[str]],
+            MetricMultipleColumnParameterAnnotation(
+                source_dataset_parameter_key="dataset",
+                allowed_column_types=[
+                    ScalarType(dtype=DType.INT),
+                    ScalarType(dtype=DType.BOOL),
+                    ScalarType(dtype=DType.STRING),
+                    ScalarType(dtype=DType.UUID),
+                ],
+                tag_hints=[],
+                friendly_name="Segmentation Columns",
+                description="All columns to include as dimensions for segmentation.",
+                optional=True,
+            ),
+        ] = None,
     ) -> list[NumericMetric]:
+        """Executed SQL with no segmentation columns:
+                SELECT time_bucket(INTERVAL '5 minutes', {escaped_timestamp_col}) as ts, \
+                SUM(POW({escaped_prediction_col} - {escaped_ground_truth_col}, 2)) as squared_error, \
+                COUNT(*) as count \
+                FROM {dataset.dataset_table_name} \
+                WHERE {escaped_prediction_col} IS NOT NULL \
+                AND {escaped_ground_truth_col} IS NOT NULL \
+                GROUP BY ts order by ts desc \
+                """
+        segmentation_cols = [] if not segmentation_cols else segmentation_cols
         escaped_timestamp_col = escape_identifier(timestamp_col)
         escaped_prediction_col = escape_identifier(prediction_col)
         escaped_ground_truth_col = escape_identifier(ground_truth_col)
-        count_query = f" \
-            SELECT time_bucket(INTERVAL '5 minutes', {escaped_timestamp_col}) as ts, \
-            SUM(POW({escaped_prediction_col} - {escaped_ground_truth_col}, 2)) as squared_error, \
-            COUNT(*) as count \
-            FROM {dataset.dataset_table_name} \
-            WHERE {escaped_prediction_col} IS NOT NULL \
-            AND {escaped_ground_truth_col} IS NOT NULL \
-            GROUP BY ts order by ts desc \
-        "
 
-        results = ddb_conn.sql(count_query).df()
-        count_series = self.dimensionless_query_results_to_numeric_metrics(
+        # build query components with segmentation columns
+        escaped_segmentation_cols = [
+            escape_identifier(col) for col in segmentation_cols
+        ]
+        all_select_clause_cols = [
+            f"time_bucket(INTERVAL '5 minutes', {escaped_timestamp_col}) as ts",
+            f"SUM(POW({escaped_prediction_col} - {escaped_ground_truth_col}, 2)) as squared_error",
+            f"COUNT(*) as count",
+        ] + escaped_segmentation_cols
+        all_group_by_cols = ["ts"] + escaped_segmentation_cols
+
+        # build query
+        mse_query = f"""
+            SELECT {", ".join(all_select_clause_cols)}
+            FROM {dataset.dataset_table_name}
+            WHERE {escaped_prediction_col} IS NOT NULL
+                  AND {escaped_ground_truth_col} IS NOT NULL
+            GROUP BY {", ".join(all_group_by_cols)} order by ts desc
+        """
+
+        results = ddb_conn.sql(mse_query).df()
+        count_series = self.group_query_results_to_numeric_metrics(
             results,
             "count",
+            segmentation_cols,
             "ts",
         )
-        squared_error_series = self.dimensionless_query_results_to_numeric_metrics(
+        squared_error_series = self.group_query_results_to_numeric_metrics(
             results,
             "squared_error",
+            segmentation_cols,
             "ts",
         )
 
-        count_metric = self.series_to_metric("squared_error_count", [count_series])
+        count_metric = self.series_to_metric("squared_error_count", count_series)
         absolute_error_metric = self.series_to_metric(
             "squared_error_sum",
-            [squared_error_series],
+            squared_error_series,
         )
 
         return [count_metric, absolute_error_metric]

arthur_common/aggregations/functions/multiclass_confusion_matrix.py

@@ -1,4 +1,4 @@
-from typing import Annotated
+from typing import Annotated, Optional
 from uuid import UUID
 
 from arthur_common.aggregations.aggregator import NumericAggregationFunction
@@ -9,6 +9,7 @@ from arthur_common.models.schema_definitions import (
     MetricColumnParameterAnnotation,
     MetricDatasetParameterAnnotation,
     MetricLiteralParameterAnnotation,
+    MetricMultipleColumnParameterAnnotation,
     ScalarType,
     ScopeSchemaTag,
 )
@@ -90,7 +91,24 @@ class MulticlassClassifierStringLabelSingleClassConfusionMatrixAggregationFuncti
                 description="The label indicating a positive class.",
             ),
         ],
+        segmentation_cols: Annotated[
+            Optional[list[str]],
+            MetricMultipleColumnParameterAnnotation(
+                source_dataset_parameter_key="dataset",
+                allowed_column_types=[
+                    ScalarType(dtype=DType.INT),
+                    ScalarType(dtype=DType.BOOL),
+                    ScalarType(dtype=DType.STRING),
+                    ScalarType(dtype=DType.UUID),
+                ],
+                tag_hints=[],
+                friendly_name="Segmentation Columns",
+                description="All columns to include as dimensions for segmentation.",
+                optional=True,
+            ),
+        ] = None,
     ) -> list[NumericMetric]:
+        segmentation_cols = [] if not segmentation_cols else segmentation_cols
         escaped_positive_class_label = escape_str_literal(positive_class_label)
         normalization_case = f"""
                 CASE
@@ -107,6 +125,7 @@ class MulticlassClassifierStringLabelSingleClassConfusionMatrixAggregationFuncti
             normalization_case,
             dataset,
             escaped_positive_class_label,
+            segmentation_cols,
         )
 
     def generate_confusion_matrix_metrics(
@@ -119,6 +138,7 @@ class MulticlassClassifierStringLabelSingleClassConfusionMatrixAggregationFuncti
         gt_normalization_case: str,
         dataset: DatasetReference,
         escaped_positive_class_label: str,
+        segmentation_cols: list[str],
     ) -> list[NumericMetric]:
         """
         Generate a SQL query to compute confusion matrix metrics over time.
@@ -132,58 +152,92 @@ class MulticlassClassifierStringLabelSingleClassConfusionMatrixAggregationFuncti
             gt_normalization_case: SQL CASE statement for normalizing ground truth values to 0 / 1 / null using 'value' as the target column name
             dataset: DatasetReference containing dataset metadata
             escaped_positive_class_label: escaped label for the class to include in the dimensions
+            segmentation_cols: List of columns to segment by
 
         Returns:
             str: SQL query that computes confusion matrix metrics
+            Returns the following SQL with no segmentation:
+            WITH normalized_data AS (
+                    SELECT
+                        {escaped_timestamp_col} AS timestamp,
+                        {prediction_normalization_case.replace('value', escaped_prediction_col)} AS prediction,
+                        {gt_normalization_case.replace('value', escaped_gt_values_col)} AS actual_value
+                    FROM {dataset.dataset_table_name}
+                    WHERE {escaped_timestamp_col} IS NOT NULL
+                )
+                SELECT
+                    time_bucket(INTERVAL '5 minutes', timestamp) AS ts,
+                    SUM(CASE WHEN prediction = 1 AND actual_value = 1 THEN 1 ELSE 0 END) AS true_positive_count,
+                    SUM(CASE WHEN prediction = 1 AND actual_value = 0 THEN 1 ELSE 0 END) AS false_positive_count,
+                    SUM(CASE WHEN prediction = 0 AND actual_value = 1 THEN 1 ELSE 0 END) AS false_negative_count,
+                    SUM(CASE WHEN prediction = 0 AND actual_value = 0 THEN 1 ELSE 0 END) AS true_negative_count,
+                    any_value({escaped_positive_class_label}) as class_label
+                FROM normalized_data
+                GROUP BY ts
+                ORDER BY ts
+
         """
         escaped_timestamp_col = escape_identifier(timestamp_col)
         escaped_prediction_col = escape_identifier(prediction_col)
         escaped_gt_values_col = escape_identifier(gt_values_col)
+
+        # build query components with segmentation columns
+        escaped_segmentation_cols = [
+            escape_identifier(col) for col in segmentation_cols
+        ]
+        first_subquery_select_cols = [
+            f"{escaped_timestamp_col} AS timestamp",
+            f"{prediction_normalization_case.replace('value', escaped_prediction_col)} AS prediction",
+            f"{gt_normalization_case.replace('value', escaped_gt_values_col)} AS actual_value",
+        ] + escaped_segmentation_cols
+        second_subquery_select_cols = [
+            "time_bucket(INTERVAL '5 minutes', timestamp) AS ts",
+            "SUM(CASE WHEN prediction = 1 AND actual_value = 1 THEN 1 ELSE 0 END) AS true_positive_count",
+            "SUM(CASE WHEN prediction = 1 AND actual_value = 0 THEN 1 ELSE 0 END) AS false_positive_count",
+            "SUM(CASE WHEN prediction = 0 AND actual_value = 1 THEN 1 ELSE 0 END) AS false_negative_count",
+            "SUM(CASE WHEN prediction = 0 AND actual_value = 0 THEN 1 ELSE 0 END) AS true_negative_count",
+            f"any_value({escaped_positive_class_label}) as class_label",
+        ] + escaped_segmentation_cols
+        second_subquery_group_by_cols = ["ts"] + escaped_segmentation_cols
+        extra_dims = ["class_label"]
+
+        # build query
         confusion_matrix_query = f"""
-            WITH normalized_data AS (
-                SELECT
-                    {escaped_timestamp_col} AS timestamp,
-                    {prediction_normalization_case.replace('value', escaped_prediction_col)} AS prediction,
-                    {gt_normalization_case.replace('value', escaped_gt_values_col)} AS actual_value
-                FROM {dataset.dataset_table_name}
-                WHERE {escaped_timestamp_col} IS NOT NULL
-            )
-            SELECT
-                time_bucket(INTERVAL '5 minutes', timestamp) AS ts,
-                SUM(CASE WHEN prediction = 1 AND actual_value = 1 THEN 1 ELSE 0 END) AS true_positive_count,
-                SUM(CASE WHEN prediction = 1 AND actual_value = 0 THEN 1 ELSE 0 END) AS false_positive_count,
-                SUM(CASE WHEN prediction = 0 AND actual_value = 1 THEN 1 ELSE 0 END) AS false_negative_count,
-                SUM(CASE WHEN prediction = 0 AND actual_value = 0 THEN 1 ELSE 0 END) AS true_negative_count,
-                any_value({escaped_positive_class_label}) as class_label
-            FROM normalized_data
-            GROUP BY ts
-            ORDER BY ts
-        """
+        WITH normalized_data AS (
+            SELECT {", ".join(first_subquery_select_cols)}
+            FROM {dataset.dataset_table_name}
+            WHERE {escaped_timestamp_col} IS NOT NULL
+        )
+        SELECT {", ".join(second_subquery_select_cols)}
+        FROM normalized_data
+        GROUP BY {", ".join(second_subquery_group_by_cols)}
+        ORDER BY ts
+"""
 
         results = ddb_conn.sql(confusion_matrix_query).df()
 
         tp = self.group_query_results_to_numeric_metrics(
             results,
             "true_positive_count",
-            dim_columns=["class_label"],
+            dim_columns=segmentation_cols + extra_dims,
             timestamp_col="ts",
         )
         fp = self.group_query_results_to_numeric_metrics(
             results,
             "false_positive_count",
-            dim_columns=["class_label"],
+            dim_columns=segmentation_cols + extra_dims,
             timestamp_col="ts",
         )
         fn = self.group_query_results_to_numeric_metrics(
             results,
             "false_negative_count",
-            dim_columns=["class_label"],
+            dim_columns=segmentation_cols + extra_dims,
             timestamp_col="ts",
         )
         tn = self.group_query_results_to_numeric_metrics(
             results,
             "true_negative_count",
-            dim_columns=["class_label"],
+            dim_columns=segmentation_cols + extra_dims,
             timestamp_col="ts",
         )
         tp_metric = self.series_to_metric(

arthur_common/aggregations/functions/multiclass_inference_count_by_class.py

@@ -1,4 +1,4 @@
-from typing import Annotated
+from typing import Annotated, Optional
 from uuid import UUID
 
 from arthur_common.aggregations.functions.inference_count_by_class import (
@@ -10,6 +10,7 @@ from arthur_common.models.schema_definitions import (
     DType,
     MetricColumnParameterAnnotation,
     MetricDatasetParameterAnnotation,
+    MetricMultipleColumnParameterAnnotation,
     ScalarType,
     ScopeSchemaTag,
 )
@@ -81,10 +82,27 @@ class MulticlassClassifierCountByClassAggregationFunction(
                 description="A column containing boolean, integer, or string labelled prediction values.",
             ),
         ],
+        segmentation_cols: Annotated[
+            Optional[list[str]],
+            MetricMultipleColumnParameterAnnotation(
+                source_dataset_parameter_key="dataset",
+                allowed_column_types=[
+                    ScalarType(dtype=DType.INT),
+                    ScalarType(dtype=DType.BOOL),
+                    ScalarType(dtype=DType.STRING),
+                    ScalarType(dtype=DType.UUID),
+                ],
+                tag_hints=[],
+                friendly_name="Segmentation Columns",
+                description="All columns to include as dimensions for segmentation.",
+                optional=True,
+            ),
+        ] = None,
     ) -> list[NumericMetric]:
         return super().aggregate(
             ddb_conn=ddb_conn,
             dataset=dataset,
             timestamp_col=timestamp_col,
             prediction_col=prediction_col,
+            segmentation_cols=segmentation_cols,
         )

arthur_common/aggregations/functions/numeric_stats.py

@@ -1,4 +1,4 @@
-from typing import Annotated
+from typing import Annotated, Optional
 from uuid import UUID
 
 from arthur_common.aggregations.aggregator import SketchAggregationFunction
@@ -7,6 +7,7 @@ from arthur_common.models.schema_definitions import (
     DType,
     MetricColumnParameterAnnotation,
     MetricDatasetParameterAnnotation,
+    MetricMultipleColumnParameterAnnotation,
     ScalarType,
     ScopeSchemaTag,
 )
@@ -66,23 +67,59 @@ class NumericSketchAggregationFunction(SketchAggregationFunction):
                 description="A column containing numeric values to calculate a data sketch on.",
             ),
         ],
+        segmentation_cols: Annotated[
+            Optional[list[str]],
+            MetricMultipleColumnParameterAnnotation(
+                source_dataset_parameter_key="dataset",
+                allowed_column_types=[
+                    ScalarType(dtype=DType.INT),
+                    ScalarType(dtype=DType.BOOL),
+                    ScalarType(dtype=DType.STRING),
+                    ScalarType(dtype=DType.UUID),
+                ],
+                tag_hints=[],
+                friendly_name="Segmentation Columns",
+                description="All columns to include as dimensions for segmentation.",
+                optional=True,
+            ),
+        ] = None,
     ) -> list[SketchMetric]:
+        """Executed SQL with no segmentation columns:
+                    select {escaped_timestamp_col_id} as ts, \
+                       {escaped_numeric_col_id}, \
+                       {numeric_col_name_str} as column_name \
+                from {dataset.dataset_table_name} \
+                where {escaped_numeric_col_id} is not null \
+        """
+        segmentation_cols = [] if not segmentation_cols else segmentation_cols
         escaped_timestamp_col_id = escape_identifier(timestamp_col)
         escaped_numeric_col_id = escape_identifier(numeric_col)
         numeric_col_name_str = escape_str_literal(numeric_col)
-        data_query = f" \
-            select {escaped_timestamp_col_id} as ts, \
-                   {escaped_numeric_col_id}, \
-                   {numeric_col_name_str} as column_name \
-            from {dataset.dataset_table_name} \
-            where {escaped_numeric_col_id} is not null \
-        "
+
+        # build query components with segmentation columns
+        escaped_segmentation_cols = [
+            escape_identifier(col) for col in segmentation_cols
+        ]
+        all_select_clause_cols = [
+            f"{escaped_timestamp_col_id} as ts",
+            f"{escaped_numeric_col_id}",
+            f"{numeric_col_name_str} as column_name",
+        ] + escaped_segmentation_cols
+        extra_dims = ["column_name"]
+
+        # build query
+        data_query = f"""
+                    select {", ".join(all_select_clause_cols)}
+                    from {dataset.dataset_table_name}
+                    where {escaped_numeric_col_id} is not null
+                """
+
         results = ddb_conn.sql(data_query).df()
 
         series = self.group_query_results_to_sketch_metrics(
             results,
             numeric_col,
-            ["column_name"],
+            segmentation_cols + extra_dims,
             "ts",
         )
 

arthur_common/aggregations/functions/numeric_sum.py

@@ -1,4 +1,4 @@
-from typing import Annotated
+from typing import Annotated, Optional
 from uuid import UUID
 
 from arthur_common.aggregations.aggregator import NumericAggregationFunction
@@ -7,6 +7,7 @@ from arthur_common.models.schema_definitions import (
     DType,
     MetricColumnParameterAnnotation,
     MetricDatasetParameterAnnotation,
+    MetricMultipleColumnParameterAnnotation,
     ScalarType,
     ScopeSchemaTag,
 )
@@ -64,24 +65,63 @@ class NumericSumAggregationFunction(NumericAggregationFunction):
                 description="A column containing numeric values to sum.",
             ),
         ],
+        segmentation_cols: Annotated[
+            Optional[list[str]],
+            MetricMultipleColumnParameterAnnotation(
+                source_dataset_parameter_key="dataset",
+                allowed_column_types=[
+                    ScalarType(dtype=DType.INT),
+                    ScalarType(dtype=DType.BOOL),
+                    ScalarType(dtype=DType.STRING),
+                    ScalarType(dtype=DType.UUID),
+                ],
+                tag_hints=[],
+                friendly_name="Segmentation Columns",
+                description="All columns to include as dimensions for segmentation.",
+                optional=True,
+            ),
+        ] = None,
     ) -> list[NumericMetric]:
+        """Executed SQL with no segmentation columns:
+                select time_bucket(INTERVAL '5 minutes', {escaped_timestamp_col}) as ts, \
+                sum({escaped_numeric_col}) as sum \
+                from {dataset.dataset_table_name} \
+                where {escaped_numeric_col} is not null \
+                group by ts \
+        """
+        segmentation_cols = [] if not segmentation_cols else segmentation_cols
         escaped_timestamp_col = escape_identifier(timestamp_col)
         escaped_numeric_col = escape_identifier(numeric_col)
-        count_query = f" \
-            select time_bucket(INTERVAL '5 minutes', {escaped_timestamp_col}) as ts, \
-            sum({escaped_numeric_col}) as sum \
-            from {dataset.dataset_table_name} \
-            where {escaped_numeric_col} is not null \
-            group by ts \
-        "
-        results = ddb_conn.sql(count_query).df()
 
-        series = self.dimensionless_query_results_to_numeric_metrics(
+        # build query components with segmentation columns
+        escaped_segmentation_cols = [
+            escape_identifier(col) for col in segmentation_cols
+        ]
+        all_select_clause_cols = [
+            f"time_bucket(INTERVAL '5 minutes', {escaped_timestamp_col}) as ts",
+            f"sum({escaped_numeric_col}) as sum",
+        ] + escaped_segmentation_cols
+        all_group_by_cols = ["ts"] + escaped_segmentation_cols
+
+        # build query
+        query = f"""
+                    select {", ".join(all_select_clause_cols)}
+                    from {dataset.dataset_table_name}
+                    where {escaped_numeric_col} is not null
+                    group by {", ".join(all_group_by_cols)}
+                """
+
+        results = ddb_conn.sql(query).df()
+
+        series = self.group_query_results_to_numeric_metrics(
             results,
             "sum",
+            segmentation_cols,
             "ts",
         )
-        series.dimensions = [Dimension(name="column_name", value=numeric_col)]
+        # preserve dimension that identifies the name of the numeric column used for the aggregation
+        for point in series:
+            point.dimensions.append(Dimension(name="column_name", value=numeric_col))
 
-        metric = self.series_to_metric(self.METRIC_NAME, [series])
+        metric = self.series_to_metric(self.METRIC_NAME, series)
         return [metric]