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

mlrun/datastore/base.py

@@ -11,11 +11,14 @@
 # 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.
+import datetime
+import os
+import os.path
 import tempfile
 import urllib.parse
 from base64 import b64encode
 from copy import copy
-from os import path, remove
+from types import ModuleType
 from typing import Optional, Union
 from urllib.parse import urlparse
 
@@ -26,6 +29,7 @@ import pyarrow
 import pytz
 import requests
 
+import mlrun.common.schemas
 import mlrun.config
 import mlrun.errors
 from mlrun.datastore.remote_client import BaseRemoteClient
@@ -156,6 +160,195 @@ class DataStore(BaseRemoteClient):
     def get_spark_options(self, path=None):
         return {}
 
+    @staticmethod
+    def _is_directory_in_range(
+        start_time: Optional[datetime.datetime],
+        end_time: Optional[datetime.datetime],
+        year: int,
+        month: Optional[int] = None,
+        day: Optional[int] = None,
+        hour: Optional[int] = None,
+        **kwargs,
+    ):
+        """Check if a partition directory (year=.., month=.., etc.) is in the time range."""
+        from dateutil.relativedelta import relativedelta
+
+        partition_start = datetime.datetime(
+            year=year,
+            month=month or 1,
+            day=day or 1,
+            hour=hour or 0,
+            tzinfo=start_time.tzinfo if start_time else end_time.tzinfo,
+        )
+        partition_end = (
+            partition_start
+            + relativedelta(
+                years=1 if month is None else 0,
+                months=1 if day is None and month is not None else 0,
+                days=1 if hour is None and day is not None else 0,
+                hours=1 if hour is not None else 0,
+            )
+            - datetime.timedelta(microseconds=1)
+        )
+
+        if (end_time and end_time < partition_start) or (
+            start_time and start_time > partition_end
+        ):
+            return False
+        return True
+
+    @staticmethod
+    def _list_partition_paths_helper(
+        paths: list[str],
+        start_time: Optional[datetime.datetime],
+        end_time: Optional[datetime.datetime],
+        current_path: str,
+        partition_level: str,
+        filesystem,
+    ):
+        directory_split = current_path.rsplit("/", 1)
+        time_unit = None
+        directory_start, directory_end = "", ""
+        if len(directory_split) == 2:
+            directory_start, directory_end = directory_split
+            time_unit = directory_end.split("=")[0] if "=" in directory_end else None
+
+        if not time_unit and directory_end.endswith((".parquet", ".pq")):
+            paths.append(directory_start.rstrip("/"))
+            return
+        elif time_unit and time_unit == partition_level:
+            paths.append(current_path.rstrip("/"))
+            return
+
+        directories = filesystem.ls(current_path, detail=True)
+        if len(directories) == 0:
+            return
+        for directory in directories:
+            current_path = directory["name"]
+            parts = [p for p in current_path.split("/") if "=" in p]
+            kwargs = {}
+            for part in parts:
+                key, value = part.split("=", 1)
+                if value.isdigit():
+                    value = int(value)
+                kwargs[key] = value
+            if DataStore._is_directory_in_range(start_time, end_time, **kwargs):
+                DataStore._list_partition_paths_helper(
+                    paths,
+                    start_time,
+                    end_time,
+                    current_path,
+                    partition_level,
+                    filesystem,
+                )
+
+    @staticmethod
+    def _list_partitioned_paths(
+        base_url: str,
+        start_time: Optional[datetime.datetime],
+        end_time: Optional[datetime.datetime],
+        partition_level: str,
+        filesystem,
+    ):
+        paths = []
+        parsed_base_url = urlparse(base_url)
+        base_path = parsed_base_url.path
+
+        if parsed_base_url.scheme not in ["v3io", "v3ios"]:
+            base_path = parsed_base_url.netloc + base_path
+
+        DataStore._list_partition_paths_helper(
+            paths, start_time, end_time, base_path, partition_level, filesystem
+        )
+        paths = [
+            DataStore._reconstruct_path_from_base_url(parsed_base_url, path)
+            for path in paths
+        ]
+        return paths
+
+    @staticmethod
+    def _reconstruct_path_from_base_url(
+        parsed_base_url: urllib.parse.ParseResult, returned_path: str
+    ) -> str:
+        scheme = parsed_base_url.scheme
+        authority = parsed_base_url.netloc
+        returned_path = returned_path.lstrip("/")
+        if scheme == "v3io":
+            return f"{scheme}://{authority}/{returned_path}"
+        else:
+            return f"{scheme}://{returned_path}"
+
+    @staticmethod
+    def _clean_filters_for_partitions(
+        filters: list[list[tuple]],
+        partition_keys: list[str],
+    ):
+        """
+        Remove partition keys from filters.
+
+        :param filters: pandas-style filters
+                Example: [[('year','=',2025),('month','=',11),('timestamp','>',ts1)]]
+        :param partition_keys: partition columns handled via directory
+
+        :return list of list of tuples: cleaned filters without partition keys
+        """
+        cleaned_filters = []
+        for group in filters:
+            new_group = [f for f in group if f[0] not in partition_keys]
+            if new_group:
+                cleaned_filters.append(new_group)
+        return cleaned_filters
+
+    @staticmethod
+    def _read_partitioned_parquet(
+        base_url: str,
+        start_time: Optional[datetime.datetime],
+        end_time: Optional[datetime.datetime],
+        partition_keys: list[str],
+        df_module: ModuleType,
+        filesystem: fsspec.AbstractFileSystem,
+        **kwargs,
+    ):
+        """
+        Reads only the relevant partitions and concatenates the results.
+        Note that partition_keys cannot be empty.
+        """
+        logger.debug(f"Starting partition discovery process for {base_url}")
+
+        paths = DataStore._list_partitioned_paths(
+            base_url,
+            start_time,
+            end_time,
+            partition_keys[-1],
+            filesystem,
+        )
+
+        dfs = []
+        for current_path in paths:
+            try:
+                kwargs["filters"] = DataStore._clean_filters_for_partitions(
+                    kwargs["filters"], partition_keys
+                )
+                df = df_module.read_parquet(current_path, **kwargs)
+                logger.debug(
+                    "Finished reading DataFrame from subpath",
+                    url=current_path,
+                )
+                dfs.append(df)
+            except FileNotFoundError as e:
+                # Skip partitions that don't exist or have no data
+                logger.warning(
+                    "Failed to read DataFrame", url=current_path, exception=e
+                )
+
+        final_df = pd.concat(dfs) if dfs else pd.DataFrame()
+        logger.debug(
+            "Finished reading partitioned parquet files",
+            url=base_url,
+            columns=final_df.columns,
+        )
+        return final_df
+
     @staticmethod
     def _parquet_reader(
         df_module,
@@ -165,6 +358,7 @@ class DataStore(BaseRemoteClient):
         start_time,
         end_time,
         additional_filters,
+        optimize_discovery,
     ):
         from storey.utils import find_filters, find_partitions
 
@@ -203,7 +397,10 @@ class DataStore(BaseRemoteClient):
                 )
 
             if start_time or end_time or additional_filters:
-                partitions_time_attributes = find_partitions(url, file_system)
+                partitions_time_attributes, partitions = find_partitions(
+                    url, file_system
+                )
+                logger.debug("Partitioned parquet read", partitions=partitions)
                 set_filters(
                     partitions_time_attributes,
                     start_time,
@@ -211,8 +408,28 @@ class DataStore(BaseRemoteClient):
                     additional_filters,
                     kwargs,
                 )
+
                 try:
-                    return df_module.read_parquet(*args, **kwargs)
+                    if (
+                        optimize_discovery
+                        and partitions_time_attributes
+                        and DataStore._verify_path_partition_level(
+                            urlparse(url).path, partitions
+                        )
+                        and (start_time or end_time)
+                    ):
+                        return DataStore._read_partitioned_parquet(
+                            url,
+                            start_time,
+                            end_time,
+                            partitions_time_attributes,
+                            df_module,
+                            file_system,
+                            **kwargs,
+                        )
+
+                    else:
+                        return df_module.read_parquet(*args, **kwargs)
                 except pyarrow.lib.ArrowInvalid as ex:
                     if not str(ex).startswith(
                         "Cannot compare timestamp with timezone to timestamp without timezone"
@@ -238,7 +455,24 @@ class DataStore(BaseRemoteClient):
                         additional_filters,
                         kwargs,
                     )
-                    return df_module.read_parquet(*args, **kwargs)
+                    if (
+                        optimize_discovery
+                        and partitions_time_attributes
+                        and DataStore._verify_path_partition_level(
+                            urlparse(url).path, partitions
+                        )
+                    ):
+                        return DataStore._read_partitioned_parquet(
+                            url,
+                            start_time_inner,
+                            end_time_inner,
+                            partitions_time_attributes,
+                            df_module,
+                            file_system,
+                            **kwargs,
+                        )
+                    else:
+                        return df_module.read_parquet(*args, **kwargs)
             else:
                 return df_module.read_parquet(*args, **kwargs)
 
@@ -261,6 +495,10 @@ class DataStore(BaseRemoteClient):
         file_url = self._sanitize_url(url)
         is_csv, is_json, drop_time_column = False, False, False
         file_system = self.filesystem
+
+        # Feature flag optimize partition discovery by providing specific partition levels urls to the parquet reader
+        optimize_discovery = kwargs.pop("optimize_discovery", True)
+
         if file_url.endswith(".csv") or format == "csv":
             is_csv = True
             drop_time_column = False
@@ -322,6 +560,7 @@ class DataStore(BaseRemoteClient):
                 start_time,
                 end_time,
                 additional_filters,
+                optimize_discovery,
             )
 
         elif file_url.endswith(".json") or format == "json":
@@ -347,7 +586,7 @@ class DataStore(BaseRemoteClient):
             temp_file = tempfile.NamedTemporaryFile(delete=False)
             self.download(self._join(subpath), temp_file.name)
             df = reader(temp_file.name, **kwargs)
-            remove(temp_file.name)
+            os.remove(temp_file.name)
 
         if is_json or is_csv:
             # for parquet file the time filtering is executed in `reader`
@@ -387,6 +626,26 @@ class DataStore(BaseRemoteClient):
         except ImportError:
             return False
 
+    @staticmethod
+    def _verify_path_partition_level(base_path: str, partitions: list[str]) -> bool:
+        if not partitions:
+            return False
+
+        path_parts = base_path.strip("/").split("/")
+        path_parts = [part.split("=")[0] for part in path_parts if "=" in part]
+        if "hour" in partitions:
+            hour_index = partitions.index("hour")
+        else:
+            return False
+        for i, part in enumerate(partitions):
+            if not (
+                part in path_parts
+                or part in ["year", "month", "day", "hour"]
+                or i > hour_index
+            ):
+                return False
+        return True
+
 
 class DataItem:
     """Data input/output class abstracting access to various local/remote data sources
@@ -439,7 +698,7 @@ class DataItem:
     @property
     def suffix(self):
         """DataItem suffix (file extension) e.g. '.png'"""
-        _, file_ext = path.splitext(self._path)
+        _, file_ext = os.path.splitext(self._path)
         return file_ext
 
     @property
@@ -548,7 +807,7 @@ class DataItem:
             return
 
         if self._local_path:
-            remove(self._local_path)
+            os.remove(self._local_path)
             self._local_path = ""
 
     def as_df(
@@ -648,8 +907,10 @@ def basic_auth_header(user, password):
     username = user.encode("latin1")
     password = password.encode("latin1")
     base = b64encode(b":".join((username, password))).strip()
-    authstr = "Basic " + base.decode("ascii")
-    return {"Authorization": authstr}
+    authstr = mlrun.common.schemas.AuthorizationHeaderPrefixes.basic + base.decode(
+        "ascii"
+    )
+    return {mlrun.common.schemas.HeaderNames.authorization: authstr}
 
 
 class HttpStore(DataStore):
@@ -696,7 +957,10 @@ class HttpStore(DataStore):
         token = self._get_secret_or_env("HTTPS_AUTH_TOKEN")
         if token:
             self._https_auth_token = token
-            self._headers.setdefault("Authorization", f"Bearer {token}")
+            self._headers.setdefault(
+                mlrun.common.schemas.HeaderNames.authorization,
+                f"{mlrun.common.schemas.AuthorizationHeaderPrefixes.bearer}{token}",
+            )
 
     def _validate_https_token(self):
         if self._https_auth_token and self._schema in ["http"]:

mlrun/datastore/datastore.py

@@ -47,7 +47,7 @@ from .v3io import V3ioStore
 in_memory_store = InMemoryStore()
 
 
-def schema_to_store(schema) -> DataStore.__subclasses__():
+def schema_to_store(schema) -> type[DataStore]:
     # import store classes inside to enable making their dependencies optional (package extras)
 
     if not schema or schema in get_local_file_schema():

mlrun/datastore/datastore_profile.py

@@ -16,7 +16,7 @@ import ast
 import base64
 import json
 import typing
-from urllib.parse import ParseResult, urlparse
+from urllib.parse import ParseResult, quote, unquote, urlparse
 
 import pydantic.v1
 from deprecated import deprecated
@@ -283,8 +283,9 @@ class DatastoreProfileRedis(DatastoreProfile):
 
     def url_with_credentials(self):
         parsed_url = urlparse(self.endpoint_url)
-        username = self.username
-        password = self.password
+        # URL-encode username and password to handle special characters like @, :, /
+        username = quote(self.username, safe="") if self.username else None
+        password = quote(self.password, safe="") if self.password else None
         netloc = parsed_url.hostname
         if username:
             if password:
@@ -448,40 +449,71 @@ class DatastoreProfileHdfs(DatastoreProfile):
         return f"webhdfs://{self.host}:{self.http_port}{subpath}"
 
 
-class DatastoreProfileTDEngine(DatastoreProfile):
+class DatastoreProfilePostgreSQL(DatastoreProfile):
     """
-    A profile that holds the required parameters for a TDEngine database, with the websocket scheme.
-    https://docs.tdengine.com/developer-guide/connecting-to-tdengine/#websocket-connection
+    A profile that holds the required parameters for a PostgreSQL database.
+    PostgreSQL uses standard PostgreSQL connection parameters.
     """
 
-    type: str = pydantic.v1.Field("taosws")
+    type: str = pydantic.v1.Field("postgresql")
     _private_attributes = ["password"]
     user: str
     # The password cannot be empty in real world scenarios. It's here just because of the profiles completion design.
     password: typing.Optional[str]
     host: str
     port: int
+    database: str = "postgres"  # Default PostgreSQL admin database
 
-    def dsn(self) -> str:
-        """Get the Data Source Name of the configured TDEngine profile."""
-        return f"{self.type}://{self.user}:{self.password}@{self.host}:{self.port}"
+    def dsn(self, database: typing.Optional[str] = None) -> str:
+        """
+        Get the Data Source Name of the configured PostgreSQL profile.
+
+        :param database: Optional database name to use instead of the configured one.
+                        If None, uses the configured database.
+        :return: The DSN string.
+        """
+        db = database or self.database
+        # URL-encode credentials and database to handle special characters
+        user = quote(self.user, safe="")
+        password = quote(self.password or "", safe="")
+        db_encoded = quote(db, safe="")
+        return f"{self.type}://{user}:{password}@{self.host}:{self.port}/{db_encoded}"
+
+    def admin_dsn(self) -> str:
+        """
+        Get DSN for administrative operations using the 'postgres' database.
+
+        Assumes the default 'postgres' database exists (standard PostgreSQL setup).
+        Used for admin tasks like creating/dropping databases.
+
+        :return: DSN pointing to the 'postgres' database.
+        """
+        return self.dsn(database="postgres")
 
     @classmethod
-    def from_dsn(cls, dsn: str, profile_name: str) -> "DatastoreProfileTDEngine":
+    def from_dsn(cls, dsn: str, profile_name: str) -> "DatastoreProfilePostgreSQL":
         """
-        Construct a TDEngine profile from DSN (connection string) and a name for the profile.
+        Construct a PostgreSQL profile from DSN (connection string) and a name for the profile.
 
-        :param dsn:          The DSN (Data Source Name) of the TDEngine database, e.g.: ``"taosws://root:taosdata@localhost:6041"``.
+        :param dsn:          The DSN (Data Source Name) of the PostgreSQL database,
+                            e.g.: ``"postgresql://user:password@localhost:5432/mydb"``.
         :param profile_name: The new profile's name.
-        :return:             The TDEngine profile.
+        :return:             The PostgreSQL profile.
         """
         parsed_url = urlparse(dsn)
+        # URL-decode username, password, and database (urlparse doesn't decode them)
+        username = unquote(parsed_url.username) if parsed_url.username else None
+        password = unquote(parsed_url.password) if parsed_url.password else None
+        database = (
+            unquote(parsed_url.path.lstrip("/")) if parsed_url.path else "postgres"
+        )
         return cls(
             name=profile_name,
-            user=parsed_url.username,
-            password=parsed_url.password,
+            user=username,
+            password=password,
             host=parsed_url.hostname,
             port=parsed_url.port,
+            database=database or "postgres",
         )
 
 
@@ -552,7 +584,7 @@ _DATASTORE_TYPE_TO_PROFILE_CLASS: dict[str, type[DatastoreProfile]] = {
     "gcs": DatastoreProfileGCS,
     "az": DatastoreProfileAzureBlob,
     "hdfs": DatastoreProfileHdfs,
-    "taosws": DatastoreProfileTDEngine,
+    "postgresql": DatastoreProfilePostgreSQL,
     "config": ConfigProfile,
     "openai": OpenAIProfile,
     "huggingface": HuggingFaceProfile,

mlrun/datastore/model_provider/huggingface_provider.py

@@ -11,7 +11,7 @@
 # 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.
-
+import threading
 from typing import TYPE_CHECKING, Any, Optional, Union
 
 import mlrun
@@ -41,6 +41,9 @@ class HuggingFaceProvider(ModelProvider):
     into memory for inference. Ensure you have the required CPU/GPU and memory to use this operation.
     """
 
+    #  locks for threading use cases
+    _client_lock = threading.Lock()
+
     def __init__(
         self,
         parent,
@@ -224,7 +227,8 @@ class HuggingFaceProvider(ModelProvider):
 
             self.options["model_kwargs"] = self.options.get("model_kwargs", {})
             self.options["model_kwargs"]["local_files_only"] = True
-            self._client = pipeline(model=self.model, **self.options)
+            with self._client_lock:
+                self._client = pipeline(model=self.model, **self.options)
             self._expected_operation_type = Pipeline
         except ImportError as exc:
             raise ImportError("transformers package is not installed") from exc

mlrun/datastore/model_provider/model_provider.py

@@ -11,8 +11,8 @@
 # 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 collections.abc import Awaitable
-from typing import Any, Callable, Optional, Union
+from collections.abc import Awaitable, Callable
+from typing import Any, Optional, Union
 
 import mlrun.errors
 from mlrun.common.types import StrEnum

mlrun/datastore/model_provider/openai_provider.py

@@ -12,8 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import inspect
-from collections.abc import Awaitable
-from typing import TYPE_CHECKING, Any, Callable, Optional, Union
+from collections.abc import Awaitable, Callable
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 import mlrun
 from mlrun.datastore.model_provider.model_provider import (

mlrun/datastore/s3.py

@@ -18,12 +18,16 @@ from typing import Optional
 from urllib.parse import urlparse
 
 import boto3
+import botocore.exceptions
 from boto3.s3.transfer import TransferConfig
 from fsspec.registry import get_filesystem_class
 
 import mlrun.errors
 
 from .base import DataStore, FileStats, make_datastore_schema_sanitizer
+from .utils import parse_s3_bucket_and_key
+
+__all__ = ["parse_s3_bucket_and_key"]
 
 
 class S3Store(DataStore):
@@ -225,9 +229,17 @@ class S3Store(DataStore):
     def get(self, key, size=None, offset=0):
         bucket, key = self.get_bucket_and_key(key)
         obj = self.s3.Object(bucket, key)
-        if size or offset:
-            return obj.get(Range=S3Store.get_range(size, offset))["Body"].read()
-        return obj.get()["Body"].read()
+        try:
+            if size or offset:
+                return obj.get(Range=S3Store.get_range(size, offset))["Body"].read()
+            return obj.get()["Body"].read()
+
+        except botocore.exceptions.ClientError as exc:
+            if exc.response["Error"]["Code"] == "NoSuchKey":
+                # "NoSuchKey" errors codes - equivalent to `FileNotFoundError`
+                raise FileNotFoundError(f"s3://{bucket}/{key}") from exc
+            # Other errors are raised as-is
+            raise
 
     def put(self, key, data, append=False):
         data, _ = self._prepare_put_data(data, append)
@@ -259,16 +271,3 @@ class S3Store(DataStore):
         #  In order to raise an error if there is connection error, ML-7056.
         self.filesystem.exists(path=path)
         self.filesystem.rm(path=path, recursive=recursive, maxdepth=maxdepth)
-
-
-def parse_s3_bucket_and_key(s3_path):
-    try:
-        path_parts = s3_path.replace("s3://", "").split("/")
-        bucket = path_parts.pop(0)
-        key = "/".join(path_parts)
-    except Exception as exc:
-        raise mlrun.errors.MLRunInvalidArgumentError(
-            "failed to parse s3 bucket and key"
-        ) from exc
-
-    return bucket, key

mlrun/datastore/sources.py

@@ -460,7 +460,7 @@ class ParquetSource(BaseSourceDriver):
             if not filter_tuple:
                 continue
             col_name, op, value = filter_tuple
-            if op.lower() in ("in", "not in") and isinstance(value, (list, tuple, set)):
+            if op.lower() in ("in", "not in") and isinstance(value, list | tuple | set):
                 none_exists = False
                 value = list(value)
                 for sub_value in value:

mlrun/datastore/store_resources.py

@@ -76,9 +76,9 @@ class ResourceCache:
             return self._tabels[uri]
 
         if uri.startswith("v3io://") or uri.startswith("v3ios://"):
-            endpoint, uri = parse_path(uri)
+            endpoint, path = parse_path(uri)
             self._tabels[uri] = Table(
-                uri,
+                path,
                 V3ioDriver(webapi=endpoint or mlrun.mlconf.v3io_api),
                 flush_interval_secs=mlrun.mlconf.feature_store.flush_interval,
             )
@@ -87,10 +87,10 @@ class ResourceCache:
         if uri.startswith("redis://") or uri.startswith("rediss://"):
             from storey.redis_driver import RedisDriver
 
-            endpoint, uri = parse_path(uri)
+            endpoint, path = parse_path(uri)
             endpoint = endpoint or mlrun.mlconf.redis.url
             self._tabels[uri] = Table(
-                uri,
+                path,
                 RedisDriver(redis_url=endpoint, key_prefix="/"),
                 flush_interval_secs=mlrun.mlconf.feature_store.flush_interval,
             )