diff --git a/snuba/web/rpc/v1/resolvers/common/aggregation.py b/snuba/web/rpc/v1/resolvers/common/aggregation.py index ec67bda617..6c049b5565 100644 --- a/snuba/web/rpc/v1/resolvers/common/aggregation.py +++ b/snuba/web/rpc/v1/resolvers/common/aggregation.py @@ -43,16 +43,12 @@ class ExtrapolationContext(ABC): confidence_interval: Any average_sample_rate: float sample_count: int + is_extrapolated: bool @property def is_data_present(self) -> bool: return self.sample_count > 0 - @property - @abstractmethod - def is_extrapolated(self) -> bool: - raise NotImplementedError - @property @abstractmethod def reliability(self) -> Reliability.ValueType: @@ -64,6 +60,7 @@ def from_row( row_data: Dict[str, Any], ) -> ExtrapolationContext: value = row_data[column_label] + is_extrapolated = False confidence_interval = None average_sample_rate = 0 @@ -88,6 +85,7 @@ def from_row( continue if custom_column_information.custom_column_id == "confidence_interval": + is_extrapolated = True confidence_interval = col_value is_percentile = custom_column_information.metadata.get( @@ -116,6 +114,7 @@ def from_row( percentile=percentile, granularity=granularity, width=width, + is_extrapolated=is_extrapolated, ) return GenericExtrapolationContext( @@ -123,18 +122,12 @@ def from_row( confidence_interval=confidence_interval, average_sample_rate=average_sample_rate, sample_count=sample_count, + is_extrapolated=is_extrapolated, ) @dataclass(frozen=True) class GenericExtrapolationContext(ExtrapolationContext): - @property - def is_extrapolated(self) -> bool: - # We infer if a column is extrapolated or not by the presence of the - # confidence interval. It will be present for extrapolated aggregates - # but not for non-extrapolated aggregates and scalars. - return self.confidence_interval is not None - @cached_property def reliability(self) -> Reliability.ValueType: if not self.is_extrapolated or not self.is_data_present: diff --git a/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table_extrapolation.py b/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table_extrapolation.py index 8499e67ded..75dcb42bd4 100644 --- a/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table_extrapolation.py +++ b/tests/web/rpc/v1/test_endpoint_trace_item_table/test_endpoint_trace_item_table_extrapolation.py @@ -784,3 +784,105 @@ def test_formula(self) -> None: ], ), ] + + def test_aggregation_with_nulls(self) -> None: + spans_storage = get_storage(StorageKey("eap_spans")) + start = BASE_TIME + messages_a = [ + gen_message( + start - timedelta(minutes=i), + measurements={ + "custom_measurement": {"value": 1}, + "server_sample_rate": {"value": 1.0}, + }, + tags={"custom_tag": "a"}, + ) + for i in range(5) + ] + messages_b = [ + gen_message( + start - timedelta(minutes=i), + measurements={ + "custom_measurement2": {"value": 1}, + "server_sample_rate": {"value": 1.0}, + }, + tags={"custom_tag": "b"}, + ) + for i in range(5) + ] + write_raw_unprocessed_events(spans_storage, messages_a + messages_b) # type: ignore + + ts = Timestamp(seconds=int(BASE_TIME.timestamp())) + hour_ago = int((BASE_TIME - timedelta(hours=1)).timestamp()) + message = TraceItemTableRequest( + meta=RequestMeta( + project_ids=[1], + organization_id=1, + cogs_category="something", + referrer="something", + start_timestamp=Timestamp(seconds=hour_ago), + end_timestamp=ts, + trace_item_type=TraceItemType.TRACE_ITEM_TYPE_SPAN, + ), + columns=[ + Column( + key=AttributeKey(type=AttributeKey.TYPE_STRING, name="custom_tag") + ), + Column( + aggregation=AttributeAggregation( + aggregate=Function.FUNCTION_SUM, + key=AttributeKey( + type=AttributeKey.TYPE_DOUBLE, name="custom_measurement" + ), + label="sum(custom_measurement)", + extrapolation_mode=ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED, + ) + ), + Column( + aggregation=AttributeAggregation( + aggregate=Function.FUNCTION_SUM, + key=AttributeKey( + type=AttributeKey.TYPE_DOUBLE, name="custom_measurement2" + ), + label="sum(custom_measurement2)", + extrapolation_mode=ExtrapolationMode.EXTRAPOLATION_MODE_SAMPLE_WEIGHTED, + ) + ), + ], + group_by=[ + AttributeKey(type=AttributeKey.TYPE_STRING, name="custom_tag"), + ], + order_by=[ + TraceItemTableRequest.OrderBy( + column=Column( + key=AttributeKey( + type=AttributeKey.TYPE_STRING, name="custom_tag" + ) + ), + ), + ], + limit=5, + ) + response = EndpointTraceItemTable().execute(message) + assert response.column_values == [ + TraceItemColumnValues( + attribute_name="custom_tag", + results=[AttributeValue(val_str="a"), AttributeValue(val_str="b")], + ), + TraceItemColumnValues( + attribute_name="sum(custom_measurement)", + results=[AttributeValue(val_double=5), AttributeValue(is_null=True)], + reliabilities=[ + Reliability.RELIABILITY_LOW, + Reliability.RELIABILITY_UNSPECIFIED, + ], + ), + TraceItemColumnValues( + attribute_name="sum(custom_measurement2)", + results=[AttributeValue(is_null=True), AttributeValue(val_double=5)], + reliabilities=[ + Reliability.RELIABILITY_UNSPECIFIED, + Reliability.RELIABILITY_LOW, + ], + ), + ]