added as_of interface to the feature store

Author: yogesh266

ads/feature_store/common/spark_session_singleton.py

@@ -84,6 +84,7 @@ def __init__(self, metastore_id: str = None):
             )
             .enableHiveSupport()
         )
+        _managed_table_location = None
 
         if not developer_enabled() and metastore_id:
             # Get the authentication credentials for the OCI data catalog service
@@ -94,12 +95,11 @@ def __init__(self, metastore_id: str = None):
 
             data_catalog_client = OCIClientFactory(**auth).data_catalog
             metastore = data_catalog_client.get_metastore(metastore_id).data
+            _managed_table_location = metastore.default_managed_table_location
             # Configure the Spark session builder object to use the specified metastore
             spark_builder.config(
                 "spark.hadoop.oracle.dcat.metastore.id", metastore_id
-            ).config(
-                "spark.sql.warehouse.dir", metastore.default_managed_table_location
-            ).config(
+            ).config("spark.sql.warehouse.dir", _managed_table_location).config(
                 "spark.driver.memory", "16G"
             )
 
@@ -114,7 +114,12 @@ def __init__(self, metastore_id: str = None):
 
         self.spark_session.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
         self.spark_session.sparkContext.setLogLevel("OFF")
+        self.managed_table_location = _managed_table_location
 
     def get_spark_session(self):
         """Access method to get the spark session."""
         return self.spark_session
+
+    def get_managed_table_location(self):
+        """Returns the managed table location for the spark"""
+        return self.managed_table_location

ads/feature_store/common/utils/transformation_utils.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 # -*- coding: utf-8; -*-
 import json
+
 # 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/
 
@@ -64,9 +65,13 @@ def apply_transformation(
             dataframe.createOrReplaceTempView(temporary_table_view)
 
             transformed_data = spark.sql(
-                transformation_function_caller(temporary_table_view, **transformation_kwargs_dict)
+                transformation_function_caller(
+                    temporary_table_view, **transformation_kwargs_dict
+                )
             )
         elif transformation.transformation_mode == TransformationMode.PANDAS.value:
-            transformed_data = transformation_function_caller(dataframe, **transformation_kwargs_dict)
+            transformed_data = transformation_function_caller(
+                dataframe, **transformation_kwargs_dict
+            )
 
         return transformed_data

ads/feature_store/dataset.py

@@ -8,13 +8,16 @@
 
 import pandas
 from great_expectations.core import ExpectationSuite
+
+from ads import deprecated
 from ads.common import utils
 from ads.common.oci_mixin import OCIModelMixin
 from ads.feature_store.common.enums import (
     ExecutionEngine,
     ExpectationType,
     EntityType,
 )
+from ads.feature_store.common.exceptions import NotMaterializedError
 from ads.feature_store.common.utils.utility import (
     get_metastore_id,
     validate_delta_format_parameters,
@@ -475,6 +478,20 @@ def with_statistics_config(
             self.CONST_STATISTICS_CONFIG, statistics_config_in.to_dict()
         )
 
+    def target_delta_table(self):
+        """
+        Returns the fully-qualified name of the target table for storing delta data.
+
+        The name of the target table is constructed by concatenating the entity ID
+        and the name of the table, separated by a dot. The resulting string has the
+        format 'entity_id.table_name'.
+
+        Returns:
+            str: The fully-qualified name of the target delta table.
+        """
+        target_table = f"{self.entity_id}.{self.name}"
+        return target_table
+
     @property
     def model_details(self) -> "ModelDetails":
         return self.get_spec(self.CONST_MODEL_DETAILS)
@@ -560,7 +577,9 @@ def add_models(self, model_details: ModelDetails) -> "Dataset":
                 f"Dataset update Failed with : {type(ex)} with error message: {ex}"
             )
             if existing_model_details:
-                self.with_model_details(ModelDetails().with_items(existing_model_details["items"]))
+                self.with_model_details(
+                    ModelDetails().with_items(existing_model_details["items"])
+                )
             else:
                 self.with_model_details(ModelDetails().with_items([]))
                 return self
@@ -773,6 +792,7 @@ def materialise(
 
         dataset_execution_strategy.ingest_dataset(self, dataset_job)
 
+    @deprecated(details="preview functionality is deprecated. Please use as_of.")
     def preview(
         self,
         row_count: int = 10,
@@ -797,6 +817,8 @@ def preview(
         spark dataframe
             The preview result in spark dataframe
         """
+        self.check_resource_materialization()
+
         validate_delta_format_parameters(timestamp, version_number)
         target_table = f"{self.entity_id}.{self.name}"
 
@@ -806,6 +828,43 @@ def preview(
 
         return self.spark_engine.sql(sql_query)
 
+    def check_resource_materialization(self):
+        """Checks whether the target Delta table for this resource has been materialized in Spark.
+        If the target Delta table doesn't exist, raises a NotMaterializedError with the type and name of this resource.
+        """
+        if not self.spark_engine.is_delta_table_exists(self.target_delta_table()):
+            raise NotMaterializedError(self.type, self.name)
+
+    def as_of(
+        self,
+        version_number: int = None,
+        timestamp: datetime = None,
+    ):
+        """preview the feature definition and return the response in dataframe.
+
+        Parameters
+        ----------
+        timestamp: datetime
+            commit date time to preview in format yyyy-MM-dd or yyyy-MM-dd HH:mm:ss
+            commit date time is maintained for every ingestion commit using delta lake
+        version_number: int
+            commit version number for the preview. Version numbers are automatically versioned for every ingestion
+            commit using delta lake
+
+        Returns
+        -------
+        spark dataframe
+            The preview result in spark dataframe
+        """
+        self.check_resource_materialization()
+
+        validate_delta_format_parameters(timestamp, version_number)
+        target_table = self.target_delta_table()
+
+        return self.spark_engine.get_time_version_data(
+            target_table, version_number, timestamp
+        )
+
     def profile(self):
         """Get the dataset profile information and return the response in dataframe.
 
@@ -814,6 +873,8 @@ def profile(self):
         spark dataframe
             The profile result in spark dataframe
         """
+        self.check_resource_materialization()
+
         target_table = f"{self.entity_id}.{self.name}"
         sql_query = f"DESCRIBE DETAIL {target_table}"
 
@@ -835,6 +896,8 @@ def restore(self, version_number: int = None, timestamp: datetime = None):
         spark dataframe
             The restore output as spark dataframe
         """
+        self.check_resource_materialization()
+
         validate_delta_format_parameters(timestamp, version_number, True)
         target_table = f"{self.entity_id}.{self.name}"
         if version_number is not None:

ads/feature_store/execution_strategy/engine/spark_engine.py

@@ -5,6 +5,7 @@
 # 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
 
@@ -17,7 +18,7 @@
     )
 except Exception as e:
     raise
-from typing import List
+from typing import List, Dict
 
 from ads.feature_store.common.utils.feature_schema_mapper import (
     map_spark_type_to_feature_type,
@@ -36,6 +37,69 @@ def __init__(self, metastore_id: str = None, spark_session: SparkSession = None)
         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,
+    ):
+        # Get the Delta table path
+        delta_table_path = (
+            self.managed_table_location + delta_table_name
+            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,

ads/feature_store/feature_group.py

@@ -13,6 +13,7 @@
 import pandas as pd
 from great_expectations.core import ExpectationSuite
 
+from ads import deprecated
 from ads.common import utils
 from ads.common.decorator.runtime_dependency import OptionalDependency
 from ads.common.oci_mixin import OCIModelMixin
@@ -989,6 +990,7 @@ def filter(self, f: Union[Filter, Logic]):
         """
         return self.select().filter(f)
 
+    @deprecated(details="preview functionality is deprecated. Please use as_of.")
     def preview(
         self,
         row_count: int = 10,
@@ -1020,10 +1022,41 @@ def preview(
 
         if version_number is not None:
             logger.warning("Time travel queries are not supported in current version")
+
         sql_query = f"select * from {target_table} LIMIT {row_count}"
 
         return self.spark_engine.sql(sql_query)
 
+    def as_of(
+        self,
+        version_number: int = None,
+        timestamp: datetime = None,
+    ):
+        """preview the feature definition and return the response in dataframe.
+
+        Parameters
+        ----------
+        timestamp: datetime
+            commit date time to preview in format yyyy-MM-dd or yyyy-MM-dd HH:mm:ss
+            commit date time is maintained for every ingestion commit using delta lake
+        version_number: int
+            commit version number for the preview. Version numbers are automatically versioned for every ingestion
+            commit using delta lake
+
+        Returns
+        -------
+        spark dataframe
+            The preview result in spark dataframe
+        """
+        self.check_resource_materialization()
+
+        validate_delta_format_parameters(timestamp, version_number)
+        target_table = self.target_delta_table()
+
+        return self.spark_engine.get_time_version_data(
+            target_table, version_number, timestamp
+        )
+
     def profile(self):
         """get the profile information for feature definition and return the response in dataframe.
 
@@ -1085,7 +1118,6 @@ def check_resource_materialization(self):
         """Checks whether the target Delta table for this resource has been materialized in Spark.
         If the target Delta table doesn't exist, raises a NotMaterializedError with the type and name of this resource.
         """
-        print(self.target_delta_table())
         if not self.spark_engine.is_delta_table_exists(self.target_delta_table()):
             raise NotMaterializedError(self.type, self.name)