mlrun 1.10.0rc16__py3-none-any.whl → 1.10.1rc4__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
 
@@ -156,6 +159,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 +357,7 @@ class DataStore(BaseRemoteClient):
         start_time,
         end_time,
         additional_filters,
+        optimize_discovery,
     ):
         from storey.utils import find_filters, find_partitions
 
@@ -203,7 +396,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, True
+                )
+                logger.debug("Partitioned parquet read", partitions=partitions)
                 set_filters(
                     partitions_time_attributes,
                     start_time,
@@ -211,8 +407,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 +454,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 +494,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 +559,7 @@ class DataStore(BaseRemoteClient):
                 start_time,
                 end_time,
                 additional_filters,
+                optimize_discovery,
             )
 
         elif file_url.endswith(".json") or format == "json":
@@ -347,7 +585,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 +625,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 +697,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 +806,7 @@ class DataItem:
             return
 
         if self._local_path:
-            remove(self._local_path)
+            os.remove(self._local_path)
             self._local_path = ""
 
     def as_df(

mlrun/datastore/datastore.py

@@ -38,6 +38,8 @@ from ..utils import DB_SCHEMA, RunKeys
 from .base import DataItem, DataStore, HttpStore
 from .filestore import FileStore
 from .inmem import InMemoryStore
+from .model_provider.huggingface_provider import HuggingFaceProvider
+from .model_provider.mock_model_provider import MockModelProvider
 from .model_provider.openai_provider import OpenAIProvider
 from .store_resources import get_store_resource, is_store_uri
 from .v3io import V3ioStore
@@ -45,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():
@@ -102,8 +104,11 @@ def schema_to_store(schema) -> DataStore.__subclasses__():
 def schema_to_model_provider(
     schema: str, raise_missing_schema_exception=True
 ) -> type[ModelProvider]:
-    #  TODO add hugging face and http
-    schema_dict = {"openai": OpenAIProvider}
+    schema_dict = {
+        "openai": OpenAIProvider,
+        "huggingface": HuggingFaceProvider,
+        "mock": MockModelProvider,
+    }
     provider_class = schema_dict.get(schema, None)
     if not provider_class:
         if raise_missing_schema_exception:
@@ -247,7 +252,7 @@ class StoreManager:
 
         if schema == "ds":
             datastore_profile = datastore_profile_read(url, project_name, secrets)
-            secrets = merge(secrets or {}, datastore_profile.secrets() or {})
+            secrets = merge({}, secrets or {}, datastore_profile.secrets() or {})
             url = datastore_profile.url(subpath)
             schema, endpoint, parsed_url = parse_url(url)
             subpath = parsed_url.path
@@ -281,7 +286,7 @@ class StoreManager:
                 endpoint, subpath
             )
             remote_client = remote_client_class(
-                self, schema, cache_key, parsed_url.netloc, secrets=secrets, **kwargs
+                self, schema, cache_key, endpoint, secrets=secrets, **kwargs
             )
             if not secrets and not mlrun.config.is_running_as_api():
                 cache[cache_key] = remote_client

mlrun/datastore/datastore_profile.py

@@ -19,6 +19,7 @@ import typing
 from urllib.parse import ParseResult, urlparse
 
 import pydantic.v1
+from deprecated import deprecated
 from mergedeep import merge
 
 import mlrun
@@ -138,6 +139,15 @@ class ConfigProfile(DatastoreProfile):
         return res
 
 
+# TODO: Remove in 1.12.0
+@deprecated(
+    version="1.10.0",
+    reason=(
+        "This class is deprecated from mlrun 1.10.0, and will be removed in 1.12.0. "
+        "Use `DatastoreProfileKafkaStream` instead."
+    ),
+    category=FutureWarning,
+)
 class DatastoreProfileKafkaTarget(DatastoreProfile):
     type: str = pydantic.v1.Field("kafka_target")
     _private_attributes = "kwargs_private"
@@ -158,8 +168,8 @@ class DatastoreProfileKafkaTarget(DatastoreProfile):
         return attributes
 
 
-class DatastoreProfileKafkaSource(DatastoreProfile):
-    type: str = pydantic.v1.Field("kafka_source")
+class DatastoreProfileKafkaStream(DatastoreProfile):
+    type: str = pydantic.v1.Field("kafka_stream")
     _private_attributes = ("kwargs_private", "sasl_user", "sasl_pass")
     brokers: typing.Union[str, list[str]]
     topics: typing.Union[str, list[str]]
@@ -198,6 +208,19 @@ class DatastoreProfileKafkaSource(DatastoreProfile):
         return attributes
 
 
+# TODO: Remove in 1.12.0
+@deprecated(
+    version="1.10.0",
+    reason=(
+        "This class is deprecated from mlrun 1.10.0, and will be removed in 1.12.0. "
+        "Use `DatastoreProfileKafkaStream` instead."
+    ),
+    category=FutureWarning,
+)
+class DatastoreProfileKafkaSource(DatastoreProfileKafkaStream):
+    type: str = pydantic.v1.Field("kafka_source")
+
+
 class DatastoreProfileV3io(DatastoreProfile):
     type: str = pydantic.v1.Field("v3io")
     v3io_access_key: typing.Optional[str] = None
@@ -232,7 +255,7 @@ class DatastoreProfileS3(DatastoreProfile):
         if self.secret_key:
             res["AWS_SECRET_ACCESS_KEY"] = self.secret_key
         if self.endpoint_url:
-            res["S3_ENDPOINT_URL"] = self.endpoint_url
+            res["AWS_ENDPOINT_URL_S3"] = self.endpoint_url
         if self.force_non_anonymous:
             res["S3_NON_ANONYMOUS"] = self.force_non_anonymous
         if self.profile_name:
@@ -333,7 +356,9 @@ class DatastoreProfileGCS(DatastoreProfile):
             #  in gcs the path after schema is starts with bucket, wherefore it should not start with "/".
             subpath = subpath[1:]
         if self.bucket:
-            return f"gcs://{self.bucket}/{subpath}"
+            return (
+                f"gcs://{self.bucket}/{subpath}" if subpath else f"gcs://{self.bucket}"
+            )
         else:
             return f"gcs://{subpath}"
 
@@ -370,7 +395,11 @@ class DatastoreProfileAzureBlob(DatastoreProfile):
             #  in azure the path after schema is starts with container, wherefore it should not start with "/".
             subpath = subpath[1:]
         if self.container:
-            return f"az://{self.container}/{subpath}"
+            return (
+                f"az://{self.container}/{subpath}"
+                if subpath
+                else f"az://{self.container}"
+            )
         else:
             return f"az://{subpath}"
 
@@ -486,6 +515,31 @@ class OpenAIProfile(DatastoreProfile):
         return f"{self.type}://{subpath.lstrip('/')}"
 
 
+class HuggingFaceProfile(DatastoreProfile):
+    type: str = pydantic.v1.Field("huggingface")
+    _private_attributes = ("token", "model_kwargs")
+    task: typing.Optional[str] = None
+    token: typing.Optional[str] = None
+    device: typing.Optional[typing.Union[int, str]] = None
+    device_map: typing.Union[str, dict[str, typing.Union[int, str]], None] = None
+    trust_remote_code: bool = None
+    model_kwargs: typing.Optional[dict[str, typing.Any]] = None
+
+    def secrets(self) -> dict:
+        keys = {
+            "HF_TASK": self.task,
+            "HF_TOKEN": self.token,
+            "HF_DEVICE": self.device,
+            "HF_DEVICE_MAP": self.device_map,
+            "HF_TRUST_REMOTE_CODE": self.trust_remote_code,
+            "HF_MODEL_KWARGS": self.model_kwargs,
+        }
+        return {k: v for k, v in keys.items() if v}
+
+    def url(self, subpath):
+        return f"{self.type}://{subpath.lstrip('/')}"
+
+
 _DATASTORE_TYPE_TO_PROFILE_CLASS: dict[str, type[DatastoreProfile]] = {
     "v3io": DatastoreProfileV3io,
     "s3": DatastoreProfileS3,
@@ -493,6 +547,7 @@ _DATASTORE_TYPE_TO_PROFILE_CLASS: dict[str, type[DatastoreProfile]] = {
     "basic": DatastoreProfileBasic,
     "kafka_target": DatastoreProfileKafkaTarget,
     "kafka_source": DatastoreProfileKafkaSource,
+    "kafka_stream": DatastoreProfileKafkaStream,
     "dbfs": DatastoreProfileDBFS,
     "gcs": DatastoreProfileGCS,
     "az": DatastoreProfileAzureBlob,
@@ -500,6 +555,7 @@ _DATASTORE_TYPE_TO_PROFILE_CLASS: dict[str, type[DatastoreProfile]] = {
     "taosws": DatastoreProfileTDEngine,
     "config": ConfigProfile,
     "openai": OpenAIProfile,
+    "huggingface": HuggingFaceProfile,
 }