oracle-ads 2.13.9rc0__py3-none-any.whl → 2.13.10rc0__py3-none-any.whl

ads/feature_store/execution_strategy/delta_lake/delta_lake_service.py

@@ -0,0 +1,375 @@
+#!/usr/bin/env python
+# -*- coding: utf-8; -*-
+
+# Copyright (c) 2023 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+
+import logging
+
+from ads.common.decorator.runtime_dependency import OptionalDependency
+from ads.feature_store.common.enums import BatchIngestionMode
+from ads.feature_store.execution_strategy.engine.spark_engine import SparkEngine
+
+try:
+    from delta.tables import *
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        f"The `delta` module was not found. Please run `pip install "
+        f"{OptionalDependency.SPARK}` to install delta and spark."
+    )
+except Exception as e:
+    raise
+
+logger = logging.getLogger(__name__)
+
+
+class DeltaLakeService:
+    target_delta_table_alias = "target_delta_table"
+    source_delta_table_alias = "source_delta_table"
+    DELTA_SCHEMA_EVOLUTION_OPTIONS = ["mergeSchema", "overwriteSchema"]
+
+    def __init__(self, spark_session: SparkSession):
+        self._spark_session = spark_session
+        self.spark_engine = SparkEngine(spark_session=self._spark_session)
+
+    def write_dataframe_to_delta_lake(
+        self,
+        dataflow_output,
+        target_table_name,
+        delta_table_primary_key,
+        partition_keys,
+        ingestion_mode,
+        raw_schema,
+        feature_options=None,
+    ):
+        """Writes the given data flow output to the Delta table.
+
+        Args:
+            dataflow_output (DataFrame): The data frame that needs to be written to the Delta table.
+            target_table_name (str): The name of the target Delta table.
+            delta_table_primary_key (List[dict]): The list of primary keys for the target Delta table.
+            partition_keys(List[dict]): The List of partition Keys.
+            ingestion_mode (str): The ingestion mode for the data load.
+            raw_schema (StructType): The schema of the raw data being ingested.
+            feature_options (Dict[str, Union[str, int, float, bool]]): Optional. The dictionary containing feature options.
+
+        Returns:
+            None.
+        """
+        logger.info(f"target table name {target_table_name}")
+
+        if (
+            self.spark_engine.is_delta_table_exists(target_table_name)
+            and ingestion_mode.upper() == BatchIngestionMode.UPSERT.value
+        ):
+            logger.info(f"Upsert ops for target table {target_table_name} begin")
+
+            if raw_schema is not None:
+                # Get the source and target table columns
+                source_table_columns = raw_schema.names
+                target_table_columns = [
+                    column_details.get("name")
+                    for column_details in self.spark_engine.get_columns_from_table(
+                        target_table_name
+                    )
+                ]
+
+                logger.info(f"source table columns {source_table_columns}")
+                logger.info(f"target table columns {target_table_columns}")
+
+                if all(
+                    feature in target_table_columns for feature in source_table_columns
+                ):
+                    logger.info(
+                        f"execute upsert for select columns  {source_table_columns}"
+                    )
+
+                    self.__execute_delta_merge_insert_update(
+                        delta_table_primary_key,
+                        target_table_name,
+                        dataflow_output,
+                        source_table_columns,
+                        feature_options,
+                    )
+                else:
+                    logger.info(
+                        f"execute upsert for all columns {source_table_columns}"
+                    )
+                    self.__execute_delta_merge_insert_update_all(
+                        delta_table_primary_key,
+                        target_table_name,
+                        dataflow_output,
+                        feature_options,
+                    )
+            else:
+                self.__execute_delta_merge_insert_update_all(
+                    delta_table_primary_key,
+                    target_table_name,
+                    dataflow_output,
+                    feature_options,
+                )
+            logger.info(f"Upsert ops for target table {target_table_name} ended")
+        else:
+            self.save_delta_dataframe(
+                dataflow_output,
+                ingestion_mode,
+                target_table_name,
+                feature_options,
+                partition_keys,
+            )
+
+    def __execute_delta_merge_insert_update(
+        self,
+        primary_keys,
+        target_table_name,
+        source_dataframe,
+        source_table_columns,
+        feature_options,
+    ):
+        """Executes a Delta merge, insert, or update operation based on the given parameters.
+
+        Parameters
+        ----------
+        primary_keys
+            A list of primary key column names used for the merge operation.
+        target_table_name
+            The name of the target Delta table.
+        source_dataframe
+            The DataFrame representing the source data for the merge operation.
+        source_table_columns
+            A list of column names for the source table.
+        feature_options
+            A dictionary containing feature options.
+        """
+
+        # Enable the schema evolution
+        self.__enable_schema_evolution(feature_options)
+
+        target_delta_table = DeltaTable.forName(self._spark_session, target_table_name)
+        source_dataframe.registerTempTable(self.source_delta_table_alias)
+
+        insert_update_set = self.__get_insert_update_query_expression(
+            source_table_columns, self.source_delta_table_alias
+        )
+
+        on_condition = self.__get_delta_table_on_condition(
+            self.target_delta_table_alias, self.source_delta_table_alias, primary_keys
+        )
+
+        return (
+            target_delta_table.alias(self.target_delta_table_alias)
+            .merge(
+                source_dataframe.alias(self.source_delta_table_alias), f"{on_condition}"
+            )
+            .whenMatchedUpdate(set=insert_update_set)
+            .whenNotMatchedInsert(values=insert_update_set)
+            .execute()
+        )
+
+    def __execute_delta_merge_insert_update_all(
+        self, primary_keys, target_table_name, dataframe, feature_options=None
+    ):
+        """Perform a merge operation on Delta tables to update and insert rows into the target table. This method
+        merges all columns from the source table, irrespective of whether they exist in the target table.
+
+        Args:
+            primary_keys (List[dict]): A dictionary containing the primary keys for the target Delta table. The dictionary
+                should have a single key 'items', which maps to a list of dictionaries containing the names of the
+                primary key columns.
+            target_table_name (str): The name of the target Delta table.
+            dataframe (DataFrame): The DataFrame containing the data to merge into the target Delta table.
+            feature_options (dict): A dictionary containing feature options.
+
+        Returns:
+            None
+        """
+        # Enable the schema
+        self.__enable_schema_evolution(feature_options)
+        target_delta_table = DeltaTable.forName(self._spark_session, target_table_name)
+        on_condition = self.__get_delta_table_on_condition(
+            self.target_delta_table_alias, self.source_delta_table_alias, primary_keys
+        )
+
+        return (
+            target_delta_table.alias(self.target_delta_table_alias)
+            .merge(dataframe.alias(self.source_delta_table_alias), f"{on_condition}")
+            .whenMatchedUpdateAll()
+            .whenNotMatchedInsertAll()
+            .execute()
+        )
+
+    def save_delta_dataframe(
+        self,
+        dataframe,
+        dataframe_ingestion_mode,
+        table_name,
+        feature_options=None,
+        partition_keys=None,
+    ):
+        """
+        Saves a DataFrame to a Delta table with the specified options.
+
+        Args:
+            dataframe (pandas.DataFrame): The DataFrame to save.
+            dataframe_ingestion_mode (str): The mode to use when ingesting the DataFrame.
+            table_name (str): The name of the Delta table to save the DataFrame to.
+            feature_options (dict): Optional feature options to use when saving the DataFrame.
+
+        """
+        delta_partition_keys = []
+
+        if partition_keys:
+            partition_keys_items = partition_keys["items"]
+            if partition_keys_items:
+                delta_partition_keys = [
+                    partition_key.get("name") for partition_key in partition_keys_items
+                ]
+
+        if feature_options and feature_options.get("featureOptionWriteConfigDetails"):
+            feature_delta_write_option_config = feature_options.get(
+                "featureOptionWriteConfigDetails"
+            )
+
+            logger.info(
+                f"feature options write config details: {feature_delta_write_option_config}"
+            )
+
+            dataframe.write.format("delta").options(
+                **self.get_delta_write_config(feature_delta_write_option_config)
+            ).mode(dataframe_ingestion_mode).partitionBy(
+                delta_partition_keys
+            ).saveAsTable(
+                table_name
+            )
+        else:
+            dataframe.write.format("delta").mode(dataframe_ingestion_mode).partitionBy(
+                delta_partition_keys
+            ).saveAsTable(table_name)
+
+    def get_delta_write_config(self, feature_delta_write_option_config):
+        """Returns a dictionary containing delta schema configuration options based on a given dictionary of feature
+        delta write options.
+
+        Parameters
+        ----------
+        feature_delta_write_option_config
+            A dictionary containing feature delta write options.
+
+        Returns
+        -------
+        A dictionary containing delta schema configuration options.
+        """
+        delta_schema_config = {}
+
+        if feature_delta_write_option_config:
+            for key in self.DELTA_SCHEMA_EVOLUTION_OPTIONS:
+                if (
+                    key in feature_delta_write_option_config
+                    and feature_delta_write_option_config[key] is not None
+                ):
+                    delta_schema_config[key] = str(
+                        feature_delta_write_option_config[key]
+                    )
+
+        return delta_schema_config
+
+    def __enable_schema_evolution(self, feature_options):
+        """Enables schema evolution for Delta tables based on the given feature options.
+
+        Parameters
+        ----------
+        feature_options
+            A dictionary containing feature options.
+
+        Returns
+        -------
+        A dictionary containing delta schema configuration options
+        """
+
+        if feature_options and feature_options.get("featureOptionWriteConfigDetails"):
+            # enable auto merge schema for the spark session
+            self._spark_session.conf.set(
+                "spark.databricks.delta.schema.autoMerge.enabled", "true"
+            )
+
+    @staticmethod
+    def __get_delta_table_on_condition(
+        target_delta_table, source_delta_table, primary_keys
+    ):
+        """Returns the ON condition for the merge operation between target_delta_table and source_delta_table,
+        based on the primary keys defined in primary_keys.
+
+        Parameters
+        ----------
+        target_delta_table
+            the name of the target delta table.
+        source_delta_table
+            the name of the source delta table.
+        primary_keys
+            primary key information of the table.
+
+        Returns
+        --------
+        str
+            The ON condition for the merge operation.
+        """
+        primary_key_items = primary_keys["items"]
+        output = " AND ".join(
+            f"{target_delta_table}.{pk['name']} = {source_delta_table}.{pk['name']}"
+            for pk in primary_key_items
+        )
+        logger.info(f"Primary key on condition: {output}")
+
+        return output
+
+    @staticmethod
+    def __get_insert_update_query_expression(feature_data_source_columns, table_name):
+        """Generates an insert/update query expression to merge data from a source dataframe to a target Delta table.
+
+        Args:
+            feature_data_source_columns (list): List of columns from the source dataframe.
+            table_name (str): Name of the target Delta table.
+
+        Returns:
+            dict: A dictionary containing the update set expressions for each feature column.
+
+        """
+        feature_data_update_set = {}
+
+        for feature_column in feature_data_source_columns:
+            target_column_field = str(table_name) + "." + feature_column
+            feature_data_update_set[feature_column] = target_column_field
+
+        logger.info(f"get_insert_update_query_expression {feature_data_update_set}")
+        return feature_data_update_set
+
+    def write_stream_dataframe_to_delta_lake(
+        self,
+        stream_dataframe,
+        target_table,
+        output_mode,
+        query_name,
+        await_termination,
+        timeout,
+        checkpoint_dir,
+        feature_option_details,
+    ):
+        if query_name is None:
+            query_name = "insert_stream_" + target_table.split(".")[1]
+
+        query = (
+            stream_dataframe.writeStream.outputMode(output_mode)
+            .format("delta")
+            .option(
+                "checkpointLocation",
+                checkpoint_dir,
+            )
+            .options(**self.get_delta_write_config(feature_option_details))
+            .queryName(query_name)
+            .toTable(target_table)
+        )
+
+        if await_termination:
+            query.awaitTermination(timeout)
+
+        return query

ads/feature_store/execution_strategy/engine/spark_engine.py

@@ -0,0 +1,316 @@
+#!/usr/bin/env python
+# -*- coding: utf-8; -*-
+
+# Copyright (c) 2023 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+
+import logging
+from datetime import datetime
+
+from ads.common.decorator.runtime_dependency import OptionalDependency
+
+try:
+    from pyspark.sql import SparkSession
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        f"The `pyspark` module was not found. Please run `pip install "
+        f"{OptionalDependency.SPARK}`."
+    )
+except Exception as e:
+    raise
+from typing import List, Dict
+
+from ads.feature_store.common.utils.feature_schema_mapper import (
+    map_spark_type_to_feature_type,
+)
+
+from ads.feature_store.common.enums import DataFrameType
+from ads.feature_store.common.spark_session_singleton import SparkSessionSingleton
+
+logger = logging.getLogger(__name__)
+
+
+class SparkEngine:
+    def __init__(self, metastore_id: str = None, spark_session: SparkSession = None):
+        if spark_session:
+            self.spark = spark_session
+        else:
+            self.spark = SparkSessionSingleton(metastore_id).get_spark_session()
+
+        self.managed_table_location = (
+            SparkSessionSingleton().get_managed_table_location()
+        )
+
+    def get_time_version_data(
+        self,
+        delta_table_name: str,
+        version_number: int = None,
+        timestamp: datetime = None,
+    ):
+        split_db_name = delta_table_name.split(".")
+
+        # Get the Delta table path
+        delta_table_path = (
+            f"{self.managed_table_location}/{split_db_name[0].lower()}.db/{split_db_name[1]}"
+            if self.managed_table_location
+            else self._get_delta_table_path(delta_table_name)
+        )
+
+        # Set read options based on version_number and timestamp
+        read_options = {}
+        if version_number is not None:
+            read_options["versionAsOf"] = version_number
+        if timestamp:
+            read_options["timestampAsOf"] = timestamp
+
+        # Load the data from the Delta table using specified read options
+        df = self._read_delta_table(delta_table_path, read_options)
+        return df
+
+    def _get_delta_table_path(self, delta_table_name: str) -> str:
+        """
+        Get the path of the Delta table using DESCRIBE EXTENDED SQL command.
+
+        Args:
+            delta_table_name (str): The name of the Delta table.
+
+        Returns:
+            str: The path of the Delta table.
+        """
+        delta_table_path = (
+            self.spark.sql(f"DESCRIBE EXTENDED {delta_table_name}")
+            .filter("col_name = 'Location'")
+            .collect()[0][1]
+        )
+        return delta_table_path
+
+    def _read_delta_table(self, delta_table_path: str, read_options: Dict):
+        """
+        Read the Delta table using specified read options.
+
+        Args:
+            delta_table_path (str): The path of the Delta table.
+            read_options (dict): Dictionary of read options for Delta table.
+
+        Returns:
+            DataFrame: The loaded DataFrame from the Delta table.
+        """
+        df = (
+            self.spark.read.format("delta")
+            .options(**read_options)
+            .load(delta_table_path)
+        )
+        return df
+
+    def sql(
+        self,
+        query: str,
+        dataframe_type: DataFrameType = DataFrameType.SPARK,
+        is_online: bool = False,
+    ):
+        """Execute SQL command on the offline or online feature store database
+
+        Arguments
+            query: The SQL query to execute.
+            dataframe_type: The type of the returned dataframe. Defaults to "default".
+            is_online: Set to true to execute the query against the online feature store.
+                Defaults to False.
+
+        Returns
+            `DataFrame`: DataFrame depending on the chosen type.
+        """
+        if is_online:
+            raise ValueError("Online query is not supported.")
+
+        response_spark_df = self.spark.sql(query)
+        response_df = (
+            response_spark_df.toPandas()
+            if dataframe_type == DataFrameType.PANDAS
+            else response_spark_df
+        )
+
+        return response_df
+
+    def is_database_exists(self, database):
+        """Checks whether the database exists or not.
+
+        Args:
+          database: A string specifying the name of the database.
+
+        Returns:
+          bool: True if the database exists, False otherwise.
+        """
+        databases = self.spark.catalog.listDatabases()
+
+        return any(db.name == database for db in databases)
+
+    def is_delta_table_exists(self, table_name) -> bool:
+        """Checks whether the delta table exists or not.
+
+        Args:
+          table_name: A string specifying the name of the table.
+
+        Returns:
+          bool: True if the Delta table exists, False otherwise.
+        """
+        table_exist = False
+
+        try:
+            # Check if spark can read the table
+            self.spark.read.table(table_name)
+            table_exist = True
+
+        except:
+            pass
+
+        return table_exist
+
+    def get_tables_from_database(self, database):
+        """Get a list of tables in the specified database using Spark SQL.
+
+        Args:
+          database: A string specifying the name of the database.
+
+        Returns:
+          List: A list of strings containing the names of the tables in the database.
+        """
+        permanent_tables = None
+
+        try:
+            tables_list = self.spark.catalog.listTables(database)
+
+            # tables_list contains temporary tables also, so we need to filter it.
+            permanent_tables = [table for table in tables_list if not table.isTemporary]
+        except Exception as e:
+            logger.error("Error: The database does not exist. ", e)
+
+        return permanent_tables
+
+    def get_output_columns_from_table_or_dataframe(
+        self, table_name: str = None, dataframe=None
+    ):
+        """Returns the column(features) along with type from the given table.
+
+        Args:
+          table_name(str): A string specifying the name of table name for which columns should be returned.
+          dataframe: Dataframe containing the transformed dataframe.
+
+        Returns:
+         List[{"name": "<feature_name>","featureType": "<feature_type>"}]
+         Returns the List of dictionary of column with name and type from the given table.
+
+        """
+        if table_name is None and dataframe is None:
+            raise ValueError(
+                "Either 'table_name' or 'dataframe' must be provided to retrieve output columns."
+            )
+
+        if dataframe is not None:
+            feature_data_target = dataframe
+        else:
+            feature_data_target = self.spark.sql(f"SELECT * FROM {table_name} LIMIT 1")
+
+        target_table_columns = []
+
+        for field in feature_data_target.schema.fields:
+            target_table_columns.append(
+                {
+                    "name": field.name,
+                    "featureType": map_spark_type_to_feature_type(field.dataType).value,
+                }
+            )
+        return target_table_columns
+
+    def convert_from_pandas_to_spark_dataframe(self, dataframe):
+        """Converts a pandas DataFrame to an Apache Spark DataFrame.
+
+        Args:
+          dataframe (pandas.DataFrame): The pandas DataFrame to convert.
+
+        Returns:
+         pyspark.sql.DataFrame: The converted Apache Spark DataFrame.
+        """
+        return self.spark.createDataFrame(dataframe)
+
+    def delete_spark_table(self, table_name: str):
+        """
+        Delete the specified Spark table from the Spark session.
+
+        Args:
+            table_name (str): The full name of the table to delete.
+
+        Returns:
+            None
+        """
+
+        # Construct SQL queries to drop the table (if it exists)
+        drop_table_query = f"DROP TABLE IF EXISTS {table_name}"
+        logger.info("Deleting the table with query: ", drop_table_query)
+        self.spark.sql(drop_table_query)
+
+    def delete_spark_database(self, database_name: str):
+        """
+        Delete the specified Spark database from the Spark session.
+
+        Args:
+            database_name (str): The name of the database to delete.
+
+        Returns:
+            None
+        """
+        # Construct SQL queries to drop the database (if it exists)
+        drop_database_query = f"DROP DATABASE IF EXISTS {database_name}"
+        self.spark.sql(drop_database_query)
+
+    def remove_table_and_database(self, database: str, table: str) -> None:
+        """
+        Remove the specified Spark table and its database (if it exists) from the current Spark session.
+
+        Args:
+            database (str): The name of the database containing the table to remove.
+            table (str): The name of the table to remove.
+
+        Returns:
+            None.
+
+        Raises:
+            None.
+
+        This method first checks if the specified database exists in the current Spark session by calling the
+        `is_database_exists()` method. If the database exists, it deletes the specified table from the database
+        by calling the `delete_spark_table()` method.
+
+        Next, the method checks if there are any tables left in the database by calling the `get_tables_from_database()` method.
+        If the method returns a list of tables, it checks if the list is empty, and if so, it deletes the database
+        by calling the `delete_spark_database()` method.
+
+        If the `get_tables_from_database()` method returns `None`, no action is taken, and the method simply returns.
+
+        If the specified database does not exist in the current Spark session, the method simply returns without doing anything.
+        """
+        if not self.is_database_exists(database):
+            return
+
+        table_name = f"{database}.{table}"
+        self.delete_spark_table(table_name)
+
+        tables_list = self.get_tables_from_database(database)
+        if tables_list is not None and len(tables_list) == 0:
+            self.delete_spark_database(database)
+
+    def create_database(self, database: str) -> None:
+        """
+        Creates a database in Spark if it does not already exist.
+
+        Args:
+            database (str): The name of the database to create.
+
+        Returns:
+            None.
+
+        """
+        # Construct a SQL query to create the database if it does not exist
+        create_database_query = f"CREATE DATABASE IF NOT EXISTS {database}"
+
+        # Use the `spark.sql()` method to execute the SQL query
+        self.spark.sql(create_database_query)