oracle-ads 2.13.8__py3-none-any.whl → 2.13.9rc0__py3-none-any.whl

ads/feature_store/execution_strategy/delta_lake/delta_lake_service.py

@@ -1,375 +0,0 @@
-#!/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

@@ -1,316 +0,0 @@
-#!/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)