mlrun 1.10.0rc14__py3-none-any.whl → 1.10.0rc16__py3-none-any.whl

mlrun/model_monitoring/applications/base.py

@@ -17,7 +17,7 @@ import socket
 from abc import ABC, abstractmethod
 from collections import defaultdict
 from collections.abc import Iterator
-from contextlib import contextmanager
+from contextlib import contextmanager, nullcontext
 from datetime import datetime, timedelta
 from typing import Any, Literal, Optional, Union, cast
 
@@ -31,6 +31,7 @@ import mlrun.errors
 import mlrun.model_monitoring.api as mm_api
 import mlrun.model_monitoring.applications.context as mm_context
 import mlrun.model_monitoring.applications.results as mm_results
+import mlrun.model_monitoring.db._schedules as mm_schedules
 import mlrun.model_monitoring.helpers as mm_helpers
 from mlrun.serving.utils import MonitoringApplicationToDict
 from mlrun.utils import logger
@@ -183,14 +184,27 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
         cls,
         *,
         write_output: bool,
+        application_name: str,
+        artifact_path: str,
         stream_profile: Optional[ds_profile.DatastoreProfile],
         project: "mlrun.MlrunProject",
-    ) -> Iterator[dict[str, list[tuple]]]:
+    ) -> Iterator[
+        tuple[
+            dict[str, list[tuple]],
+            Optional[mm_schedules.ModelMonitoringSchedulesFileApplication],
+        ]
+    ]:
         endpoints_output: dict[str, list[tuple]] = defaultdict(list)
+        application_schedules = nullcontext()
         if write_output:
             cls._check_writer_is_up(project)
+            application_schedules = (
+                mm_schedules.ModelMonitoringSchedulesFileApplication(
+                    artifact_path, application=application_name
+                )
+            )
         try:
-            yield endpoints_output
+            yield endpoints_output, application_schedules.__enter__()
         finally:
             if write_output:
                 logger.debug(
@@ -218,6 +232,12 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
                     endpoints_output=endpoints_output,
                 )
 
+                logger.debug(
+                    "Saving the application schedules",
+                    application_name=application_name,
+                )
+                application_schedules.__exit__(None, None, None)
+
     def _handler(
         self,
         context: "mlrun.MLClientCtx",
@@ -230,6 +250,7 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
         end: Optional[str] = None,
         base_period: Optional[int] = None,
         write_output: bool = False,
+        allow_unordered_data: bool = False,
         stream_profile: Optional[ds_profile.DatastoreProfile] = None,
     ):
         """
@@ -250,6 +271,8 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
                 "working with endpoints, without any custom data-frame input"
             )
 
+        application_name = self.__class__.__name__
+
         feature_stats = (
             mm_api.get_sample_set_statistics(reference_data)
             if reference_data is not None
@@ -257,8 +280,12 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
         )
 
         with self._push_to_writer(
-            write_output=write_output, stream_profile=stream_profile, project=project
-        ) as endpoints_output:
+            write_output=write_output,
+            stream_profile=stream_profile,
+            application_name=application_name,
+            artifact_path=context.artifact_path,
+            project=project,
+        ) as (endpoints_output, application_schedules):
 
             def call_do_tracking(event: Optional[dict] = None):
                 nonlocal endpoints_output
@@ -268,7 +295,7 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
                 monitoring_context = (
                     mm_context.MonitoringApplicationContext._from_ml_ctx(
                         event=event,
-                        application_name=self.__class__.__name__,
+                        application_name=application_name,
                         context=context,
                         project=project,
                         sample_df=sample_data,
@@ -285,10 +312,16 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
                 resolved_endpoints = self._handle_endpoints_type_evaluate(
                     project=project, endpoints=endpoints
                 )
-                for window_start, window_end in self._window_generator(
-                    start, end, base_period
-                ):
-                    for endpoint_name, endpoint_id in resolved_endpoints:
+                for endpoint_name, endpoint_id in resolved_endpoints:
+                    for window_start, window_end in self._window_generator(
+                        start=start,
+                        end=end,
+                        base_period=base_period,
+                        application_schedules=application_schedules,
+                        endpoint_id=endpoint_id,
+                        application_name=application_name,
+                        allow_unordered_data=allow_unordered_data,
+                    ):
                         result = call_do_tracking(
                             event={
                                 mm_constants.ApplicationEvent.ENDPOINT_NAME: endpoint_name,
@@ -370,8 +403,103 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
             )
 
     @staticmethod
+    def _validate_and_get_window_length(
+        *, base_period: int, start_dt: datetime, end_dt: datetime
+    ) -> timedelta:
+        if not isinstance(base_period, int) or base_period <= 0:
+            raise mlrun.errors.MLRunValueError(
+                "`base_period` must be a nonnegative integer - the number of minutes in a monitoring window"
+            )
+
+        window_length = timedelta(minutes=base_period)
+
+        full_interval_length = end_dt - start_dt
+        remainder = full_interval_length % window_length
+        if remainder:
+            if full_interval_length < window_length:
+                extra_msg = (
+                    "The `base_period` is longer than the difference between `end` and `start`: "
+                    f"{full_interval_length}. Consider not specifying `base_period`."
+                )
+            else:
+                extra_msg = (
+                    f"Consider changing the `end` time to `end`={end_dt - remainder}"
+                )
+            raise mlrun.errors.MLRunValueError(
+                "The difference between `end` and `start` must be a multiple of `base_period`: "
+                f"`base_period`={window_length}, `start`={start_dt}, `end`={end_dt}. "
+                f"{extra_msg}"
+            )
+        return window_length
+
+    @staticmethod
+    def _validate_monotonically_increasing_data(
+        *,
+        application_schedules: Optional[
+            mm_schedules.ModelMonitoringSchedulesFileApplication
+        ],
+        endpoint_id: str,
+        start_dt: datetime,
+        end_dt: datetime,
+        base_period: Optional[int],
+        application_name: str,
+        allow_unordered_data: bool,
+    ) -> datetime:
+        """Make sure that the (app, endpoint) pair doesn't write output before the last analyzed window"""
+        if application_schedules:
+            last_analyzed = application_schedules.get_endpoint_last_analyzed(
+                endpoint_id
+            )
+            if last_analyzed:
+                if start_dt < last_analyzed:
+                    if allow_unordered_data:
+                        if last_analyzed < end_dt and base_period is None:
+                            logger.warn(
+                                "Setting the start time to last_analyzed since the original start time precedes "
+                                "last_analyzed",
+                                original_start=start_dt,
+                                new_start=last_analyzed,
+                                application_name=application_name,
+                                endpoint_id=endpoint_id,
+                            )
+                            start_dt = last_analyzed
+                        else:
+                            raise mlrun.errors.MLRunValueError(
+                                "The start time for the application and endpoint precedes the last analyzed time: "
+                                f"{start_dt=}, {last_analyzed=}, {application_name=}, {endpoint_id=}. "
+                                "Writing data out of order is not supported, and the start time could not be "
+                                "dynamically reset, as last_analyzed is later than the given end time or that "
+                                f"base_period was specified ({end_dt=}, {base_period=})."
+                            )
+                    else:
+                        raise mlrun.errors.MLRunValueError(
+                            "The start time for the application and endpoint precedes the last analyzed time: "
+                            f"{start_dt=}, {last_analyzed=}, {application_name=}, {endpoint_id=}. "
+                            "Writing data out of order is not supported. You should change the start time to "
+                            f"'{last_analyzed}' or later."
+                        )
+            else:
+                logger.debug(
+                    "The application is running on the endpoint for the first time",
+                    endpoint_id=endpoint_id,
+                    start_dt=start_dt,
+                    application_name=application_name,
+                )
+        return start_dt
+
+    @classmethod
     def _window_generator(
-        start: Optional[str], end: Optional[str], base_period: Optional[int]
+        cls,
+        *,
+        start: Optional[str],
+        end: Optional[str],
+        base_period: Optional[int],
+        application_schedules: Optional[
+            mm_schedules.ModelMonitoringSchedulesFileApplication
+        ],
+        endpoint_id: str,
+        application_name: str,
+        allow_unordered_data: bool,
     ) -> Iterator[tuple[Optional[datetime], Optional[datetime]]]:
         if start is None or end is None:
             # A single window based on the `sample_data` input - see `_handler`.
@@ -381,20 +509,36 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
         start_dt = datetime.fromisoformat(start)
         end_dt = datetime.fromisoformat(end)
 
+        start_dt = cls._validate_monotonically_increasing_data(
+            application_schedules=application_schedules,
+            endpoint_id=endpoint_id,
+            start_dt=start_dt,
+            end_dt=end_dt,
+            base_period=base_period,
+            application_name=application_name,
+            allow_unordered_data=allow_unordered_data,
+        )
+
         if base_period is None:
             yield start_dt, end_dt
+            if application_schedules:
+                application_schedules.update_endpoint_last_analyzed(
+                    endpoint_uid=endpoint_id, last_analyzed=end_dt
+                )
             return
 
-        if not isinstance(base_period, int) or base_period <= 0:
-            raise mlrun.errors.MLRunValueError(
-                "`base_period` must be a nonnegative integer - the number of minutes in a monitoring window"
-            )
+        window_length = cls._validate_and_get_window_length(
+            base_period=base_period, start_dt=start_dt, end_dt=end_dt
+        )
 
-        window_length = timedelta(minutes=base_period)
         current_start_time = start_dt
         while current_start_time < end_dt:
             current_end_time = min(current_start_time + window_length, end_dt)
             yield current_start_time, current_end_time
+            if application_schedules:
+                application_schedules.update_endpoint_last_analyzed(
+                    endpoint_uid=endpoint_id, last_analyzed=current_end_time
+                )
             current_start_time = current_end_time
 
     @classmethod
@@ -484,6 +628,7 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
         * ``end``, ``datetime``
         * ``base_period``, ``int``
         * ``write_output``, ``bool``
+        * ``allow_unordered_data``, ``bool``
 
         For Git sources, add the source archive to the returned job and change the handler:
 
@@ -567,6 +712,7 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
         end: Optional[datetime] = None,
         base_period: Optional[int] = None,
         write_output: bool = False,
+        allow_unordered_data: bool = False,
         stream_profile: Optional[ds_profile.DatastoreProfile] = None,
     ) -> "mlrun.RunObject":
         """
@@ -608,6 +754,8 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
         :param start:             The start time of the endpoint's data, not included.
                                   If you want the model endpoint's data at ``start`` included, you need to subtract a
                                   small ``datetime.timedelta`` from it.
+                                  Make sure to include the time zone when constructing `datetime.datetime` objects
+                                  manually.
         :param end:               The end time of the endpoint's data, included.
                                   Please note: when ``start`` and ``end`` are set, they create a left-open time interval
                                   ("window") :math:`(\\operatorname{start}, \\operatorname{end}]` that excludes the
@@ -616,17 +764,24 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
                                   taken in the window's data.
         :param base_period:       The window length in minutes. If ``None``, the whole window from ``start`` to ``end``
                                   is taken. If an integer is specified, the application is run from ``start`` to ``end``
-                                  in ``base_period`` length windows, except for the last window that ends at ``end`` and
-                                  therefore may be shorter:
+                                  in ``base_period`` length windows:
                                   :math:`(\\operatorname{start}, \\operatorname{start} + \\operatorname{base\\_period}],
                                   (\\operatorname{start} + \\operatorname{base\\_period},
                                   \\operatorname{start} + 2\\cdot\\operatorname{base\\_period}],
                                   ..., (\\operatorname{start} +
-                                  m\\cdot\\operatorname{base\\_period}, \\operatorname{end}]`,
-                                  where :math:`m` is some positive integer.
+                                  (m - 1)\\cdot\\operatorname{base\\_period}, \\operatorname{end}]`,
+                                  where :math:`m` is a positive integer and :math:`\\operatorname{end} =
+                                  \\operatorname{start} + m\\cdot\\operatorname{base\\_period}`.
+                                  Please note that the difference between ``end`` and ``start`` must be a multiple of
+                                  ``base_period``.
         :param write_output:      Whether to write the results and metrics to the time-series DB. Can be ``True`` only
                                   if ``endpoints`` are passed.
                                   Note: the model monitoring infrastructure must be up for the writing to work.
+        :param allow_unordered_data: Relevant only when writing outputs to the database. When ``False``, and the
+                                     requested ``start`` time precedes the ``end`` time of a previous run that also
+                                     wrote to the database - an error is raised.
+                                     If ``True``, when the previously described situation occurs, the relevant time
+                                     window is cut so that it starts at the earliest possible time after ``start``.
         :param stream_profile:    The stream datastore profile. It should be provided only when running locally and
                                   writing the outputs to the database (i.e., when both ``run_local`` and
                                   ``write_output`` are set to ``True``).
@@ -666,6 +821,7 @@ class ModelMonitoringApplicationBase(MonitoringApplicationToDict, ABC):
                 params["end"] = end.isoformat() if isinstance(end, datetime) else end
                 params["base_period"] = base_period
                 params["write_output"] = write_output
+                params["allow_unordered_data"] = allow_unordered_data
                 if stream_profile:
                     if not run_local:
                         raise mlrun.errors.MLRunValueError(

mlrun/model_monitoring/db/_schedules.py

@@ -13,25 +13,36 @@
 # limitations under the License.
 
 import json
+import sys
 from abc import ABC, abstractmethod
 from contextlib import AbstractContextManager
+from datetime import datetime, timezone
 from types import TracebackType
-from typing import Final, Optional
+from typing import TYPE_CHECKING, Final, Optional
 
 import botocore.exceptions
 
+import mlrun
 import mlrun.common.schemas as schemas
 import mlrun.errors
 import mlrun.model_monitoring.helpers
+import mlrun.utils.helpers
 from mlrun.utils import logger
 
+if TYPE_CHECKING:
+    if sys.version_info >= (3, 11):
+        from typing import Self
+    else:
+        from typing_extensions import Self
+
 
 class ModelMonitoringSchedulesFileBase(AbstractContextManager, ABC):
     DEFAULT_SCHEDULES: Final = {}
     INITIAL_CONTENT = json.dumps(DEFAULT_SCHEDULES)
     ENCODING = "utf-8"
 
-    def __init__(self):
+    def __init__(self) -> None:
+        # `self._item` is the persistent version of the monitoring schedules.
         self._item = self.get_data_item_object()
         if self._item:
             self._path = self._item.url
@@ -43,9 +54,16 @@ class ModelMonitoringSchedulesFileBase(AbstractContextManager, ABC):
             self._open_schedules = False
 
     @abstractmethod
-    def get_data_item_object(self) -> mlrun.DataItem:
+    def get_data_item_object(self) -> "mlrun.DataItem":
         pass
 
+    def _exists(self) -> bool:
+        """Return whether the file exists or not"""
+        return (
+            self._fs is None  # In-memory store
+            or self._fs.exists(self._path)
+        )
+
     def create(self) -> None:
         """Create a schedules file with initial content - an empty dictionary"""
         logger.debug("Creating model monitoring schedules file", path=self._item.url)
@@ -53,10 +71,7 @@ class ModelMonitoringSchedulesFileBase(AbstractContextManager, ABC):
 
     def delete(self) -> None:
         """Delete schedules file if it exists"""
-        if (
-            self._fs is None  # In-memory store
-            or self._fs.exists(self._path)
-        ):
+        if self._exists():
             logger.debug(
                 "Deleting model monitoring schedules file", path=self._item.url
             )
@@ -100,7 +115,7 @@ class ModelMonitoringSchedulesFileBase(AbstractContextManager, ABC):
         self._schedules = self.DEFAULT_SCHEDULES
         self._open_schedules = False
 
-    def __enter__(self) -> "ModelMonitoringSchedulesFileBase":
+    def __enter__(self) -> "Self":
         self._open()
         return super().__enter__()
 
@@ -129,12 +144,11 @@ class ModelMonitoringSchedulesFileEndpoint(ModelMonitoringSchedulesFileBase):
         :param project:     The project name.
         :param endpoint_id: The endpoint ID.
         """
-        # `self._item` is the persistent version of the monitoring schedules.
         self._project = project
         self._endpoint_id = endpoint_id
         super().__init__()
 
-    def get_data_item_object(self) -> mlrun.DataItem:
+    def get_data_item_object(self) -> "mlrun.DataItem":
         return mlrun.model_monitoring.helpers.get_monitoring_schedules_endpoint_data(
             project=self._project, endpoint_id=self._endpoint_id
         )
@@ -179,7 +193,7 @@ class ModelMonitoringSchedulesFileChief(ModelMonitoringSchedulesFileBase):
         self._project = project
         super().__init__()
 
-    def get_data_item_object(self) -> mlrun.DataItem:
+    def get_data_item_object(self) -> "mlrun.DataItem":
         return mlrun.model_monitoring.helpers.get_monitoring_schedules_chief_data(
             project=self._project
         )
@@ -216,22 +230,49 @@ class ModelMonitoringSchedulesFileChief(ModelMonitoringSchedulesFileBase):
         return set(self._schedules.keys())
 
     def get_or_create(self) -> None:
-        try:
-            self._open()
-        except (
-            mlrun.errors.MLRunNotFoundError,
-            # Different errors are raised for S3 or local storage, see ML-8042
-            botocore.exceptions.ClientError,
-            FileNotFoundError,
-        ):
+        if not self._exists():
             self.create()
 
 
-def delete_model_monitoring_schedules_folder(project: str) -> None:
-    """Delete the model monitoring schedules folder of the project"""
-    folder = mlrun.model_monitoring.helpers._get_monitoring_schedules_folder_path(
-        project
-    )
+class ModelMonitoringSchedulesFileApplication(ModelMonitoringSchedulesFileBase):
+    def __init__(self, out_path: str, application: str) -> None:
+        self._out_path = out_path
+        self._application = application
+        super().__init__()
+
+    def get_data_item_object(self) -> "mlrun.DataItem":
+        return mlrun.model_monitoring.helpers.get_monitoring_schedules_user_application_data(
+            out_path=self._out_path, application=self._application
+        )
+
+    def _open(self) -> None:
+        if not self._exists():
+            # Create the file when it is needed the first time
+            logger.info(
+                "Creating the application schedules file",
+                application=self._application,
+                path=self._path,
+            )
+            self.create()
+        super()._open()
+
+    def get_endpoint_last_analyzed(self, endpoint_uid: str) -> Optional[datetime]:
+        self._check_open_schedules()
+        if endpoint_uid in self._schedules:
+            return datetime.fromisoformat(self._schedules[endpoint_uid])
+        else:
+            return None
+
+    def update_endpoint_last_analyzed(
+        self, endpoint_uid: str, last_analyzed: datetime
+    ) -> None:
+        self._check_open_schedules()
+        self._schedules[endpoint_uid] = last_analyzed.astimezone(
+            timezone.utc
+        ).isoformat()
+
+
+def _delete_folder(folder: str) -> None:
     fs = mlrun.datastore.store_manager.object(folder).store.filesystem
     if fs and fs.exists(folder):
         logger.debug("Deleting model monitoring schedules folder", folder=folder)
@@ -240,3 +281,22 @@ def delete_model_monitoring_schedules_folder(project: str) -> None:
         raise mlrun.errors.MLRunValueError(
             "Cannot delete a folder without a file-system"
         )
+
+
+def delete_model_monitoring_schedules_folder(project: str) -> None:
+    """Delete the model monitoring schedules folder of the project"""
+    folder = mlrun.model_monitoring.helpers._get_monitoring_schedules_folder_path(
+        project
+    )
+    _delete_folder(folder)
+
+
+def delete_model_monitoring_schedules_user_folder(project: str) -> None:
+    """Delete the user created schedules folder (created through `app.evaluate`)"""
+    out_path = mlrun.utils.helpers.template_artifact_path(
+        mlrun.mlconf.artifact_path, project=project
+    )
+    folder = mlrun.model_monitoring.helpers._get_monitoring_schedules_user_folder_path(
+        out_path
+    )
+    _delete_folder(folder)

mlrun/model_monitoring/db/tsdb/base.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import Callable, ClassVar, Literal, Optional, Union
 
 import pandas as pd
@@ -80,6 +80,21 @@ class TSDBConnector(ABC):
         :raise mlrun.errors.MLRunRuntimeError: If an error occurred while writing the event.
         """
 
+    @abstractmethod
+    def get_drift_data(
+        self,
+        start: datetime,
+        end: datetime,
+    ) -> mm_schemas.ModelEndpointDriftValues:
+        """
+        Fetches drift counts per interval in the specified time range.
+
+        :param start: The start time of the query.
+        :param end:   The end time of the query.
+
+        :return: A ModelEndpointDriftValues object containing drift data.
+        """
+
     @abstractmethod
     def delete_tsdb_records(
         self,
@@ -703,3 +718,59 @@ class TSDBConnector(ABC):
                 )
             )
         return {dict_key: metrics}
+
+    @staticmethod
+    def _prepare_aligned_start_end(
+        start: datetime, end: datetime
+    ) -> tuple[datetime, datetime, str]:
+        delta = end - start
+        if delta <= timedelta(hours=6):
+            interval = "10m"
+            start = start.replace(
+                minute=start.minute // 10 * 10, second=0, microsecond=0
+            )
+        elif delta <= timedelta(hours=72):
+            interval = "1h"
+            start = start.replace(minute=0, second=0, microsecond=0)
+        else:
+            interval = "1d"
+            start = start.replace(hour=0, minute=0, second=0, microsecond=0)
+
+        interval_map = {
+            "10m": timedelta(minutes=10),
+            "1h": timedelta(hours=1),
+            "1d": timedelta(days=1),
+        }
+        delta = end - start
+        interval_td = interval_map[interval]
+        end = start + (delta // interval_td) * interval_td
+        return start, end, interval
+
+    @staticmethod
+    def _df_to_drift_data(df: pd.DataFrame) -> mm_schemas.ModelEndpointDriftValues:
+        suspected_val = mm_schemas.constants.ResultStatusApp.potential_detection.value
+        detected_val = mm_schemas.constants.ResultStatusApp.detected.value
+        aggregated_df = (
+            df.groupby(["_wstart", f"max({mm_schemas.ResultData.RESULT_STATUS})"])
+            .size()  # add size column for each interval x result-status combination
+            .unstack()  # create a size column for each result-status
+            .reindex(
+                columns=[suspected_val, detected_val], fill_value=0
+            )  # ensure both columns exists
+            .fillna(0)
+            .astype(int)
+            .rename(
+                columns={
+                    suspected_val: "count_suspected",
+                    detected_val: "count_detected",
+                }
+            )
+        )
+        values = list(
+            zip(
+                aggregated_df.index,
+                aggregated_df["count_suspected"],
+                aggregated_df["count_detected"],
+            )
+        )
+        return mm_schemas.ModelEndpointDriftValues(values=values)

mlrun/model_monitoring/db/tsdb/tdengine/schemas.py

@@ -165,6 +165,7 @@ class TDEngineSchema:
         preform_agg_funcs_columns: Optional[list[str]] = None,
         order_by: Optional[str] = None,
         desc: Optional[bool] = None,
+        partition_by: Optional[str] = None,
     ) -> str:
         if agg_funcs and not columns_to_filter:
             raise mlrun.errors.MLRunInvalidArgumentError(
@@ -176,7 +177,10 @@ class TDEngineSchema:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "`agg_funcs` must be provided when using interval"
             )
-
+        if partition_by and not agg_funcs:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "`agg_funcs` must be provided when using partition by"
+            )
         if sliding_window_step and not interval:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "`interval` must be provided when using sliding window"
@@ -232,6 +236,8 @@ class TDEngineSchema:
                 if isinstance(group_by, list):
                     group_by = ", ".join(group_by)
                 query.write(f" GROUP BY {group_by}")
+            if partition_by:
+                query.write(f" PARTITION BY {partition_by}")
             if order_by:
                 desc = " DESC" if desc else ""
                 query.write(f" ORDER BY {order_by}{desc}")