mlrun 1.10.0rc18__py3-none-any.whl → 1.11.0rc16__py3-none-any.whl

mlrun/model_monitoring/db/tsdb/timescaledb/queries/timescaledb_predictions_queries.py

@@ -0,0 +1,376 @@
+# Copyright 2025 Iguazio
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from datetime import datetime, timedelta
+from typing import Optional, Union
+
+import pandas as pd
+import v3io_frames.client
+
+import mlrun
+import mlrun.common.schemas.model_monitoring as mm_schemas
+import mlrun.errors
+import mlrun.model_monitoring.db.tsdb.timescaledb.timescaledb_schema as timescaledb_schema
+import mlrun.utils
+from mlrun.common.schemas.model_monitoring.model_endpoints import _MetricPoint
+from mlrun.model_monitoring.db.tsdb.timescaledb.utils.timescaledb_dataframe_processor import (
+    TimescaleDBDataFrameProcessor,
+)
+from mlrun.model_monitoring.db.tsdb.timescaledb.utils.timescaledb_query_builder import (
+    TimescaleDBQueryBuilder,
+)
+from mlrun.model_monitoring.helpers import get_invocations_fqn
+
+
+class TimescaleDBPredictionsQueries:
+    """
+    Query class containing predictions-related query methods for TimescaleDB.
+
+    Can be used as a mixin or standalone instance with proper initialization.
+    """
+
+    def __init__(
+        self,
+        project: Optional[str] = None,
+        connection=None,
+        pre_aggregate_manager=None,
+        tables: Optional[dict] = None,
+    ):
+        """
+        Initialize TimescaleDB predictions query handler.
+
+        :param project: Project name
+        :param connection: TimescaleDB connection instance
+        :param pre_aggregate_manager: PreAggregateManager instance
+        :param tables: Dictionary of table schemas
+        """
+        self.project = project
+        self._connection = connection
+        self._pre_aggregate_manager = pre_aggregate_manager
+        self.tables = tables
+
+    def read_predictions_impl(
+        self,
+        *,
+        endpoint_id: Optional[str] = None,
+        start: datetime,
+        end: datetime,
+        columns: Optional[list[str]] = None,
+        aggregation_window: Optional[str] = None,
+        agg_funcs: Optional[list[str]] = None,
+        limit: Optional[int] = None,
+        use_pre_aggregates: bool = True,
+        timestamp_column: Optional[str] = None,
+    ) -> pd.DataFrame:
+        """Read predictions data from TimescaleDB (predictions table) - returns DataFrame.
+
+        :param endpoint_id: Endpoint ID to filter by, or None to get all endpoints
+        :param start: Start time
+        :param end: End time
+        :param columns: Optional list of specific columns to return
+        :param aggregation_window: Optional aggregation window (e.g., "1h", "1d")
+        :param agg_funcs: Optional list of aggregation functions (e.g., ["avg", "max"])
+        :param limit: Optional limit on number of results
+        :param use_pre_aggregates: Whether to use pre-aggregates if available
+        :param timestamp_column: Optional timestamp column to use for time filtering
+        :return: DataFrame with predictions data
+        """
+        if (agg_funcs and not aggregation_window) or (
+            aggregation_window and not agg_funcs
+        ):
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "both or neither of `aggregation_window` and `agg_funcs` must be provided"
+            )
+
+        # Align times if aggregation window is provided
+        start, end = self._pre_aggregate_manager.align_time_range(
+            start, end, aggregation_window
+        )
+
+        # Check if we can use pre-aggregates
+        can_use_pre_aggregates = (
+            use_pre_aggregates
+            and self._pre_aggregate_manager.can_use_pre_aggregates(
+                interval=aggregation_window, agg_funcs=agg_funcs
+            )
+        )
+
+        table_schema = self.tables[mm_schemas.TimescaleDBTables.PREDICTIONS]
+        filter_query = TimescaleDBQueryBuilder.build_endpoint_filter(endpoint_id)
+
+        query = table_schema._get_records_query(
+            start=start,
+            end=end,
+            columns_to_filter=columns,
+            filter_query=filter_query,
+            interval=aggregation_window if can_use_pre_aggregates else None,
+            agg_funcs=agg_funcs if can_use_pre_aggregates else None,
+            limit=limit,
+            use_pre_aggregates=can_use_pre_aggregates,
+            timestamp_column=timestamp_column,
+        )
+
+        result = self._connection.run(query=query)
+        df = TimescaleDBDataFrameProcessor.from_query_result(result)
+
+        if not df.empty:
+            # Set up time index based on whether we used aggregation
+            if aggregation_window and can_use_pre_aggregates:
+                time_col = timescaledb_schema.TIME_BUCKET_COLUMN
+            else:
+                time_col = table_schema.time_column
+
+            if time_col in df.columns:
+                df[time_col] = pd.to_datetime(df[time_col])
+                df.set_index(time_col, inplace=True)
+
+        return df
+
+    def read_predictions(
+        self,
+        *,
+        endpoint_id: str,
+        start: datetime,
+        end: datetime,
+        aggregation_window: Optional[str] = None,
+        agg_funcs: Optional[list[str]] = None,
+        limit: Optional[int] = None,
+        use_pre_aggregates: bool = True,
+    ) -> Union[
+        mm_schemas.ModelEndpointMonitoringMetricValues,
+        mm_schemas.ModelEndpointMonitoringMetricNoData,
+    ]:
+        """Read predictions with optional pre-aggregate optimization."""
+
+        table_schema = self.tables[mm_schemas.TimescaleDBTables.PREDICTIONS]
+        columns = [
+            table_schema.time_column,
+            mm_schemas.EventFieldType.ESTIMATED_PREDICTION_COUNT,
+        ]
+
+        # Get raw DataFrame from read_predictions_impl
+        df = self.read_predictions_impl(
+            endpoint_id=endpoint_id,
+            start=start,
+            end=end,
+            columns=columns,
+            aggregation_window=aggregation_window,
+            agg_funcs=agg_funcs,
+            limit=limit,
+            use_pre_aggregates=use_pre_aggregates,
+        )
+
+        # Convert to domain objects
+        full_name = get_invocations_fqn(self.project)
+
+        if df.empty:
+            return TimescaleDBDataFrameProcessor.handle_empty_dataframe(full_name)
+
+        # Determine value column name based on whether aggregation was used
+        can_use_pre_aggregates = (
+            use_pre_aggregates
+            and aggregation_window
+            and agg_funcs
+            and self._pre_aggregate_manager.can_use_pre_aggregates(
+                interval=aggregation_window, agg_funcs=agg_funcs
+            )
+        )
+
+        if agg_funcs and can_use_pre_aggregates:
+            value_col = (
+                f"{agg_funcs[0]}_{mm_schemas.EventFieldType.ESTIMATED_PREDICTION_COUNT}"
+            )
+        else:
+            value_col = mm_schemas.EventFieldType.ESTIMATED_PREDICTION_COUNT
+
+        return mm_schemas.ModelEndpointMonitoringMetricValues(
+            full_name=full_name,
+            values=[
+                _MetricPoint(timestamp=timestamp, value=value)
+                for timestamp, value in zip(df.index, df[value_col])
+            ],
+        )
+
+    def get_last_request(
+        self,
+        endpoint_ids: Union[str, list[str]],
+        start: Optional[datetime] = None,
+        end: Optional[datetime] = None,
+        interval: Optional[str] = None,
+    ) -> pd.DataFrame:
+        """Get last request timestamp with optional pre-aggregate optimization."""
+
+        # Prepare time range and interval (no auto-determination since interval may be None)
+        start, end, interval = TimescaleDBQueryBuilder.prepare_time_range_and_interval(
+            self._pre_aggregate_manager,
+            start,
+            end,
+            interval,
+            auto_determine_interval=False,
+        )
+        use_pre_aggregates = self._pre_aggregate_manager.can_use_pre_aggregates(
+            interval=interval
+        )
+
+        table_schema = self.tables[mm_schemas.TimescaleDBTables.PREDICTIONS]
+        filter_query = TimescaleDBQueryBuilder.build_endpoint_filter(endpoint_ids)
+
+        if use_pre_aggregates:
+            # Calculate latest (MAX) timestamp and corresponding latency per endpoint
+            # Use subquery to get time-bucketed data, then MAX over those results
+            subquery = table_schema._get_records_query(
+                start=start,
+                end=end,
+                columns_to_filter=[
+                    timescaledb_schema.TIME_BUCKET_COLUMN,
+                    f"max_{table_schema.time_column}",
+                    f"max_{mm_schemas.EventFieldType.LATENCY}",
+                    mm_schemas.WriterEvent.ENDPOINT_ID,
+                ],
+                filter_query=filter_query,
+                agg_funcs=["max"],
+                interval=interval,
+                use_pre_aggregates=True,
+            )
+
+            # Use helper to build endpoint aggregation query
+            query = TimescaleDBQueryBuilder.build_endpoint_aggregation_query(
+                subquery=subquery,
+                aggregation_columns={
+                    mm_schemas.EventFieldType.LAST_REQUEST: f"MAX(max_{table_schema.time_column})",
+                    "last_latency": f"MAX(max_{mm_schemas.EventFieldType.LATENCY})",
+                },
+            )
+
+            result = self._connection.run(query=query)
+            df = TimescaleDBDataFrameProcessor.from_query_result(result)
+        else:
+            # Use PostgreSQL DISTINCT ON for raw data - most efficient approach
+            query = f"""
+            SELECT DISTINCT ON ({mm_schemas.WriterEvent.ENDPOINT_ID})
+                {mm_schemas.WriterEvent.ENDPOINT_ID} AS endpoint_id,
+                {table_schema.time_column} AS {mm_schemas.EventFieldType.LAST_REQUEST},
+                {mm_schemas.EventFieldType.LATENCY} AS last_latency
+            FROM {table_schema.full_name()}
+            WHERE {filter_query}
+            AND {table_schema.time_column} >= '{start}'
+            AND {table_schema.time_column} <= '{end}'
+            ORDER BY {mm_schemas.WriterEvent.ENDPOINT_ID}, {table_schema.time_column} DESC;
+            """
+
+            result = self._connection.run(query=query)
+            df = TimescaleDBDataFrameProcessor.from_query_result(result)
+
+        # Convert timestamp to proper format (common for both paths)
+        if not df.empty and mm_schemas.EventFieldType.LAST_REQUEST in df.columns:
+            df[mm_schemas.EventFieldType.LAST_REQUEST] = pd.to_datetime(
+                df[mm_schemas.EventFieldType.LAST_REQUEST], errors="coerce", utc=True
+            )
+
+        return df
+
+    def get_avg_latency(
+        self,
+        endpoint_ids: Union[str, list[str]],
+        start: Optional[datetime] = None,
+        end: Optional[datetime] = None,
+        get_raw: bool = False,
+    ) -> Union[pd.DataFrame, list[v3io_frames.client.RawFrame]]:
+        """Get average latency with automatic pre-aggregate optimization, returning single value per endpoint."""
+
+        # Convert single endpoint to list for consistent handling
+        if isinstance(endpoint_ids, str):
+            endpoint_ids = [endpoint_ids]
+
+        # Set default start time and get end time
+        start = start or (mlrun.utils.datetime_now() - timedelta(hours=24))
+        # Prepare time range with auto-determined interval
+        start, end, interval = TimescaleDBQueryBuilder.prepare_time_range_and_interval(
+            self._pre_aggregate_manager, start, end
+        )
+
+        table_schema = self.tables[mm_schemas.TimescaleDBTables.PREDICTIONS]
+        filter_query = TimescaleDBQueryBuilder.build_endpoint_filter(endpoint_ids)
+
+        def build_pre_agg_query():
+            # Calculate overall average in SQL across all time buckets
+            # Use subquery to get time-bucketed data, then AVG over those results
+            subquery = table_schema._get_records_query(
+                start=start,
+                end=end,
+                columns_to_filter=[
+                    timescaledb_schema.TIME_BUCKET_COLUMN,
+                    mm_schemas.ModelEndpointSchema.AVG_LATENCY,
+                    mm_schemas.WriterEvent.ENDPOINT_ID,
+                ],
+                filter_query=filter_query,
+                agg_funcs=["avg"],
+                interval=interval,
+                use_pre_aggregates=True,
+            )
+
+            # Use helper to build endpoint aggregation query
+            return TimescaleDBQueryBuilder.build_endpoint_aggregation_query(
+                subquery=subquery,
+                aggregation_columns={
+                    mm_schemas.ModelEndpointSchema.AVG_LATENCY: f"AVG({mm_schemas.ModelEndpointSchema.AVG_LATENCY})"
+                },
+            )
+
+        def build_raw_query():
+            # Single aggregated value across entire time range
+            columns = [
+                f"{mm_schemas.WriterEvent.ENDPOINT_ID} AS {mm_schemas.WriterEvent.ENDPOINT_ID}",
+                f"AVG({mm_schemas.EventFieldType.LATENCY}) AS {mm_schemas.ModelEndpointSchema.AVG_LATENCY}",
+            ]
+            group_by_columns = [mm_schemas.WriterEvent.ENDPOINT_ID]
+
+            # Add additional filter to exclude invalid latency values
+            latency_col = mm_schemas.EventFieldType.LATENCY
+            latency_filter = f"{latency_col} IS NOT NULL AND {latency_col} > 0"
+            enhanced_filter_query = (
+                f"{filter_query} AND {latency_filter}"
+                if filter_query
+                else latency_filter
+            )
+
+            return table_schema._get_records_query(
+                start=start,
+                end=end,
+                columns_to_filter=columns,
+                filter_query=enhanced_filter_query,
+                group_by=group_by_columns,
+                order_by=mm_schemas.WriterEvent.ENDPOINT_ID,
+            )
+
+        # Column mapping rules for results (both pre-agg and raw return same structure now)
+        column_mapping_rules = {
+            mm_schemas.ModelEndpointSchema.AVG_LATENCY: [
+                mm_schemas.ModelEndpointSchema.AVG_LATENCY,
+                "average_latency",
+                mm_schemas.EventFieldType.LATENCY,
+            ],
+            mm_schemas.WriterEvent.ENDPOINT_ID: [mm_schemas.WriterEvent.ENDPOINT_ID],
+        }
+
+        # Both queries now return single value per endpoint, no post-processing needed
+        return self._connection.execute_with_fallback(
+            self._pre_aggregate_manager,
+            build_pre_agg_query,
+            build_raw_query,
+            interval=interval,
+            agg_funcs=["avg"],
+            column_mapping_rules=column_mapping_rules,
+            debug_name="avg_latency",
+        )