openmetadata-ingestion 1.5.13.0__py3-none-any.whl → 1.6.0.0rc2__py3-none-any.whl

metadata/profiler/interface/pandas/profiler_interface.py

@@ -16,28 +16,33 @@ supporting sqlalchemy abstraction layer
 """
 import traceback
 from collections import defaultdict
-from copy import deepcopy
 from datetime import datetime
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Union
 
 from sqlalchemy import Column
 
 from metadata.generated.schema.entity.data.table import (
     CustomMetricProfile,
     DataType,
-    TableData,
+    Table,
 )
 from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
     DatalakeConnection,
 )
+from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
+from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline import (
+    DatabaseServiceProfilerPipeline,
+)
 from metadata.generated.schema.tests.customMetric import CustomMetric
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
 from metadata.mixins.pandas.pandas_mixin import PandasInterfaceMixin
 from metadata.profiler.api.models import ThreadPoolMetrics
 from metadata.profiler.interface.profiler_interface import ProfilerInterface
 from metadata.profiler.metrics.core import MetricTypes
 from metadata.profiler.metrics.registry import Metrics
 from metadata.profiler.processor.metric_filter import MetricFilter
-from metadata.utils.constants import COMPLEX_COLUMN_SEPARATOR, SAMPLE_DATA_DEFAULT_COUNT
+from metadata.sampler.pandas.sampler import DatalakeSampler
+from metadata.utils.constants import COMPLEX_COLUMN_SEPARATOR
 from metadata.utils.datalake.datalake_utils import GenericDataFrameColumnParser
 from metadata.utils.logger import profiler_interface_registry_logger
 from metadata.utils.sqa_like_column import SQALikeColumn
@@ -55,47 +60,30 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
 
     def __init__(
         self,
-        service_connection_config,
-        ometa_client,
-        entity,
-        storage_config,
-        profile_sample_config,
-        source_config,
-        sample_query,
-        table_partition_config,
+        service_connection_config: Union[DatabaseConnection, DatalakeConnection],
+        ometa_client: OpenMetadata,
+        entity: Table,
+        source_config: DatabaseServiceProfilerPipeline,
+        sampler: DatalakeSampler,
         thread_count: int = 5,
         timeout_seconds: int = 43200,
-        sample_data_count: int = SAMPLE_DATA_DEFAULT_COUNT,
         **kwargs,
     ):
         """Instantiate Pandas Interface object"""
 
         super().__init__(
-            service_connection_config,
-            ometa_client,
-            entity,
-            storage_config,
-            profile_sample_config,
-            source_config,
-            sample_query,
-            table_partition_config,
-            thread_count,
-            timeout_seconds,
-            sample_data_count,
+            service_connection_config=service_connection_config,
+            ometa_client=ometa_client,
+            entity=entity,
+            source_config=source_config,
+            sampler=sampler,
+            thread_count=thread_count,
+            timeout_seconds=timeout_seconds,
             **kwargs,
         )
 
-        self.client = self.connection.client
-        self.dfs = self.return_ometa_dataframes_sampled(
-            service_connection_config=self.service_connection_config,
-            client=self.client._client,
-            table=self.table_entity,
-            profile_sample_config=profile_sample_config,
-        )
-        self.sampler = self._get_sampler()
-        self.complex_dataframe_sample = deepcopy(
-            self.sampler.random_sample(is_sampled=True)
-        )
+        self.client = self.sampler.client
+        self.dataset = self.sampler.get_dataset()
         self.complex_df()
 
     def complex_df(self):
@@ -104,7 +92,7 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
         data_formats = (
             GenericDataFrameColumnParser._data_formats  # pylint: disable=protected-access
         )
-        for index, df in enumerate(self.complex_dataframe_sample):
+        for index, df in enumerate(self.dataset):
             if index == 0:
                 for col in self.table.columns:
                     coltype = next(
@@ -121,29 +109,14 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
                         coltype_mapping_df.append("object")
 
             try:
-                self.complex_dataframe_sample[index] = df.astype(
+                self.dataset[index] = df.astype(
                     dict(zip(df.keys(), coltype_mapping_df))
                 )
             except (TypeError, ValueError) as err:
-                self.complex_dataframe_sample[index] = df
+                self.dataset[index] = df
                 logger.warning(f"NaN/NoneType found in the Dataframe: {err}")
                 break
 
-    def _get_sampler(self):
-        """Get dataframe sampler from config"""
-        from metadata.profiler.processor.sampler.sampler_factory import (  # pylint: disable=import-outside-toplevel
-            sampler_factory_,
-        )
-
-        return sampler_factory_.create(
-            DatalakeConnection.__name__,
-            client=self.client._client,  # pylint: disable=W0212
-            table=self.dfs,
-            profile_sample_config=self.profile_sample_config,
-            partition_details=self.partition_details,
-            profile_sample_query=self.profile_query,
-        )
-
     def _compute_table_metrics(
         self,
         metrics: List[Metrics],
@@ -152,7 +125,8 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
         **kwargs,
     ):
         """Given a list of metrics, compute the given results
-        and returns the values
+        and returns the values. Table metrics are computed on the
+        entire dataset omitting the sampling and partitioning
 
         Args:
             metrics: list of metrics to compute
@@ -163,7 +137,7 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
 
         try:
             row_dict = {}
-            df_list = [df.where(pd.notnull(df), None) for df in runner]
+            df_list = [df.where(pd.notnull(df), None) for df in self.dataset]
             for metric in metrics:
                 row_dict[metric.name()] = metric().df_fn(df_list)
             return row_dict
@@ -306,10 +280,10 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
         logger.debug(f"Running profiler for {metric_func.table.name.root}")
         try:
             row = None
-            if self.complex_dataframe_sample:
+            if self.dataset:
                 row = self._get_metric_fn[metric_func.metric_type.value](
                     metric_func.metrics,
-                    self.complex_dataframe_sample,
+                    self.dataset,
                     column=metric_func.column,
                 )
         except Exception as exc:
@@ -326,18 +300,6 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
             column = None
         return row, column, metric_func.metric_type.value
 
-    def fetch_sample_data(self, table, columns: SQALikeColumn) -> TableData:
-        """Fetch sample data from database
-
-        Args:
-            table: ORM declarative table
-
-        Returns:
-            TableData: sample table data
-        """
-        sampler = self._get_sampler()
-        return sampler.fetch_sample_data(columns)
-
     def get_composed_metrics(
         self, column: Column, metric: Metrics, column_results: Dict
     ):
@@ -358,9 +320,7 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
             logger.warning(f"Unexpected exception computing metrics: {exc}")
             return None
 
-    def get_hybrid_metrics(
-        self, column: Column, metric: Metrics, column_results: Dict, **kwargs
-    ):
+    def get_hybrid_metrics(self, column: Column, metric: Metrics, column_results: Dict):
         """Given a list of metrics, compute the given results
         and returns the values
 
@@ -372,7 +332,7 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
             dictionary of results
         """
         try:
-            return metric(column).df_fn(column_results, self.complex_dataframe_sample)
+            return metric(column).df_fn(column_results, self.dataset)
         except Exception as exc:
             logger.debug(traceback.format_exc())
             logger.warning(f"Unexpected exception computing metrics: {exc}")
@@ -417,15 +377,15 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
     def get_columns(self) -> List[Optional[SQALikeColumn]]:
         """Get SQALikeColumns for datalake to be passed for metric computation"""
         sqalike_columns = []
-        if self.complex_dataframe_sample:
-            for column_name in self.complex_dataframe_sample[0].columns:
+        if self.dataset:
+            for column_name in self.dataset[0].columns:
                 complex_col_name = None
                 if COMPLEX_COLUMN_SEPARATOR in column_name:
                     complex_col_name = ".".join(
                         column_name.split(COMPLEX_COLUMN_SEPARATOR)[1:]
                     )
                     if complex_col_name:
-                        for df in self.complex_dataframe_sample:
+                        for df in self.dataset:
                             df.rename(
                                 columns={column_name: complex_col_name}, inplace=True
                             )
@@ -434,7 +394,7 @@ class PandasProfilerInterface(ProfilerInterface, PandasInterfaceMixin):
                     SQALikeColumn(
                         column_name,
                         GenericDataFrameColumnParser.fetch_col_types(
-                            self.complex_dataframe_sample[0], column_name
+                            self.dataset[0], column_name
                         ),
                     )
                 )

metadata/profiler/interface/profiler_interface.py

@@ -15,33 +15,15 @@ supporting sqlalchemy abstraction layer
 """
 
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Type, Union
 
 from sqlalchemy import Column
 
-from metadata.generated.schema.entity.data.database import (
-    Database,
-    DatabaseProfilerConfig,
-)
-from metadata.generated.schema.entity.data.databaseSchema import (
-    DatabaseSchema,
-    DatabaseSchemaProfilerConfig,
-)
-from metadata.generated.schema.entity.data.table import (
-    PartitionProfilerConfig,
-    Table,
-    TableData,
-)
-from metadata.generated.schema.entity.services.connections.connectionBasicType import (
-    DataStorageConfig,
-)
+from metadata.generated.schema.entity.data.table import SystemProfile, Table
 from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
     DatalakeConnection,
 )
-from metadata.generated.schema.entity.services.databaseService import (
-    DatabaseConnection,
-    DatabaseService,
-)
+from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
 from metadata.generated.schema.entity.services.ingestionPipelines.status import (
     StackTraceError,
 )
@@ -51,17 +33,12 @@ from metadata.generated.schema.metadataIngestion.databaseServiceProfilerPipeline
 from metadata.generated.schema.tests.customMetric import CustomMetric
 from metadata.ingestion.api.status import Status
 from metadata.ingestion.ometa.ometa_api import OpenMetadata
-from metadata.profiler.api.models import (
-    DatabaseAndSchemaConfig,
-    ProfilerProcessorConfig,
-    ProfileSampleConfig,
-    TableConfig,
-)
 from metadata.profiler.metrics.core import MetricTypes
 from metadata.profiler.metrics.registry import Metrics
+from metadata.profiler.metrics.system.system import System
 from metadata.profiler.processor.runner import QueryRunner
-from metadata.utils.constants import SAMPLE_DATA_DEFAULT_COUNT
-from metadata.utils.partition import get_partition_details
+from metadata.sampler.sampler_interface import SamplerInterface
+from metadata.utils.collaborative_super import Root
 from metadata.utils.ssl_manager import get_ssl_connection
 
 
@@ -83,30 +60,23 @@ class ProfilerProcessorStatus(Status):
         )
 
 
-# pylint: disable=too-many-instance-attributes
-class ProfilerInterface(ABC):
+class ProfilerInterface(Root, ABC):
     """Protocol interface for the profiler processor"""
 
-    # pylint: disable=too-many-arguments,unused-argument
-    def __init__(
+    def __init__(  # pylint: disable=too-many-arguments
         self,
         service_connection_config: Union[DatabaseConnection, DatalakeConnection],
         ometa_client: OpenMetadata,
         entity: Table,
-        storage_config: DataStorageConfig,
-        profile_sample_config: Optional[ProfileSampleConfig],
         source_config: DatabaseServiceProfilerPipeline,
-        sample_query: Optional[str],
-        table_partition_config: Optional[PartitionProfilerConfig],
+        sampler: SamplerInterface,
         thread_count: int = 5,
         timeout_seconds: int = 43200,
-        sample_data_count: Optional[int] = SAMPLE_DATA_DEFAULT_COUNT,
         **kwargs,
     ):
         """Required attribute for the interface"""
         self._thread_count = thread_count
         self.table_entity = entity
-        self.storage_config = storage_config
         self.ometa_client = ometa_client
         self.source_config = source_config
         self.service_connection_config = service_connection_config
@@ -118,13 +88,8 @@ class ProfilerInterface(ABC):
             self.status.entity = None
         else:
             self.status.entity = fqn.root if fqn else None
-        self.profile_sample_config = profile_sample_config
-        self.profile_query = sample_query
-        self.partition_details = (
-            table_partition_config if not self.profile_query else None
-        )
+        self.sampler = sampler
         self.timeout_seconds = timeout_seconds
-        self.sample_data_count = sample_data_count
 
         self._get_metric_fn = {
             MetricTypes.Table.value: self._compute_table_metrics,
@@ -135,23 +100,24 @@ class ProfilerInterface(ABC):
             MetricTypes.Custom.value: self._compute_custom_metrics,
         }
 
-    @abstractmethod
-    def _get_sampler(self):
-        """Get the sampler"""
-        raise NotImplementedError
+        super().__init__(
+            service_connection_config=service_connection_config,
+            ometa_client=ometa_client,
+            entity=entity,
+            source_config=source_config,
+            sampler=sampler,
+            thread_count=thread_count,
+            timeout_seconds=timeout_seconds,
+            **kwargs,
+        )
 
-    # pylint: disable=too-many-locals
     @classmethod
     def create(
         cls,
         entity: Table,
-        database_schema: DatabaseSchema,
-        database: Database,
-        database_service: DatabaseService,
-        entity_config: Optional[TableConfig],
-        profiler_config: Optional[ProfilerProcessorConfig],
         source_config: DatabaseServiceProfilerPipeline,
         service_connection_config,
+        sampler: SamplerInterface,
         ometa_client: Optional[OpenMetadata],
         **kwargs,
     ) -> "ProfilerInterface":
@@ -174,233 +140,18 @@ class ProfilerInterface(ABC):
         """
         thread_count = source_config.threadCount
         timeout_seconds = source_config.timeoutSeconds
-        database_profiler_config = cls.get_database_profiler_config(
-            database_entity=database
-        )
-        schema_profiler_config = cls.get_schema_profiler_config(
-            schema_entity=database_schema
-        )
-        storage_config = cls.get_storage_config_for_table(
-            entity=entity,
-            schema_profiler_config=schema_profiler_config,
-            database_profiler_config=database_profiler_config,
-            db_service=database_service,
-            profiler_config=profiler_config,
-        )
 
-        if not cls.get_profile_query(entity, entity_config):
-            profile_sample_config = cls.get_profile_sample_config(
-                entity,
-                schema_profiler_config,
-                database_profiler_config,
-                entity_config,
-                source_config,
-            )
-            table_partition_config = cls.get_partition_details(entity, entity_config)
-            sample_query = None
-        else:
-            sample_query = cls.get_profile_query(entity, entity_config)
-            profile_sample_config = None
-            table_partition_config = None
-
-        sample_data_count = cls.get_sample_data_count_config(
-            entity,
-            schema_profiler_config,
-            database_profiler_config,
-            entity_config,
-            source_config,
-        )
         return cls(
             service_connection_config=service_connection_config,
             ometa_client=ometa_client,
             entity=entity,
-            storage_config=storage_config,
-            profile_sample_config=profile_sample_config,
             source_config=source_config,
-            sample_query=sample_query,
-            table_partition_config=table_partition_config,
+            sampler=sampler,
             thread_count=thread_count,
             timeout_seconds=timeout_seconds,
-            sample_data_count=sample_data_count,
             **kwargs,
         )
 
-    @staticmethod
-    def get_schema_profiler_config(
-        schema_entity: Optional[DatabaseSchema],
-    ) -> DatabaseSchemaProfilerConfig:
-        if schema_entity and schema_entity.databaseSchemaProfilerConfig:
-            return schema_entity.databaseSchemaProfilerConfig
-        return None
-
-    @staticmethod
-    def get_database_profiler_config(
-        database_entity: Optional[Database],
-    ) -> DatabaseProfilerConfig:
-        if database_entity and database_entity.databaseProfilerConfig:
-            return database_entity.databaseProfilerConfig
-        return None
-
-    @staticmethod
-    def _get_sample_storage_config(
-        config: Union[
-            DatabaseSchemaProfilerConfig,
-            DatabaseProfilerConfig,
-            DatabaseAndSchemaConfig,
-        ]
-    ) -> Optional[DataStorageConfig]:
-        if (
-            config
-            and config.sampleDataStorageConfig
-            and config.sampleDataStorageConfig.config
-        ):
-            return config.sampleDataStorageConfig.config
-        return None
-
-    @staticmethod
-    def get_storage_config_for_table(
-        entity: Table,
-        schema_profiler_config: Optional[DatabaseSchemaProfilerConfig],
-        database_profiler_config: Optional[DatabaseProfilerConfig],
-        db_service: Optional[DatabaseService],
-        profiler_config: ProfilerProcessorConfig,
-    ) -> Optional[DataStorageConfig]:
-        """Get config for a specific entity
-
-        Args:
-            entity: table entity
-        """
-        for schema_config in profiler_config.schemaConfig:
-            if (
-                schema_config.fullyQualifiedName.root
-                == entity.databaseSchema.fullyQualifiedName
-                and ProfilerInterface._get_sample_storage_config(schema_config)
-            ):
-                return ProfilerInterface._get_sample_storage_config(schema_config)
-
-        for database_config in profiler_config.databaseConfig:
-            if (
-                database_config.fullyQualifiedName.root
-                == entity.database.fullyQualifiedName
-                and ProfilerInterface._get_sample_storage_config(database_config)
-            ):
-                return ProfilerInterface._get_sample_storage_config(database_config)
-
-        if ProfilerInterface._get_sample_storage_config(schema_profiler_config):
-            return ProfilerInterface._get_sample_storage_config(schema_profiler_config)
-
-        if ProfilerInterface._get_sample_storage_config(database_profiler_config):
-            return ProfilerInterface._get_sample_storage_config(
-                database_profiler_config
-            )
-
-        try:
-            return db_service.connection.config.sampleDataStorageConfig.config
-        except AttributeError:
-            pass
-
-        return None
-
-    @staticmethod
-    def get_profile_sample_config(
-        entity: Table,
-        schema_profiler_config: Optional[DatabaseSchemaProfilerConfig],
-        database_profiler_config: Optional[DatabaseProfilerConfig],
-        entity_config: Optional[Union[TableConfig, DatabaseAndSchemaConfig]],
-        source_config: DatabaseServiceProfilerPipeline,
-    ) -> Optional[ProfileSampleConfig]:
-        """_summary_
-        Args:
-            entity (Table): table entity object
-            entity_config (Optional[TableConfig]): table config object from yaml/json file
-        Returns:
-            Optional[dict]: dict
-        """
-        for config in (
-            entity_config,
-            entity.tableProfilerConfig,
-            schema_profiler_config,
-            database_profiler_config,
-            source_config,
-        ):
-            try:
-                if config and config.profileSample:
-                    return ProfileSampleConfig(
-                        profile_sample=config.profileSample,
-                        profile_sample_type=config.profileSampleType,
-                    )
-            except AttributeError:
-                pass
-
-        return None
-
-    @staticmethod
-    def get_sample_data_count_config(
-        entity: Table,
-        schema_profiler_config: Optional[DatabaseSchemaProfilerConfig],
-        database_profiler_config: Optional[DatabaseProfilerConfig],
-        entity_config: Optional[TableConfig],
-        source_config: DatabaseServiceProfilerPipeline,
-    ) -> Optional[int]:
-        """_summary_
-        Args:
-            entity_config (Optional[TableConfig]): table config object from yaml/json file
-            source_config DatabaseServiceProfilerPipeline: profiler pipeline details
-        Returns:
-            Optional[int]: int
-        """
-
-        for config in (
-            entity_config,
-            entity.tableProfilerConfig,
-            schema_profiler_config,
-            database_profiler_config,
-        ):
-            if config and config.sampleDataCount:
-                return config.sampleDataCount
-
-        return source_config.sampleDataCount
-
-    @staticmethod
-    def get_profile_query(
-        entity: Table, entity_config: Optional[TableConfig]
-    ) -> Optional[str]:
-        """get profile query for sampling
-
-        Args:
-            entity (Table): table entity object
-            entity_config (Optional[TableConfig]): entity configuration
-
-        Returns:
-            Optional[str]:
-        """
-        if entity_config:
-            return entity_config.profileQuery
-
-        if entity.tableProfilerConfig:
-            return entity.tableProfilerConfig.profileQuery
-
-        return None
-
-    @staticmethod
-    def get_partition_details(
-        entity: Table,
-        entity_config: Optional[TableConfig] = None,
-    ) -> Optional[PartitionProfilerConfig]:
-        """_summary_
-
-        Args:
-            entity (Table): table entity object
-            entity_config (Optional[TableConfig]): entity configuration
-
-        Returns:
-            Optional[PartitionProfilerConfig]:
-        """
-        if entity_config:
-            return entity_config.partitionConfig
-
-        return get_partition_details(entity)
-
     @property
     @abstractmethod
     def table(self):
@@ -458,11 +209,11 @@ class ProfilerInterface(ABC):
     @abstractmethod
     def _compute_system_metrics(
         self,
-        metrics: Metrics,
+        metrics: Type[System],
         runner,
         *args,
         **kwargs,
-    ):
+    ) -> List[SystemProfile]:
         """Get metrics"""
         raise NotImplementedError
 
@@ -487,16 +238,11 @@ class ProfilerInterface(ABC):
 
     @abstractmethod
     def get_hybrid_metrics(
-        self, column: Column, metric: Metrics, column_results: Dict, **kwargs
+        self, column: Column, metric: Metrics, column_results: Dict
     ) -> dict:
         """run profiler metrics"""
         raise NotImplementedError
 
-    @abstractmethod
-    def fetch_sample_data(self, table, columns: List[Column]) -> TableData:
-        """run profiler metrics"""
-        raise NotImplementedError
-
     @abstractmethod
     def close(self):
         """Clean up profiler interface"""

metadata/profiler/interface/sqlalchemy/bigquery/profiler_interface.py

@@ -43,25 +43,6 @@ class BigQueryProfilerInterface(SQAProfilerInterface):
                 columns_list.extend(col)
         return columns_list
 
-    def _get_sampler(self, **kwargs):
-        """get sampler object"""
-        from metadata.profiler.processor.sampler.sampler_factory import (  # pylint: disable=import-outside-toplevel
-            sampler_factory_,
-        )
-
-        session = kwargs.get("session")
-        table = kwargs["table"]
-
-        return sampler_factory_.create(
-            self.service_connection_config.__class__.__name__,
-            client=session or self.session,
-            table=table,
-            profile_sample_config=self.profile_sample_config,
-            partition_details=self.partition_details,
-            profile_sample_query=self.profile_query,
-            table_type=self.table_entity.tableType,
-        )
-
     def get_columns(self) -> Column:
         """Get columns from table"""
         # pylint: disable=import-outside-toplevel