mlrun 1.6.4rc8__py3-none-any.whl → 1.7.0__py3-none-any.whl

mlrun/data_types/to_pandas.py

@@ -15,23 +15,13 @@
 import warnings
 from collections import Counter
 
-from pyspark.sql.types import (
-    BooleanType,
-    ByteType,
-    DoubleType,
-    FloatType,
-    IntegerType,
-    IntegralType,
-    LongType,
-    MapType,
-    ShortType,
-    TimestampType,
-)
-
-
-def toPandas(spark_df):
+import pandas as pd
+import semver
+
+
+def _to_pandas(spark_df):
     """
-    Modified version of spark DataFrame.toPandas() –
+    Modified version of spark DataFrame.toPandas() -
     https://github.com/apache/spark/blob/v3.2.3/python/pyspark/sql/pandas/conversion.py#L35
 
     The original code (which is only replaced in pyspark 3.5.0) fails with Pandas 2 installed, with the following error:
@@ -40,6 +30,12 @@ def toPandas(spark_df):
     This modification adds the missing unit to the dtype.
     """
     from pyspark.sql.dataframe import DataFrame
+    from pyspark.sql.types import (
+        BooleanType,
+        IntegralType,
+        MapType,
+        TimestampType,
+    )
 
     assert isinstance(spark_df, DataFrame)
 
@@ -48,7 +44,6 @@ def toPandas(spark_df):
     require_minimum_pandas_version()
 
     import numpy as np
-    import pandas as pd
 
     timezone = spark_df.sql_ctx._conf.sessionLocalTimeZone()
 
@@ -65,10 +60,10 @@ def toPandas(spark_df):
                 msg = (
                     "toPandas attempted Arrow optimization because "
                     "'spark.sql.execution.arrow.pyspark.enabled' is set to true; however, "
-                    "failed by the reason below:\n  %s\n"
+                    f"failed by the reason below:\n  {e}\n"
                     "Attempting non-optimization as "
                     "'spark.sql.execution.arrow.pyspark.fallback.enabled' is set to "
-                    "true." % str(e)
+                    "true."
                 )
                 warnings.warn(msg)
                 use_arrow = False
@@ -78,7 +73,7 @@ def toPandas(spark_df):
                     "'spark.sql.execution.arrow.pyspark.enabled' is set to true, but has "
                     "reached the error below and will not continue because automatic fallback "
                     "with 'spark.sql.execution.arrow.pyspark.fallback.enabled' has been set to "
-                    "false.\n  %s" % str(e)
+                    f"false.\n  {e}"
                 )
                 warnings.warn(msg)
                 raise
@@ -94,9 +89,7 @@ def toPandas(spark_df):
                 )
 
                 # Rename columns to avoid duplicated column names.
-                tmp_column_names = [
-                    "col_{}".format(i) for i in range(len(spark_df.columns))
-                ]
+                tmp_column_names = [f"col_{i}" for i in range(len(spark_df.columns))]
                 self_destruct = spark_df.sql_ctx._conf.arrowPySparkSelfDestructEnabled()
                 batches = spark_df.toDF(*tmp_column_names)._collect_as_arrow(
                     split_batches=self_destruct
@@ -146,7 +139,7 @@ def toPandas(spark_df):
                     "reached the error below and can not continue. Note that "
                     "'spark.sql.execution.arrow.pyspark.fallback.enabled' does not have an "
                     "effect on failures in the middle of "
-                    "computation.\n  %s" % str(e)
+                    f"computation.\n  {e}"
                 )
                 warnings.warn(msg)
                 raise
@@ -156,10 +149,10 @@ def toPandas(spark_df):
     column_counter = Counter(spark_df.columns)
 
     dtype = [None] * len(spark_df.schema)
-    for fieldIdx, field in enumerate(spark_df.schema):
+    for field_idx, field in enumerate(spark_df.schema):
         # For duplicate column name, we use `iloc` to access it.
         if column_counter[field.name] > 1:
-            pandas_col = pdf.iloc[:, fieldIdx]
+            pandas_col = pdf.iloc[:, field_idx]
         else:
             pandas_col = pdf[field.name]
 
@@ -173,12 +166,12 @@ def toPandas(spark_df):
             and field.nullable
             and pandas_col.isnull().any()
         ):
-            dtype[fieldIdx] = pandas_type
+            dtype[field_idx] = pandas_type
         # Ensure we fall back to nullable numpy types, even when whole column is null:
         if isinstance(field.dataType, IntegralType) and pandas_col.isnull().any():
-            dtype[fieldIdx] = np.float64
+            dtype[field_idx] = np.float64
         if isinstance(field.dataType, BooleanType) and pandas_col.isnull().any():
-            dtype[fieldIdx] = object
+            dtype[field_idx] = object
 
     df = pd.DataFrame()
     for index, t in enumerate(dtype):
@@ -219,22 +212,59 @@ def toPandas(spark_df):
 
 def _to_corrected_pandas_type(dt):
     import numpy as np
+    from pyspark.sql.types import (
+        BooleanType,
+        ByteType,
+        DoubleType,
+        FloatType,
+        IntegerType,
+        LongType,
+        ShortType,
+        TimestampType,
+    )
 
-    if type(dt) == ByteType:
+    if isinstance(dt, ByteType):
         return np.int8
-    elif type(dt) == ShortType:
+    elif isinstance(dt, ShortType):
         return np.int16
-    elif type(dt) == IntegerType:
+    elif isinstance(dt, IntegerType):
         return np.int32
-    elif type(dt) == LongType:
+    elif isinstance(dt, LongType):
         return np.int64
-    elif type(dt) == FloatType:
+    elif isinstance(dt, FloatType):
         return np.float32
-    elif type(dt) == DoubleType:
+    elif isinstance(dt, DoubleType):
         return np.float64
-    elif type(dt) == BooleanType:
+    elif isinstance(dt, BooleanType):
         return bool
-    elif type(dt) == TimestampType:
+    elif isinstance(dt, TimestampType):
         return "datetime64[ns]"
     else:
         return None
+
+
+def spark_df_to_pandas(spark_df):
+    # as of pyspark 3.2.3, toPandas fails to convert timestamps unless we work around the issue
+    # when we upgrade pyspark, we should check whether this workaround is still necessary
+    # see https://stackoverflow.com/questions/76389694/transforming-pyspark-to-pandas-dataframe
+    if semver.parse(pd.__version__)["major"] >= 2:
+        import pyspark.sql.functions as pyspark_functions
+
+        type_conversion_dict = {}
+        for field in spark_df.schema.fields:
+            if str(field.dataType) == "TimestampType":
+                spark_df = spark_df.withColumn(
+                    field.name,
+                    pyspark_functions.date_format(
+                        pyspark_functions.to_timestamp(field.name),
+                        "yyyy-MM-dd'T'HH:mm:ss.SSSSSSSSS",
+                    ),
+                )
+                type_conversion_dict[field.name] = "datetime64[ns]"
+
+        df = _to_pandas(spark_df)
+        if type_conversion_dict:
+            df = df.astype(type_conversion_dict)
+        return df
+    else:
+        return _to_pandas(spark_df)

mlrun/datastore/__init__.py

@@ -64,7 +64,7 @@ from .store_resources import (
     parse_store_uri,
 )
 from .targets import CSVTarget, NoSqlTarget, ParquetTarget, StreamTarget
-from .utils import parse_kafka_url
+from .utils import get_kafka_brokers_from_dict, parse_kafka_url
 
 store_manager = StoreManager()
 
@@ -107,19 +107,17 @@ def get_stream_pusher(stream_path: str, **kwargs):
     :param stream_path:        path/url of stream
     """
 
-    if stream_path.startswith("kafka://") or "kafka_bootstrap_servers" in kwargs:
-        topic, bootstrap_servers = parse_kafka_url(
-            stream_path, kwargs.get("kafka_bootstrap_servers")
-        )
-        return KafkaOutputStream(
-            topic, bootstrap_servers, kwargs.get("kafka_producer_options")
-        )
+    kafka_brokers = get_kafka_brokers_from_dict(kwargs)
+    if stream_path.startswith("kafka://") or kafka_brokers:
+        topic, brokers = parse_kafka_url(stream_path, kafka_brokers)
+        return KafkaOutputStream(topic, brokers, kwargs.get("kafka_producer_options"))
     elif stream_path.startswith("http://") or stream_path.startswith("https://"):
         return HTTPOutputStream(stream_path=stream_path)
     elif "://" not in stream_path:
         return OutputStream(stream_path, **kwargs)
     elif stream_path.startswith("v3io"):
         endpoint, stream_path = parse_path(stream_path)
+        endpoint = kwargs.pop("endpoint", None) or endpoint
         return OutputStream(stream_path, endpoint=endpoint, **kwargs)
     elif stream_path.startswith("dummy://"):
         return _DummyStream(**kwargs)

mlrun/datastore/alibaba_oss.py

@@ -0,0 +1,131 @@
+# Copyright 2023 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.
+
+import time
+from datetime import datetime
+from pathlib import Path
+from urllib.parse import urlparse
+
+import oss2
+from fsspec.registry import get_filesystem_class
+
+import mlrun.errors
+
+from .base import DataStore, FileStats, make_datastore_schema_sanitizer
+
+
+class OSSStore(DataStore):
+    using_bucket = True
+
+    def __init__(self, parent, schema, name, endpoint="", secrets: dict = None):
+        super().__init__(parent, name, schema, endpoint, secrets)
+        # will be used in case user asks to assume a role and work through fsspec
+
+        access_key_id = self._get_secret_or_env("ALIBABA_ACCESS_KEY_ID")
+        secret_key = self._get_secret_or_env("ALIBABA_SECRET_ACCESS_KEY")
+        endpoint_url = self._get_secret_or_env("ALIBABA_ENDPOINT_URL")
+        if access_key_id and secret_key and endpoint_url:
+            self.auth = oss2.Auth(access_key_id, secret_key)
+            self.endpoint_url = endpoint_url
+        else:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "missing ALIBABA_ACCESS_KEY_ID or ALIBABA_SECRET_ACCESS_KEY ALIBABA_ENDPOINT_URL in environment"
+            )
+
+    @property
+    def filesystem(self):
+        """return fsspec file system object, if supported"""
+        if self._filesystem:
+            return self._filesystem
+        try:
+            import ossfs  # noqa
+        except ImportError as exc:
+            raise ImportError("ALIBABA ossfs not installed") from exc
+        filesystem_class = get_filesystem_class(protocol=self.kind)
+        self._filesystem = make_datastore_schema_sanitizer(
+            filesystem_class,
+            using_bucket=self.using_bucket,
+            **self.get_storage_options(),
+        )
+        return self._filesystem
+
+    def get_storage_options(self):
+        res = dict(
+            endpoint=self._get_secret_or_env("ALIBABA_ENDPOINT_URL"),
+            key=self._get_secret_or_env("ALIBABA_ACCESS_KEY_ID"),
+            secret=self._get_secret_or_env("ALIBABA_SECRET_ACCESS_KEY"),
+        )
+        return self._sanitize_storage_options(res)
+
+    def get_bucket_and_key(self, key):
+        path = self._join(key)[1:]
+        return self.endpoint, path
+
+    def upload(self, key, src_path):
+        bucket, key = self.get_bucket_and_key(key)
+        oss = oss2.Bucket(self.auth, self.endpoint_url, bucket)
+        oss.put_object(key, open(src_path, "rb"))
+
+    def get(self, key, size=None, offset=0):
+        bucket, key = self.get_bucket_and_key(key)
+        oss = oss2.Bucket(self.auth, self.endpoint_url, bucket)
+        if size or offset:
+            return oss.get_object(key, byte_range=self.get_range(size, offset)).read()
+        return oss.get_object(key).read()
+
+    def put(self, key, data, append=False):
+        data, _ = self._prepare_put_data(data, append)
+        bucket, key = self.get_bucket_and_key(key)
+        oss = oss2.Bucket(self.auth, self.endpoint_url, bucket)
+        oss.put_object(key, data)
+
+    def stat(self, key):
+        bucket, key = self.get_bucket_and_key(key)
+        oss = oss2.Bucket(self.auth, self.endpoint_url, bucket)
+        obj = oss.get_object_meta(key)
+        size = obj.content_length
+        modified = datetime.fromtimestamp(obj.last_modified)
+        return FileStats(size, time.mktime(modified.timetuple()))
+
+    def listdir(self, key):
+        remote_path = self._convert_key_to_remote_path(key)
+        if self.filesystem.isfile(remote_path):
+            return key
+        remote_path = f"{remote_path}/**"
+        files = self.filesystem.glob(remote_path)
+        key_length = len(key)
+        files = [
+            f.split("/", 1)[1][key_length:] for f in files if len(f.split("/")) > 1
+        ]
+        return files
+
+    def delete(self, key):
+        bucket, key = self.get_bucket_and_key(key)
+        oss = oss2.Bucket(self.auth, self.endpoint_url, bucket)
+        oss.delete_object(key)
+
+    def _convert_key_to_remote_path(self, key):
+        key = key.strip("/")
+        schema = urlparse(key).scheme
+        #  if called without passing dataitem - like in fset.purge_targets,
+        #  key will include schema.
+        if not schema:
+            key = Path(self.endpoint, key).as_posix()
+        return key
+
+    @staticmethod
+    def get_range(size, offset):
+        if size:
+            return [offset, size]
+        return [offset, None]

mlrun/datastore/azure_blob.py

@@ -16,12 +16,13 @@ import time
 from pathlib import Path
 from urllib.parse import urlparse
 
+from azure.storage.blob import BlobServiceClient
 from azure.storage.blob._shared.base_client import parse_connection_str
 from fsspec.registry import get_filesystem_class
 
 import mlrun.errors
 
-from .base import DataStore, FileStats, makeDatastoreSchemaSanitizer
+from .base import DataStore, FileStats, make_datastore_schema_sanitizer
 
 # Azure blobs will be represented with the following URL: az://<container name>. The storage account is already
 # pointed to by the connection string, so the user is not expected to specify it in any way.
@@ -29,47 +30,131 @@ from .base import DataStore, FileStats, makeDatastoreSchemaSanitizer
 
 class AzureBlobStore(DataStore):
     using_bucket = True
+    max_concurrency = 100
+    max_blocksize = 1024 * 1024 * 4
+    max_single_put_size = (
+        1024 * 1024 * 8
+    )  # for service_client property only, does not affect filesystem
 
     def __init__(self, parent, schema, name, endpoint="", secrets: dict = None):
         super().__init__(parent, name, schema, endpoint, secrets=secrets)
+        self._service_client = None
+        self._storage_options = None
+
+    def get_storage_options(self):
+        return self.storage_options
+
+    @property
+    def storage_options(self):
+        if not self._storage_options:
+            res = dict(
+                account_name=self._get_secret_or_env("account_name")
+                or self._get_secret_or_env("AZURE_STORAGE_ACCOUNT_NAME"),
+                account_key=self._get_secret_or_env("account_key")
+                or self._get_secret_or_env("AZURE_STORAGE_ACCOUNT_KEY"),
+                connection_string=self._get_secret_or_env("connection_string")
+                or self._get_secret_or_env("AZURE_STORAGE_CONNECTION_STRING"),
+                tenant_id=self._get_secret_or_env("tenant_id")
+                or self._get_secret_or_env("AZURE_STORAGE_TENANT_ID"),
+                client_id=self._get_secret_or_env("client_id")
+                or self._get_secret_or_env("AZURE_STORAGE_CLIENT_ID"),
+                client_secret=self._get_secret_or_env("client_secret")
+                or self._get_secret_or_env("AZURE_STORAGE_CLIENT_SECRET"),
+                sas_token=self._get_secret_or_env("sas_token")
+                or self._get_secret_or_env("AZURE_STORAGE_SAS_TOKEN"),
+                credential=self._get_secret_or_env("credential"),
+            )
+            self._storage_options = self._sanitize_storage_options(res)
+        return self._storage_options
 
     @property
     def filesystem(self):
         """return fsspec file system object, if supported"""
-        if self._filesystem:
-            return self._filesystem
         try:
             import adlfs  # noqa
         except ImportError as exc:
             raise ImportError("Azure adlfs not installed") from exc
-        # in order to support az and wasbs kinds.
-        filesystem_class = get_filesystem_class(protocol=self.kind)
-        self._filesystem = makeDatastoreSchemaSanitizer(
-            filesystem_class,
-            using_bucket=self.using_bucket,
-            **self.get_storage_options(),
-        )
+
+        if not self._filesystem:
+            # in order to support az and wasbs kinds
+            filesystem_class = get_filesystem_class(protocol=self.kind)
+            self._filesystem = make_datastore_schema_sanitizer(
+                filesystem_class,
+                using_bucket=self.using_bucket,
+                blocksize=self.max_blocksize,
+                **self.storage_options,
+            )
         return self._filesystem
 
-    def get_storage_options(self):
-        res = dict(
-            account_name=self._get_secret_or_env("account_name")
-            or self._get_secret_or_env("AZURE_STORAGE_ACCOUNT_NAME"),
-            account_key=self._get_secret_or_env("account_key")
-            or self._get_secret_or_env("AZURE_STORAGE_KEY"),
-            connection_string=self._get_secret_or_env("connection_string")
-            or self._get_secret_or_env("AZURE_STORAGE_CONNECTION_STRING"),
-            tenant_id=self._get_secret_or_env("tenant_id")
-            or self._get_secret_or_env("AZURE_STORAGE_TENANT_ID"),
-            client_id=self._get_secret_or_env("client_id")
-            or self._get_secret_or_env("AZURE_STORAGE_CLIENT_ID"),
-            client_secret=self._get_secret_or_env("client_secret")
-            or self._get_secret_or_env("AZURE_STORAGE_CLIENT_SECRET"),
-            sas_token=self._get_secret_or_env("sas_token")
-            or self._get_secret_or_env("AZURE_STORAGE_SAS_TOKEN"),
-            credential=self._get_secret_or_env("credential"),
-        )
-        return self._sanitize_storage_options(res)
+    @property
+    def service_client(self):
+        try:
+            import azure  # noqa
+        except ImportError as exc:
+            raise ImportError("Azure not installed") from exc
+
+        if not self._service_client:
+            self._do_connect()
+        return self._service_client
+
+    def _do_connect(self):
+        """
+
+        Creates a client for azure.
+        Raises MLRunInvalidArgumentError if none of the connection details are available
+        based on do_connect in AzureBlobFileSystem:
+        https://github.com/fsspec/adlfs/blob/2023.9.0/adlfs/spec.py#L422
+        """
+        from azure.identity import ClientSecretCredential
+
+        storage_options = self.storage_options
+        connection_string = storage_options.get("connection_string")
+        client_name = storage_options.get("account_name")
+        account_key = storage_options.get("account_key")
+        sas_token = storage_options.get("sas_token")
+        client_id = storage_options.get("client_id")
+        credential = storage_options.get("credential")
+
+        credential_from_client_id = None
+        if (
+            credential is None
+            and account_key is None
+            and sas_token is None
+            and client_id is not None
+        ):
+            credential_from_client_id = ClientSecretCredential(
+                tenant_id=storage_options.get("tenant_id"),
+                client_id=client_id,
+                client_secret=storage_options.get("client_secret"),
+            )
+        try:
+            if connection_string is not None:
+                self._service_client = BlobServiceClient.from_connection_string(
+                    conn_str=connection_string,
+                    max_block_size=self.max_blocksize,
+                    max_single_put_size=self.max_single_put_size,
+                )
+            elif client_name is not None:
+                account_url = f"https://{client_name}.blob.core.windows.net"
+                cred = credential_from_client_id or credential or account_key
+                if not cred and sas_token is not None:
+                    if not sas_token.startswith("?"):
+                        sas_token = f"?{sas_token}"
+                    account_url = account_url + sas_token
+                self._service_client = BlobServiceClient(
+                    account_url=account_url,
+                    credential=cred,
+                    max_block_size=self.max_blocksize,
+                    max_single_put_size=self.max_single_put_size,
+                )
+            else:
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    "Must provide either a connection_string or account_name with credentials"
+                )
+        except Exception as e:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                f"unable to connect to account for {e}"
+            )
 
     def _convert_key_to_remote_path(self, key):
         key = key.strip("/")
@@ -82,7 +167,15 @@ class AzureBlobStore(DataStore):
 
     def upload(self, key, src_path):
         remote_path = self._convert_key_to_remote_path(key)
-        self.filesystem.put_file(src_path, remote_path, overwrite=True)
+        container, remote_path = remote_path.split("/", 1)
+        container_client = self.service_client.get_container_client(container=container)
+        with open(file=src_path, mode="rb") as data:
+            container_client.upload_blob(
+                name=remote_path,
+                data=data,
+                overwrite=True,
+                max_concurrency=self.max_concurrency,
+            )
 
     def get(self, key, size=None, offset=0):
         remote_path = self._convert_key_to_remote_path(key)
@@ -96,12 +189,7 @@ class AzureBlobStore(DataStore):
                 "Append mode not supported for Azure blob datastore"
             )
         remote_path = self._convert_key_to_remote_path(key)
-        if isinstance(data, bytes):
-            mode = "wb"
-        elif isinstance(data, str):
-            mode = "w"
-        else:
-            raise TypeError("Data type unknown.  Unable to put in Azure!")
+        data, mode = self._prepare_put_data(data, append)
         with self.filesystem.open(remote_path, mode) as f:
             f.write(data)
 
@@ -135,7 +223,7 @@ class AzureBlobStore(DataStore):
 
     def get_spark_options(self):
         res = {}
-        st = self.get_storage_options()
+        st = self.storage_options
         service = "blob"
         primary_url = None
         if st.get("connection_string"):
@@ -158,18 +246,17 @@ class AzureBlobStore(DataStore):
                     st[key] = parsed_value
 
         account_name = st.get("account_name")
-        if not account_name:
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                "Property 'account_name' is absent both in storage settings and connection string"
-            )
         if primary_url:
             if primary_url.startswith("http://"):
                 primary_url = primary_url[len("http://") :]
             if primary_url.startswith("https://"):
                 primary_url = primary_url[len("https://") :]
             host = primary_url
-        else:
+        elif account_name:
             host = f"{account_name}.{service}.core.windows.net"
+        else:
+            return res
+
         if "account_key" in st:
             res[f"spark.hadoop.fs.azure.account.key.{host}"] = st["account_key"]
 
@@ -199,3 +286,17 @@ class AzureBlobStore(DataStore):
             )
             res[f"spark.hadoop.fs.azure.sas.fixed.token.{host}"] = st["sas_token"]
         return res
+
+    @property
+    def spark_url(self):
+        spark_options = self.get_spark_options()
+        url = f"wasbs://{self.endpoint}"
+        prefix = "spark.hadoop.fs.azure.account.key."
+        if spark_options:
+            for key in spark_options:
+                if key.startswith(prefix):
+                    account_key = key[len(prefix) :]
+                    if not url.endswith(account_key):
+                        url += f"@{account_key}"
+                    break
+        return url