openmetadata-ingestion 1.5.13.2__py3-none-any.whl → 1.6.0.0__py3-none-any.whl

metadata/profiler/api/models.py

@@ -15,85 +15,25 @@ Return types for Profiler workflow execution.
 We need to define this class as we end up having
 multiple profilers per table and columns.
 """
+
 from typing import List, Optional, Type, Union
 
-from pydantic import ConfigDict, Field
+from pydantic import ConfigDict
 from sqlalchemy import Column
 from sqlalchemy.orm import DeclarativeMeta
-from typing_extensions import Annotated
 
 from metadata.config.common import ConfigModel
 from metadata.generated.schema.api.data.createTableProfile import (
     CreateTableProfileRequest,
 )
-from metadata.generated.schema.entity.data.table import (
-    ColumnProfilerConfig,
-    PartitionProfilerConfig,
-    ProfileSampleType,
-    Table,
-    TableData,
-)
-from metadata.generated.schema.entity.services.connections.connectionBasicType import (
-    SampleDataStorageConfig,
-)
+from metadata.generated.schema.entity.data.table import Table
 from metadata.generated.schema.tests.customMetric import CustomMetric
-from metadata.generated.schema.type.basic import FullyQualifiedEntityName
-from metadata.ingestion.models.custom_pydantic import BaseModel
-from metadata.ingestion.models.table_metadata import ColumnTag
 from metadata.profiler.metrics.core import Metric, MetricTypes
 from metadata.profiler.processor.models import ProfilerDef
+from metadata.sampler.models import DatabaseAndSchemaConfig, TableConfig
 from metadata.utils.sqa_like_column import SQALikeColumn
 
 
-class ColumnConfig(ConfigModel):
-    """Column config for profiler"""
-
-    excludeColumns: Optional[List[str]] = None
-    includeColumns: Optional[List[ColumnProfilerConfig]] = None
-
-
-class BaseProfileConfig(ConfigModel):
-    """base profile config"""
-
-    fullyQualifiedName: FullyQualifiedEntityName
-    profileSample: Optional[Union[float, int]] = None
-    profileSampleType: Optional[ProfileSampleType] = None
-    sampleDataCount: Optional[int] = 100
-
-
-class TableConfig(BaseProfileConfig):
-    """table profile config"""
-
-    profileQuery: Optional[str] = None
-    partitionConfig: Optional[PartitionProfilerConfig] = None
-    columnConfig: Optional[ColumnConfig] = None
-
-    @classmethod
-    def from_database_and_schema_config(
-        cls, config: "DatabaseAndSchemaConfig", table_fqn: str
-    ):
-        table_config = TableConfig(
-            fullyQualifiedName=table_fqn,
-            profileSample=config.profileSample,
-            profileSampleType=config.profileSampleType,
-            sampleDataCount=config.sampleDataCount,
-        )
-        return table_config
-
-
-class DatabaseAndSchemaConfig(BaseProfileConfig):
-    """schema profile config"""
-
-    sampleDataStorageConfig: Optional[SampleDataStorageConfig] = None
-
-
-class ProfileSampleConfig(ConfigModel):
-    """Profile Sample Config"""
-
-    profile_sample: Optional[Union[float, int]] = None
-    profile_sample_type: Optional[ProfileSampleType] = ProfileSampleType.PERCENTAGE
-
-
 class ProfilerProcessorConfig(ConfigModel):
     """
     Defines how we read the processor information
@@ -106,15 +46,6 @@ class ProfilerProcessorConfig(ConfigModel):
     databaseConfig: Optional[List[DatabaseAndSchemaConfig]] = []
 
 
-class SampleData(BaseModel):
-    """TableData wrapper to handle ephemeral SampleData"""
-
-    data: Annotated[TableData, Field(None, description="Table Sample Data")]
-    store: Annotated[
-        bool, Field(False, description="Is the sample data should be stored or not")
-    ]
-
-
 class ProfilerResponse(ConfigModel):
     """
     ORM Profiler processor response.
@@ -125,8 +56,6 @@ class ProfilerResponse(ConfigModel):
 
     table: Table
     profile: CreateTableProfileRequest
-    sample_data: Optional[SampleData] = None
-    column_tags: Optional[List[ColumnTag]] = None
 
     def __str__(self):
         """Return the table name being processed"""

metadata/profiler/config.py

@@ -0,0 +1,39 @@
+#  Copyright 2021 Collate
+#  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.
+"""
+Profiler configuration helpers
+"""
+from typing import Optional
+
+from metadata.generated.schema.entity.data.database import (
+    Database,
+    DatabaseProfilerConfig,
+)
+from metadata.generated.schema.entity.data.databaseSchema import (
+    DatabaseSchema,
+    DatabaseSchemaProfilerConfig,
+)
+
+
+def get_database_profiler_config(
+    database_entity: Optional[Database],
+) -> Optional[DatabaseProfilerConfig]:
+    if database_entity and database_entity.databaseProfilerConfig:
+        return database_entity.databaseProfilerConfig
+    return None
+
+
+def get_schema_profiler_config(
+    schema_entity: Optional[DatabaseSchema],
+) -> Optional[DatabaseSchemaProfilerConfig]:
+    if schema_entity and schema_entity.databaseSchemaProfilerConfig:
+        return schema_entity.databaseSchemaProfilerConfig
+    return None

metadata/profiler/interface/nosql/profiler_interface.py

@@ -30,7 +30,6 @@ from metadata.profiler.interface.profiler_interface import ProfilerInterface
 from metadata.profiler.metrics.core import Metric, MetricTypes
 from metadata.profiler.metrics.registry import Metrics
 from metadata.profiler.processor.metric_filter import MetricFilter
-from metadata.profiler.processor.sampler.nosql.sampler import NoSQLSampler
 from metadata.utils.logger import profiler_interface_registry_logger
 from metadata.utils.sqa_like_column import SQALikeColumn
 
@@ -45,10 +44,6 @@ class NoSQLProfilerInterface(ProfilerInterface):
 
     # pylint: disable=too-many-arguments
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.sampler = self._get_sampler()
-
     def _compute_table_metrics(
         self,
         metrics: List[Type[Metric]],
@@ -157,32 +152,12 @@ class NoSQLProfilerInterface(ProfilerInterface):
     def fetch_sample_data(self, table, columns: List[SQALikeColumn]) -> TableData:
         return self.sampler.fetch_sample_data(columns)
 
-    def _get_sampler(self) -> NoSQLSampler:
-        """Get NoSQL sampler from config"""
-        from metadata.profiler.processor.sampler.sampler_factory import (  # pylint: disable=import-outside-toplevel
-            sampler_factory_,
-        )
-
-        return sampler_factory_.create(
-            self.service_connection_config.__class__.__name__,
-            table=self.table,
-            client=factory.create(
-                self.service_connection_config.__class__.__name__,
-                client=self.connection,
-            ),
-            profile_sample_config=self.profile_sample_config,
-            partition_details=self.partition_details,
-            profile_sample_query=self.profile_query,
-        )
-
     def get_composed_metrics(
         self, column: Column, metric: Metrics, column_results: Dict
     ):
         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):
         return None
 
     def get_all_metrics(

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
                         ),
                     )
                 )