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

mlrun/datastore/sources.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.
@@ -32,7 +32,12 @@ from ..config import config
 from ..model import DataSource
 from ..platforms.iguazio import parse_path
 from ..utils import get_class
-from .utils import store_path_to_spark
+from .utils import (
+    _generate_sql_query_with_time_filter,
+    filter_df_start_end_time,
+    select_columns_from_df,
+    store_path_to_spark,
+)
 
 
 def get_source_from_dict(source):
@@ -62,38 +67,59 @@ class BaseSourceDriver(DataSource):
     def to_step(self, key_field=None, time_field=None, context=None):
         import storey
 
+        if not self.support_storey:
+            raise mlrun.errors.MLRunRuntimeError(
+                f"{type(self).__name__} does not support storey engine"
+            )
+
         return storey.SyncEmitSource(context=context)
 
     def get_table_object(self):
         """get storey Table object"""
         return None
 
-    def to_dataframe(self):
-        return mlrun.store_manager.object(url=self.path).as_df()
-
-    def filter_df_start_end_time(self, df, time_field):
-        # give priority to source time_field over the feature set's timestamp_key
-        if self.time_field:
-            time_field = self.time_field
-
-        if self.start_time or self.end_time:
-            self.start_time = (
-                datetime.min if self.start_time is None else self.start_time
-            )
-            self.end_time = datetime.max if self.end_time is None else self.end_time
-            df = df.filter(
-                (df[time_field] > self.start_time) & (df[time_field] <= self.end_time)
-            )
-        return df
+    def to_dataframe(
+        self,
+        columns=None,
+        df_module=None,
+        entities=None,
+        start_time=None,
+        end_time=None,
+        time_field=None,
+    ):
+        """return the source data as dataframe"""
+        return mlrun.store_manager.object(url=self.path).as_df(
+            columns=columns,
+            df_module=df_module,
+            start_time=start_time or self.start_time,
+            end_time=end_time or self.end_time,
+            time_column=time_field or self.time_field,
+        )
 
-    def to_spark_df(self, session, named_view=False, time_field=None):
+    def to_spark_df(self, session, named_view=False, time_field=None, columns=None):
         if self.support_spark:
             df = session.read.load(**self.get_spark_options())
             if named_view:
                 df.createOrReplaceTempView(self.name)
-            return df
+            return self._filter_spark_df(df, time_field, columns)
         raise NotImplementedError()
 
+    def _filter_spark_df(self, df, time_field=None, columns=None):
+        if not (columns or time_field):
+            return df
+
+        from pyspark.sql.functions import col
+
+        if time_field:
+            if self.start_time:
+                df = df.filter(col(time_field) > self.start_time)
+            if self.end_time:
+                df = df.filter(col(time_field) <= self.end_time)
+
+        if columns:
+            df = df.select([col(name) for name in columns])
+        return df
+
     def get_spark_options(self):
         # options used in spark.read.load(**options)
         raise NotImplementedError()
@@ -166,7 +192,6 @@ class CSVSource(BaseSourceDriver):
 
         return storey.CSVSource(
             paths=self.path,
-            header=True,
             build_dict=True,
             key_field=self.key_field or key_field,
             storage_options=self._get_store().get_storage_options(),
@@ -182,7 +207,7 @@ class CSVSource(BaseSourceDriver):
             "inferSchema": "true",
         }
 
-    def to_spark_df(self, session, named_view=False, time_field=None):
+    def to_spark_df(self, session, named_view=False, time_field=None, columns=None):
         import pyspark.sql.functions as funcs
 
         df = session.read.load(**self.get_spark_options())
@@ -196,15 +221,28 @@ class CSVSource(BaseSourceDriver):
                 df = df.withColumn(col_name, funcs.col(col_name).cast("timestamp"))
         if named_view:
             df.createOrReplaceTempView(self.name)
-        return df
+        return self._filter_spark_df(df, time_field, columns)
 
-    def to_dataframe(self):
-        kwargs = self.attributes.get("reader_args", {})
-        chunksize = self.attributes.get("chunksize")
-        if chunksize:
-            kwargs["chunksize"] = chunksize
+    def to_dataframe(
+        self,
+        columns=None,
+        df_module=None,
+        entities=None,
+        start_time=None,
+        end_time=None,
+        time_field=None,
+    ):
+        reader_args = self.attributes.get("reader_args", {})
         return mlrun.store_manager.object(url=self.path).as_df(
-            parse_dates=self._parse_dates, **kwargs
+            columns=columns,
+            df_module=df_module,
+            format="csv",
+            start_time=start_time or self.start_time,
+            end_time=end_time or self.end_time,
+            time_column=time_field or self.time_field,
+            parse_dates=self._parse_dates,
+            chunksize=self.attributes.get("chunksize"),
+            **reader_args,
         )
 
     def is_iterator(self):
@@ -246,7 +284,6 @@ class ParquetSource(BaseSourceDriver):
         start_time: Optional[Union[datetime, str]] = None,
         end_time: Optional[Union[datetime, str]] = None,
     ):
-
         super().__init__(
             name,
             path,
@@ -312,10 +349,24 @@ class ParquetSource(BaseSourceDriver):
             "format": "parquet",
         }
 
-    def to_dataframe(self):
-        kwargs = self.attributes.get("reader_args", {})
+    def to_dataframe(
+        self,
+        columns=None,
+        df_module=None,
+        entities=None,
+        start_time=None,
+        end_time=None,
+        time_field=None,
+    ):
+        reader_args = self.attributes.get("reader_args", {})
         return mlrun.store_manager.object(url=self.path).as_df(
-            format="parquet", **kwargs
+            columns=columns,
+            df_module=df_module,
+            start_time=start_time or self.start_time,
+            end_time=end_time or self.end_time,
+            time_column=time_field or self.time_field,
+            format="parquet",
+            **reader_args,
         )
 
 
@@ -323,8 +374,13 @@ class BigQuerySource(BaseSourceDriver):
     """
     Reads Google BigQuery query results as input source for a flow.
 
+    For authentication, set the GCP_CREDENTIALS project secret to the credentials json string.
+
     example::
 
+         # set the credentials
+         project.set_secrets({"GCP_CREDENTIALS": gcp_credentials_json})
+
          # use sql query
          query_string = "SELECT * FROM `the-psf.pypi.downloads20210328` LIMIT 5000"
          source = BigQuerySource("bq1", query=query_string,
@@ -371,11 +427,21 @@ class BigQuerySource(BaseSourceDriver):
         end_time=None,
         gcp_project: str = None,
         spark_options: dict = None,
+        **kwargs,
     ):
         if query and table:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "cannot specify both table and query args"
             )
+        # Otherwise, the client library does not fully respect the limit
+        if (
+            max_results_for_table
+            and chunksize
+            and max_results_for_table % chunksize != 0
+        ):
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "max_results_for_table must be a multiple of chunksize"
+            )
         attrs = {
             "query": query,
             "table": table,
@@ -394,8 +460,8 @@ class BigQuerySource(BaseSourceDriver):
             schedule=schedule,
             start_time=start_time,
             end_time=end_time,
+            **kwargs,
         )
-        self._rows_iterator = None
 
     def _get_credentials_string(self):
         gcp_project = self.attributes.get("gcp_project", None)
@@ -417,7 +483,15 @@ class BigQuerySource(BaseSourceDriver):
             return credentials, gcp_project or gcp_cred_dict["project_id"]
         return None, gcp_project
 
-    def to_dataframe(self):
+    def to_dataframe(
+        self,
+        columns=None,
+        df_module=None,
+        entities=None,
+        start_time=None,
+        end_time=None,
+        time_field=None,
+    ):
         from google.cloud import bigquery
         from google.cloud.bigquery_storage_v1 import BigQueryReadClient
 
@@ -438,39 +512,43 @@ class BigQuerySource(BaseSourceDriver):
         if query:
             query_job = bqclient.query(query)
 
-            self._rows_iterator = query_job.result(page_size=chunksize)
-            dtypes = schema_to_dtypes(self._rows_iterator.schema)
-            if chunksize:
-                # passing bqstorage_client greatly improves performance
-                return self._rows_iterator.to_dataframe_iterable(
-                    bqstorage_client=BigQueryReadClient(), dtypes=dtypes
-                )
-            else:
-                return self._rows_iterator.to_dataframe(dtypes=dtypes)
+            rows_iterator = query_job.result(page_size=chunksize)
         elif table:
             table = self.attributes.get("table")
             max_results = self.attributes.get("max_results")
 
-            rows = bqclient.list_rows(
+            rows_iterator = bqclient.list_rows(
                 table, page_size=chunksize, max_results=max_results
             )
-            dtypes = schema_to_dtypes(rows.schema)
-            if chunksize:
-                # passing bqstorage_client greatly improves performance
-                return rows.to_dataframe_iterable(
-                    bqstorage_client=BigQueryReadClient(), dtypes=dtypes
-                )
-            else:
-                return rows.to_dataframe(dtypes=dtypes)
         else:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "table or query args must be specified"
             )
 
+        dtypes = schema_to_dtypes(rows_iterator.schema)
+        if chunksize:
+            # passing bqstorage_client greatly improves performance
+            df = rows_iterator.to_dataframe_iterable(
+                bqstorage_client=BigQueryReadClient(), dtypes=dtypes
+            )
+        else:
+            df = rows_iterator.to_dataframe(dtypes=dtypes)
+
+        # TODO : filter as part of the query
+        return select_columns_from_df(
+            filter_df_start_end_time(
+                df,
+                time_column=time_field or self.time_field,
+                start_time=start_time or self.start_time,
+                end_time=end_time or self.end_time,
+            ),
+            columns=columns,
+        )
+
     def is_iterator(self):
         return bool(self.attributes.get("chunksize"))
 
-    def to_spark_df(self, session, named_view=False, time_field=None):
+    def to_spark_df(self, session, named_view=False, time_field=None, columns=None):
         options = copy(self.attributes.get("spark_options", {}))
         credentials, gcp_project = self._get_credentials_string()
         if credentials:
@@ -500,7 +578,7 @@ class BigQuerySource(BaseSourceDriver):
         df = session.read.format("bigquery").load(**options)
         if named_view:
             df.createOrReplaceTempView(self.name)
-        return df
+        return self._filter_spark_df(df, time_field, columns)
 
 
 class SnowflakeSource(BaseSourceDriver):
@@ -555,6 +633,7 @@ class SnowflakeSource(BaseSourceDriver):
         database: str = None,
         schema: str = None,
         warehouse: str = None,
+        **kwargs,
     ):
         attrs = {
             "query": query,
@@ -573,6 +652,7 @@ class SnowflakeSource(BaseSourceDriver):
             schedule=schedule,
             start_time=start_time,
             end_time=end_time,
+            **kwargs,
         )
 
     def _get_password(self):
@@ -664,7 +744,7 @@ class DataFrameSource:
             context=self.context or context,
         )
 
-    def to_dataframe(self):
+    def to_dataframe(self, **kwargs):
         return self._df
 
     def is_iterator(self):
@@ -839,7 +919,15 @@ class KafkaSource(OnlineSource):
             attributes["sasl"] = sasl
         super().__init__(attributes=attributes, **kwargs)
 
-    def to_dataframe(self):
+    def to_dataframe(
+        self,
+        columns=None,
+        df_module=None,
+        entities=None,
+        start_time=None,
+        end_time=None,
+        time_field=None,
+    ):
         raise mlrun.MLRunInvalidArgumentError(
             "KafkaSource does not support batch processing"
         )
@@ -880,13 +968,15 @@ class SQLSource(BaseSourceDriver):
         table_name: str = None,
         spark_options: dict = None,
         time_fields: List[str] = None,
+        parse_dates: List[str] = None,
+        **kwargs,
     ):
         """
         Reads SqlDB as input source for a flow.
         example::
-            db_path = "mysql+pymysql://<username>:<password>@<host>:<port>/<db_name>"
-            source = SqlDBSource(
-                collection_name='source_name', db_path=self.db, key_field='key'
+            db_url = "mysql+pymysql://<username>:<password>@<host>:<port>/<db_name>"
+            source = SQLSource(
+                table_name='source_name', db_url=db_url, key_field='key'
             )
         :param name:            source name
         :param chunksize:       number of rows per chunk (default large single chunk)
@@ -903,19 +993,32 @@ class SQLSource(BaseSourceDriver):
                                 from the current database
         :param spark_options:   additional spark read options
         :param time_fields :    all the field to be parsed as timestamp.
+        :param parse_dates :    all the field to be parsed as timestamp.
         """
-
+        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:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "cannot specify without db_path arg or secret MLRUN_SQL__URL"
             )
+        if time_field:
+            if parse_dates:
+                time_fields.append(time_field)
+            else:
+                parse_dates = [time_field]
         attrs = {
             "chunksize": chunksize,
             "spark_options": spark_options,
             "table_name": table_name,
             "db_path": db_url,
-            "time_fields": time_fields,
+            "parse_dates": parse_dates,
         }
         attrs = {key: value for key, value in attrs.items() if value is not None}
         super().__init__(
@@ -926,24 +1029,43 @@ class SQLSource(BaseSourceDriver):
             schedule=schedule,
             start_time=start_time,
             end_time=end_time,
+            **kwargs,
         )
 
-    def to_dataframe(self):
-        import sqlalchemy as db
+    def to_dataframe(
+        self,
+        columns=None,
+        df_module=None,
+        entities=None,
+        start_time=None,
+        end_time=None,
+        time_field=None,
+    ):
+        import sqlalchemy as sqlalchemy
 
-        query = self.attributes.get("query", None)
         db_path = self.attributes.get("db_path")
         table_name = self.attributes.get("table_name")
-        if not query:
-            query = f"SELECT * FROM {table_name}"
+        parse_dates = self.attributes.get("parse_dates")
+        time_field = time_field or self.time_field
+        start_time = start_time or self.start_time
+        end_time = end_time or self.end_time
         if table_name and db_path:
-            engine = db.create_engine(db_path)
+            engine = sqlalchemy.create_engine(db_path)
+            query, parse_dates = _generate_sql_query_with_time_filter(
+                table_name=table_name,
+                engine=engine,
+                time_column=time_field,
+                parse_dates=parse_dates,
+                start_time=start_time,
+                end_time=end_time,
+            )
             with engine.connect() as con:
                 return pd.read_sql(
                     query,
                     con=con,
                     chunksize=self.attributes.get("chunksize"),
-                    parse_dates=self.attributes.get("time_fields"),
+                    parse_dates=parse_dates,
+                    columns=columns,
                 )
         else:
             raise mlrun.errors.MLRunInvalidArgumentError(

mlrun/datastore/spark_udf.py

@@ -0,0 +1,44 @@
+# 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 hashlib
+
+from pyspark.sql.functions import udf
+from pyspark.sql.types import StringType
+
+
+def _hash_list(*list_to_hash):
+    list_to_hash = [str(element) for element in list_to_hash]
+    str_concatted = "".join(list_to_hash)
+    sha1 = hashlib.sha1()
+    sha1.update(str_concatted.encode("utf8"))
+    return sha1.hexdigest()
+
+
+def _redis_stringify_key(*args):
+    if len(args) == 1:
+        key_list = args[0]
+    else:
+        key_list = list(args)
+    suffix = "}:static"
+    if isinstance(key_list, list):
+        if len(key_list) >= 3:
+            return str(key_list[0]) + "." + _hash_list(*key_list[1:]) + suffix
+        if len(key_list) == 2:
+            return str(key_list[0]) + "." + str(key_list[1]) + suffix
+        return str(key_list[0]) + suffix
+    return str(key_list) + suffix
+
+
+hash_and_concat_v3io_udf = udf(_hash_list, StringType())
+hash_and_concat_redis_udf = udf(_redis_stringify_key, StringType())

mlrun/datastore/store_resources.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.
@@ -81,7 +81,7 @@ class ResourceCache:
             endpoint, uri = parse_path(uri)
             self._tabels[uri] = Table(
                 uri,
-                V3ioDriver(webapi=endpoint),
+                V3ioDriver(webapi=endpoint or mlrun.mlconf.v3io_api),
                 flush_interval_secs=mlrun.mlconf.feature_store.flush_interval,
             )
             return self._tabels[uri]
@@ -101,8 +101,8 @@ class ResourceCache:
         if is_store_uri(uri):
             resource = get_store_resource(uri)
             if resource.kind in [
-                mlrun.api.schemas.ObjectKind.feature_set.value,
-                mlrun.api.schemas.ObjectKind.feature_vector.value,
+                mlrun.common.schemas.ObjectKind.feature_set.value,
+                mlrun.common.schemas.ObjectKind.feature_vector.value,
             ]:
                 target = get_online_target(resource)
                 if not target: