mlrun 1.3.3__py3-none-any.whl → 1.4.0__py3-none-any.whl

mlrun/datastore/targets.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# 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.
@@ -15,7 +15,9 @@ import ast
 import datetime
 import os
 import random
+import sys
 import time
+import warnings
 from collections import Counter
 from copy import copy
 from typing import Any, Dict, List, Optional, Union
@@ -34,7 +36,13 @@ from mlrun.utils.v3io_clients import get_frames_client
 from .. import errors
 from ..data_types import ValueType
 from ..platforms.iguazio import parse_path, split_path
-from .utils import parse_kafka_url, store_path_to_spark
+from .utils import (
+    _generate_sql_query_with_time_filter,
+    filter_df_start_end_time,
+    parse_kafka_url,
+    select_columns_from_df,
+    store_path_to_spark,
+)
 
 
 class TargetTypes:
@@ -74,11 +82,12 @@ def default_target_names():
     return [target.strip() for target in targets.split(",")]
 
 
-def get_default_targets():
+def get_default_targets(offline_only=False):
     """initialize the default feature set targets list"""
     return [
         DataTargetBase(target, name=str(target), partitioned=(target == "parquet"))
         for target in default_target_names()
+        if not offline_only or not target == "nosql"
     ]
 
 
@@ -986,6 +995,9 @@ class CSVTarget(BaseStoreTarget):
             df_module=df_module,
             entities=entities,
             format="csv",
+            start_time=start_time,
+            end_time=end_time,
+            time_column=time_column,
             **kwargs,
         )
         if entities:
@@ -1050,24 +1062,11 @@ class NoSqlBaseTarget(BaseStoreTarget):
             **self.attributes,
         )
 
+    def prepare_spark_df(self, df, key_columns):
+        raise NotImplementedError()
+
     def get_spark_options(self, key_column=None, timestamp_key=None, overwrite=True):
-        spark_options = {
-            "path": store_path_to_spark(self.get_target_path()),
-            "format": "io.iguaz.v3io.spark.sql.kv",
-        }
-        if isinstance(key_column, list) and len(key_column) >= 1:
-            if len(key_column) > 2:
-                raise mlrun.errors.MLRunInvalidArgumentError(
-                    f"Spark supports maximun of 2 keys and {key_column} are provided"
-                )
-            spark_options["key"] = key_column[0]
-            if len(key_column) > 1:
-                spark_options["sorting-key"] = key_column[1]
-        else:
-            spark_options["key"] = key_column
-        if not overwrite:
-            spark_options["columnUpdate"] = True
-        return spark_options
+        raise NotImplementedError()
 
     def get_dask_options(self):
         return {"format": "csv"}
@@ -1075,15 +1074,6 @@ class NoSqlBaseTarget(BaseStoreTarget):
     def as_df(self, columns=None, df_module=None, **kwargs):
         raise NotImplementedError()
 
-    def prepare_spark_df(self, df, key_columns):
-        import pyspark.sql.functions as funcs
-
-        for col_name, col_type in df.dtypes:
-            if col_type.startswith("decimal("):
-                # V3IO does not support this level of precision
-                df = df.withColumn(col_name, funcs.col(col_name).cast("double"))
-        return df
-
     def write_dataframe(
         self, df, key_column=None, timestamp_key=None, chunk_id=0, **kwargs
     ):
@@ -1123,10 +1113,52 @@ class NoSqlTarget(NoSqlBaseTarget):
         endpoint, uri = parse_path(self.get_target_path())
         return Table(
             uri,
-            V3ioDriver(webapi=endpoint),
+            V3ioDriver(webapi=endpoint or mlrun.mlconf.v3io_api),
             flush_interval_secs=mlrun.mlconf.feature_store.flush_interval,
         )
 
+    def get_spark_options(self, key_column=None, timestamp_key=None, overwrite=True):
+        spark_options = {
+            "path": store_path_to_spark(self.get_target_path()),
+            "format": "io.iguaz.v3io.spark.sql.kv",
+        }
+        if isinstance(key_column, list) and len(key_column) >= 1:
+            spark_options["key"] = key_column[0]
+            if len(key_column) > 2:
+                spark_options["sorting-key"] = "_spark_object_name"
+            if len(key_column) == 2:
+                spark_options["sorting-key"] = key_column[1]
+        else:
+            spark_options["key"] = key_column
+        if not overwrite:
+            spark_options["columnUpdate"] = True
+        return spark_options
+
+    def prepare_spark_df(self, df, key_columns):
+        from pyspark.sql.functions import col
+
+        spark_udf_directory = os.path.dirname(os.path.abspath(__file__))
+        sys.path.append(spark_udf_directory)
+        try:
+            import spark_udf
+
+            df.rdd.context.addFile(spark_udf.__file__)
+
+            for col_name, col_type in df.dtypes:
+                if col_type.startswith("decimal("):
+                    # V3IO does not support this level of precision
+                    df = df.withColumn(col_name, col(col_name).cast("double"))
+            if len(key_columns) > 2:
+                return df.withColumn(
+                    "_spark_object_name",
+                    spark_udf.hash_and_concat_v3io_udf(
+                        *[col(c) for c in key_columns[1:]]
+                    ),
+                )
+        finally:
+            sys.path.remove(spark_udf_directory)
+        return df
+
 
 class RedisNoSqlTarget(NoSqlBaseTarget):
     kind = TargetTypes.redisnosql
@@ -1186,11 +1218,23 @@ class RedisNoSqlTarget(NoSqlBaseTarget):
         return endpoint
 
     def prepare_spark_df(self, df, key_columns):
-        from pyspark.sql.functions import udf
-        from pyspark.sql.types import StringType
+        from pyspark.sql.functions import col
+
+        spark_udf_directory = os.path.dirname(os.path.abspath(__file__))
+        sys.path.append(spark_udf_directory)
+        try:
+            import spark_udf
+
+            df.rdd.context.addFile(spark_udf.__file__)
 
-        udf1 = udf(lambda x: str(x) + "}:static", StringType())
-        return df.withColumn("_spark_object_name", udf1(key_columns[0]))
+            df = df.withColumn(
+                "_spark_object_name",
+                spark_udf.hash_and_concat_redis_udf(*[col(c) for c in key_columns]),
+            )
+        finally:
+            sys.path.remove(spark_udf_directory)
+
+        return df
 
 
 class StreamTarget(BaseStoreTarget):
@@ -1224,7 +1268,7 @@ class StreamTarget(BaseStoreTarget):
             graph_shape="cylinder",
             class_name="storey.StreamTarget",
             columns=column_list,
-            storage=V3ioDriver(webapi=endpoint),
+            storage=V3ioDriver(webapi=endpoint or mlrun.mlconf.v3io_api),
             stream_path=uri,
             **self.attributes,
         )
@@ -1441,7 +1485,15 @@ class DFTarget(BaseStoreTarget):
         time_column=None,
         **kwargs,
     ):
-        return self._df
+        return select_columns_from_df(
+            filter_df_start_end_time(
+                self._df,
+                time_column=time_column,
+                start_time=start_time,
+                end_time=end_time,
+            ),
+            columns,
+        )
 
 
 class SQLTarget(BaseStoreTarget):
@@ -1472,14 +1524,15 @@ class SQLTarget(BaseStoreTarget):
         # create_according_to_data: bool = False,
         time_fields: List[str] = None,
         varchar_len: int = 50,
+        parse_dates: List[str] = None,
     ):
         """
         Write to SqlDB as output target for a flow.
         example::
-             db_path = "sqlite:///stockmarket.db"
+             db_url = "sqlite:///stockmarket.db"
              schema = {'time': datetime.datetime, 'ticker': str,
                     'bid': float, 'ask': float, 'ind': int}
-             target = SqlDBTarget(table_name=f'{name}-tatget', db_path=db_path, create_table=True,
+             target = SqlDBTarget(table_name=f'{name}-target', db_url=db_url, create_table=True,
                                    schema=schema, primary_key_column=key)
         :param name:
         :param path:
@@ -1509,8 +1562,17 @@ class SQLTarget(BaseStoreTarget):
         :param create_according_to_data:    (not valid)
         :param time_fields :    all the field to be parsed as timestamp.
         :param varchar_len :    the defalut len of the all the varchar column (using if needed to create the table).
+        :param parse_dates :    all the field to be parsed as timestamp.
         """
         create_according_to_data = False  # TODO: open for user
+        if time_fields:
+            warnings.warn(
+                "'time_fields' is deprecated, use 'parse_dates' instead. "
+                "This will be removed in 1.6.0",
+                # TODO: Remove this in 1.6.0
+                FutureWarning,
+            )
+            parse_dates = time_fields
         db_url = db_url or mlrun.mlconf.sql.url
         if db_url is None or table_name is None:
             attr = {}
@@ -1523,7 +1585,7 @@ class SQLTarget(BaseStoreTarget):
                 "db_path": db_url,
                 "create_according_to_data": create_according_to_data,
                 "if_exists": if_exists,
-                "time_fields": time_fields,
+                "parse_dates": parse_dates,
                 "varchar_len": varchar_len,
             }
             path = (
@@ -1610,16 +1672,24 @@ class SQLTarget(BaseStoreTarget):
     ):
         db_path, table_name, _, _, _, _ = self._parse_url()
         engine = sqlalchemy.create_engine(db_path)
+        parse_dates: Optional[List[str]] = self.attributes.get("parse_dates")
         with engine.connect() as conn:
+            query, parse_dates = _generate_sql_query_with_time_filter(
+                table_name=table_name,
+                engine=engine,
+                time_column=time_column,
+                parse_dates=parse_dates,
+                start_time=start_time,
+                end_time=end_time,
+            )
             df = pd.read_sql(
-                f"SELECT * FROM {self.attributes.get('table_name')}",
+                query,
                 con=conn,
-                parse_dates=self.attributes.get("time_fields"),
+                parse_dates=parse_dates,
+                columns=columns,
             )
             if self._primary_key_column:
                 df.set_index(self._primary_key_column, inplace=True)
-            if columns:
-                df = df[columns]
         return df
 
     def write_dataframe(
@@ -1730,12 +1800,12 @@ def _get_target_path(driver, resource, run_id_mode=False):
     if not suffix:
         if (
             kind == ParquetTarget.kind
-            and resource.kind == mlrun.api.schemas.ObjectKind.feature_vector
+            and resource.kind == mlrun.common.schemas.ObjectKind.feature_vector
         ):
             suffix = ".parquet"
     kind_prefix = (
         "sets"
-        if resource.kind == mlrun.api.schemas.ObjectKind.feature_set
+        if resource.kind == mlrun.common.schemas.ObjectKind.feature_set
         else "vectors"
     )
     name = resource.metadata.name

mlrun/datastore/utils.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# 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.
@@ -12,7 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from urllib.parse import urlparse
+import tarfile
+import tempfile
+import typing
+from urllib.parse import parse_qs, urlparse
+
+import pandas as pd
+import sqlalchemy
+
+import mlrun.datastore
 
 
 def store_path_to_spark(path):
@@ -36,11 +44,125 @@ def store_path_to_spark(path):
     return path
 
 
-def parse_kafka_url(url, bootstrap_servers=None):
+def parse_kafka_url(
+    url: str, bootstrap_servers: typing.List = None
+) -> typing.Tuple[str, typing.List]:
+    """Generating Kafka topic and adjusting a list of bootstrap servers.
+
+    :param url:               URL path to parse using urllib.parse.urlparse.
+    :param bootstrap_servers: List of bootstrap servers for the kafka brokers.
+
+    :return: A tuple of:
+         [0] = Kafka topic value
+         [1] = List of bootstrap servers
+    """
     bootstrap_servers = bootstrap_servers or []
+
+    # Parse the provided URL into six components according to the general structure of a URL
     url = urlparse(url)
+
+    # Add the network location to the bootstrap servers list
     if url.netloc:
         bootstrap_servers = [url.netloc] + bootstrap_servers
-    topic = url.path
-    topic = topic.lstrip("/")
+
+    # Get the topic value from the parsed url
+    query_dict = parse_qs(url.query)
+    if "topic" in query_dict:
+        topic = query_dict["topic"][0]
+    else:
+        topic = url.path
+        topic = topic.lstrip("/")
     return topic, bootstrap_servers
+
+
+def upload_tarball(source_dir, target, secrets=None):
+    # will delete the temp file
+    with tempfile.NamedTemporaryFile(suffix=".tar.gz") as temp_fh:
+        with tarfile.open(mode="w:gz", fileobj=temp_fh) as tar:
+            tar.add(source_dir, arcname="")
+        stores = mlrun.datastore.store_manager.set(secrets)
+        datastore, subpath = stores.get_or_create_store(target)
+        datastore.upload(subpath, temp_fh.name)
+
+
+def filter_df_start_end_time(
+    df: typing.Union[pd.DataFrame, typing.Iterator[pd.DataFrame]],
+    time_column: str = None,
+    start_time: pd.Timestamp = None,
+    end_time: pd.Timestamp = None,
+) -> typing.Union[pd.DataFrame, typing.Iterator[pd.DataFrame]]:
+    if not time_column or (not start_time and not end_time):
+        return df
+    if isinstance(df, pd.DataFrame):
+        return _execute_time_filter(df, time_column, start_time, end_time)
+    else:
+        return filter_df_generator(df, time_column, start_time, end_time)
+
+
+def filter_df_generator(
+    dfs: typing.Iterator[pd.DataFrame],
+    time_field: str,
+    start_time: pd.Timestamp,
+    end_time: pd.Timestamp,
+) -> typing.Iterator[pd.DataFrame]:
+    for df in dfs:
+        yield _execute_time_filter(df, time_field, start_time, end_time)
+
+
+def _execute_time_filter(
+    df: pd.DataFrame, time_column: str, start_time: pd.Timestamp, end_time: pd.Timestamp
+):
+    df[time_column] = pd.to_datetime(df[time_column])
+    if start_time:
+        df = df[df[time_column] > start_time]
+    if end_time:
+        df = df[df[time_column] <= end_time]
+    return df
+
+
+def select_columns_from_df(
+    df: typing.Union[pd.DataFrame, typing.Iterator[pd.DataFrame]],
+    columns: typing.List[str],
+) -> typing.Union[pd.DataFrame, typing.Iterator[pd.DataFrame]]:
+    if not columns:
+        return df
+    if isinstance(df, pd.DataFrame):
+        return df[columns]
+    else:
+        return select_columns_generator(df, columns)
+
+
+def select_columns_generator(
+    dfs: typing.Union[pd.DataFrame, typing.Iterator[pd.DataFrame]],
+    columns: typing.List[str],
+) -> typing.Iterator[pd.DataFrame]:
+    for df in dfs:
+        yield df[columns]
+
+
+def _generate_sql_query_with_time_filter(
+    table_name: str,
+    engine: sqlalchemy.engine.Engine,
+    time_column: str,
+    parse_dates: typing.List[str],
+    start_time: pd.Timestamp,
+    end_time: pd.Timestamp,
+):
+    table = sqlalchemy.Table(
+        table_name,
+        sqlalchemy.MetaData(),
+        autoload=True,
+        autoload_with=engine,
+    )
+    query = sqlalchemy.select(table)
+    if time_column:
+        if parse_dates and time_column not in parse_dates:
+            parse_dates.append(time_column)
+        else:
+            parse_dates = [time_column]
+        if start_time:
+            query = query.filter(getattr(table.c, time_column) > start_time)
+        if end_time:
+            query = query.filter(getattr(table.c, time_column) <= end_time)
+
+    return query, parse_dates

mlrun/datastore/v3io.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# 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.

mlrun/datastore/wasbfs/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2020 Iguazio
+# 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.

mlrun/datastore/wasbfs/fs.py

@@ -1,4 +1,4 @@
-# Copyright 2022 Iguazio
+# 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.

mlrun/db/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# 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.
@@ -18,7 +18,7 @@ from ..config import config
 from ..platforms import add_or_refresh_credentials
 from ..utils import logger
 from .base import RunDBError, RunDBInterface  # noqa
-from .filedb import FileRunDB
+from .nopdb import NopDB
 from .sqldb import SQLDB
 
 
@@ -69,12 +69,14 @@ def get_run_db(url="", secrets=None, force_reconnect=False):
     kwargs = {}
     if "://" not in str(url) or scheme in ["file", "s3", "v3io", "v3ios"]:
         logger.warning(
-            "Could not detect path to API server, Using Deprecated client interface"
+            "Could not detect path to API server, not connected to API server!"
         )
         logger.warning(
-            "Please make sure your env variable MLRUN_DBPATH is configured correctly to point to the API server!"
+            "MLRUN_DBPATH is not set. Set this environment variable to the URL of the API server"
+            " in order to connect"
         )
-        cls = FileRunDB
+        cls = NopDB
+
     elif scheme in ("http", "https"):
         # import here to avoid circular imports
         from .httpdb import HTTPRunDB