openmetadata-ingestion 1.6.0.0rc1__py3-none-any.whl → 1.6.0.0rc3__py3-none-any.whl

metadata/sampler/pandas/sampler.py

@@ -0,0 +1,239 @@
+#  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.
+"""
+Helper module to handle data sampling
+for the profiler
+"""
+import math
+import random
+from typing import List, Optional, cast
+
+from metadata.data_quality.validations.table.pandas.tableRowInsertedCountToBeBetween import (
+    TableRowInsertedCountToBeBetweenValidator,
+)
+from metadata.generated.schema.entity.data.table import (
+    PartitionIntervalTypes,
+    PartitionProfilerConfig,
+    ProfileSampleType,
+    TableData,
+)
+from metadata.mixins.pandas.pandas_mixin import PandasInterfaceMixin
+from metadata.sampler.sampler_interface import SamplerInterface
+from metadata.utils.constants import COMPLEX_COLUMN_SEPARATOR
+from metadata.utils.datalake.datalake_utils import GenericDataFrameColumnParser
+from metadata.utils.logger import profiler_logger
+from metadata.utils.sqa_like_column import SQALikeColumn
+
+logger = profiler_logger()
+
+
+class DatalakeSampler(SamplerInterface, PandasInterfaceMixin):
+    """
+    Generates a sample of the data to not
+    run the query in the whole table.
+    """
+
+    def __init__(self, *args, **kwargs):
+        """Init the pandas sampler"""
+        super().__init__(*args, **kwargs)
+        self.partition_details = cast(PartitionProfilerConfig, self.partition_details)
+        self._table = None
+
+    @property
+    def raw_dataset(self):
+        """Get the raw dataset (unpartitioned and unsampled)"""
+        if not self._table:
+            self._table = self.get_dataframes(
+                service_connection_config=self.service_connection_config,
+                client=self.client.client,
+                table=self.entity,
+            )
+        return self._table
+
+    def get_client(self):
+        return self.connection.client
+
+    def _partitioned_table(self):
+        """Get partitioned table"""
+        partition_field = self.partition_details.partitionColumnName
+        if (
+            self.partition_details.partitionIntervalType
+            == PartitionIntervalTypes.COLUMN_VALUE
+        ):
+            return [
+                df[df[partition_field].isin(self.partition_details.partitionValues)]
+                for df in self.raw_dataset
+            ]
+        if (
+            self.partition_details.partitionIntervalType
+            == PartitionIntervalTypes.INTEGER_RANGE
+        ):
+            return [
+                df[
+                    df[partition_field].between(
+                        self.partition_details.partitionIntegerRangeStart,
+                        self.partition_details.partitionIntegerRangeEnd,
+                    )
+                ]
+                for df in self.raw_dataset
+            ]
+        if (
+            self.partition_details.partitionIntervalType
+            in {PartitionIntervalTypes.INGESTION_TIME, PartitionIntervalTypes.TIME_UNIT}
+            and self.partition_details.partitionIntervalUnit
+            and self.partition_details.partitionInterval
+        ):
+            return [
+                df[
+                    df[partition_field]
+                    >= TableRowInsertedCountToBeBetweenValidator._get_threshold_date(  # pylint: disable=protected-access
+                        self.partition_details.partitionIntervalUnit.value,
+                        self.partition_details.partitionInterval,
+                    )
+                ]
+                for df in self.raw_dataset
+            ]
+        logger.warning(
+            f"{self.partition_details.partitionIntervalType} is not a supported partition type. "
+            "Partitioning will be ignored. Use one of "
+            f"{', '.join(list(PartitionIntervalTypes.__members__))}"
+        )
+        return self.raw_dataset
+
+    def _fetch_sample_data_from_user_query(self) -> TableData:
+        """Fetch sample data from user query"""
+        cols, rows = self.get_col_row(data_frame=self._rdn_sample_from_user_query())
+        return TableData(columns=cols, rows=rows)
+
+    def _rdn_sample_from_user_query(self):
+        """Generate sample from user query"""
+        return [df.query(self.sample_query) for df in self.raw_dataset]
+
+    def _get_sampled_dataframe(self):
+        """
+        returns sampled ometa dataframes
+        """
+        random.shuffle(self.raw_dataset)  # we'll shuffle the list of dataframes
+        # sampling data based on profiler config (if any)
+        if self.sample_config.profile_sample_type == ProfileSampleType.PERCENTAGE:
+            try:
+                profile_sample = (self.sample_config.profile_sample or 100) / 100
+            except TypeError:
+                # if the profile sample is not a number or is None
+                # we'll set it to 100
+                profile_sample = self.sample_config.profile_sample = 100
+            return [
+                df.sample(
+                    frac=profile_sample,
+                    random_state=random.randint(0, 100),
+                    replace=True,
+                )
+                for df in self.raw_dataset
+            ]
+
+        # we'll distribute the sample size equally among the dataframes
+        sample_rows_per_chunk: int = math.floor(
+            (self.sample_config.profile_sample or 100) / len(self.raw_dataset)
+        )
+        num_rows = sum(len(df) for df in self.raw_dataset)
+
+        # if we have less rows than the sample size
+        # we'll return the whole table
+        if sample_rows_per_chunk > num_rows:
+            return self.raw_dataset
+        return [
+            df.sample(
+                n=sample_rows_per_chunk,
+                random_state=random.randint(0, 100),
+                replace=True,
+            )
+            for df in self.raw_dataset
+        ]
+
+    def get_col_row(self, data_frame, columns: Optional[List[SQALikeColumn]] = None):
+        """
+        Fetches columns and rows from the data_frame
+        """
+        if columns:
+            cols = [col.name for col in columns]
+        else:
+            # we'll use the first dataframe to get the columns
+            cols = data_frame[0].columns.tolist()
+        rows = []
+        # Sample Data should not exceed sample limit
+        for chunk in data_frame:
+            rows.extend(self._fetch_rows(chunk[cols])[: self.sample_limit])
+            if len(rows) >= (self.sample_limit or 100):
+                break
+        return cols, rows
+
+    def get_dataset(self, **__):
+        """Generate random sample from the table
+
+        Returns:
+            List[DataFrame]
+        """
+        if self.sample_query:
+            return self._rdn_sample_from_user_query()
+
+        if self.partition_details:
+            self._table = self._partitioned_table()
+
+        if not self.sample_config.profile_sample or (
+            self.sample_config.profile_sample == 100
+            and self.sample_config.profile_sample_type == ProfileSampleType.PERCENTAGE
+        ):
+            return self.raw_dataset
+        return self._get_sampled_dataframe()
+
+    def _fetch_rows(self, data_frame):
+        return data_frame.dropna().values.tolist()
+
+    def fetch_sample_data(
+        self, columns: Optional[List[SQALikeColumn]] = None
+    ) -> TableData:
+        """Fetch sample data from the table
+
+        Returns:
+            TableData:
+        """
+        if self.sample_query:
+            return self._fetch_sample_data_from_user_query()
+
+        cols, rows = self.get_col_row(data_frame=self.raw_dataset, columns=columns)
+        return TableData(columns=cols, rows=rows)
+
+    def get_columns(self) -> List[Optional[SQALikeColumn]]:
+        """Get SQALikeColumns for datalake to be passed for metric computation"""
+        sqalike_columns = []
+        if self.raw_dataset:
+            for column_name in self.raw_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.raw_dataset:
+                            df.rename(
+                                columns={column_name: complex_col_name}, inplace=True
+                            )
+                column_name = complex_col_name or column_name
+                sqalike_columns.append(
+                    SQALikeColumn(
+                        column_name,
+                        GenericDataFrameColumnParser.fetch_col_types(
+                            self.raw_dataset[0], column_name
+                        ),
+                    )
+                )
+            return sqalike_columns
+        return []

metadata/{utils → sampler}/partition.py

@@ -25,6 +25,7 @@ from metadata.generated.schema.entity.data.table import (
 from metadata.generated.schema.entity.services.databaseService import (
     DatabaseServiceType,
 )
+from metadata.sampler.models import TableConfig
 
 
 def validate_athena_injected_partitioning(
@@ -66,7 +67,9 @@ def validate_athena_injected_partitioning(
                 )
 
 
-def get_partition_details(entity: Table) -> Optional[PartitionProfilerConfig]:
+def get_partition_details(
+    entity: Table, entity_config: Optional[TableConfig] = None
+) -> Optional[PartitionProfilerConfig]:
     """Build PartitionProfilerConfig object from entity
 
     Args:
@@ -74,6 +77,9 @@ def get_partition_details(entity: Table) -> Optional[PartitionProfilerConfig]:
     Returns:
         PartitionProfilerConfig
     """
+    if entity_config:
+        return entity_config.partitionConfig
+
     # Gather service type information
     service_type = getattr(entity, "serviceType", None)
 
@@ -100,56 +106,65 @@ def get_partition_details(entity: Table) -> Optional[PartitionProfilerConfig]:
         return profiler_partitioning_config
 
     if service_type == DatabaseServiceType.BigQuery:
-        if table_partition:
-            column_partitions: Optional[
-                List[PartitionColumnDetails]
-            ] = entity.tablePartition.columns
-            if not column_partitions:
-                raise TypeError("table partition missing. Skipping table")
-
-            partiton = column_partitions[0]
-
-            if partiton.intervalType == PartitionIntervalTypes.TIME_UNIT:
-                return PartitionProfilerConfig(
-                    enablePartitioning=True,
-                    partitionColumnName=partiton.columnName,
-                    partitionIntervalUnit=PartitionIntervalUnit.DAY
-                    if partiton.interval != "HOUR"
-                    else partiton.interval,
-                    partitionInterval=1,
-                    partitionIntervalType=partiton.intervalType.value,
-                    partitionValues=None,
-                    partitionIntegerRangeStart=None,
-                    partitionIntegerRangeEnd=None,
-                )
-            if partiton.intervalType == PartitionIntervalTypes.INGESTION_TIME:
-                return PartitionProfilerConfig(
-                    enablePartitioning=True,
-                    partitionColumnName="_PARTITIONDATE"
-                    if partiton.interval == "DAY"
-                    else "_PARTITIONTIME",
-                    partitionIntervalUnit=PartitionIntervalUnit.DAY
-                    if partiton.interval != "HOUR"
-                    else partiton.interval,
-                    partitionInterval=1,
-                    partitionIntervalType=partiton.intervalType.value,
-                    partitionValues=None,
-                    partitionIntegerRangeStart=None,
-                    partitionIntegerRangeEnd=None,
-                )
-            if partiton.intervalType == PartitionIntervalTypes.INTEGER_RANGE:
-                return PartitionProfilerConfig(
-                    enablePartitioning=True,
-                    partitionColumnName=partiton.columnName,
-                    partitionIntervalUnit=None,
-                    partitionInterval=None,
-                    partitionIntervalType=partiton.intervalType.value,
-                    partitionValues=None,
-                    partitionIntegerRangeStart=1,
-                    partitionIntegerRangeEnd=10000,
-                )
-            raise TypeError(
-                f"Unsupported partition type {partiton.intervalType}. Skipping table"
+        return _handle_bigquery_partition(entity, table_partition)
+
+    return None
+
+
+def _handle_bigquery_partition(
+    entity: Table, table_partition: TablePartition
+) -> Optional[PartitionProfilerConfig]:
+    """Bigquery specific logic for partitions"""
+    if table_partition:
+        column_partitions: Optional[
+            List[PartitionColumnDetails]
+        ] = entity.tablePartition.columns
+        if not column_partitions:
+            raise TypeError("table partition missing. Skipping table")
+
+        partition = column_partitions[0]
+
+        if partition.intervalType == PartitionIntervalTypes.TIME_UNIT:
+            return PartitionProfilerConfig(
+                enablePartitioning=True,
+                partitionColumnName=partition.columnName,
+                partitionIntervalUnit=PartitionIntervalUnit.DAY
+                if partition.interval != "HOUR"
+                else partition.interval,
+                partitionInterval=1,
+                partitionIntervalType=partition.intervalType.value,
+                partitionValues=None,
+                partitionIntegerRangeStart=None,
+                partitionIntegerRangeEnd=None,
+            )
+        if partition.intervalType == PartitionIntervalTypes.INGESTION_TIME:
+            return PartitionProfilerConfig(
+                enablePartitioning=True,
+                partitionColumnName="_PARTITIONDATE"
+                if partition.interval == "DAY"
+                else "_PARTITIONTIME",
+                partitionIntervalUnit=PartitionIntervalUnit.DAY
+                if partition.interval != "HOUR"
+                else partition.interval,
+                partitionInterval=1,
+                partitionIntervalType=partition.intervalType.value,
+                partitionValues=None,
+                partitionIntegerRangeStart=None,
+                partitionIntegerRangeEnd=None,
             )
+        if partition.intervalType == PartitionIntervalTypes.INTEGER_RANGE:
+            return PartitionProfilerConfig(
+                enablePartitioning=True,
+                partitionColumnName=partition.columnName,
+                partitionIntervalUnit=None,
+                partitionInterval=None,
+                partitionIntervalType=partition.intervalType.value,
+                partitionValues=None,
+                partitionIntegerRangeStart=1,
+                partitionIntegerRangeEnd=10000,
+            )
+        raise TypeError(
+            f"Unsupported partition type {partition.intervalType}. Skipping table"
+        )
 
     return None

metadata/sampler/processor.py

@@ -0,0 +1,154 @@
+#  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.
+"""
+Data Sampler for the PII Workflow
+"""
+import traceback
+from copy import deepcopy
+from typing import Optional, cast
+
+from metadata.generated.schema.entity.data.database import Database
+from metadata.generated.schema.entity.data.table import Table
+from metadata.generated.schema.entity.services.databaseService import DatabaseConnection
+from metadata.generated.schema.entity.services.ingestionPipelines.status import (
+    StackTraceError,
+)
+from metadata.generated.schema.entity.services.serviceType import ServiceType
+from metadata.generated.schema.metadataIngestion.databaseServiceAutoClassificationPipeline import (
+    DatabaseServiceAutoClassificationPipeline,
+)
+from metadata.generated.schema.metadataIngestion.workflow import (
+    OpenMetadataWorkflowConfig,
+)
+from metadata.ingestion.api.models import Either
+from metadata.ingestion.api.parser import parse_workflow_config_gracefully
+from metadata.ingestion.api.step import Step
+from metadata.ingestion.api.steps import Processor
+from metadata.ingestion.ometa.ometa_api import OpenMetadata
+from metadata.profiler.api.models import ProfilerProcessorConfig
+from metadata.profiler.source.metadata import ProfilerSourceAndEntity
+from metadata.sampler.config import get_config_for_table
+from metadata.sampler.models import SampleConfig, SampleData, SamplerResponse
+from metadata.sampler.sampler_interface import SamplerInterface
+from metadata.utils.profiler_utils import get_context_entities
+from metadata.utils.service_spec.service_spec import import_sampler_class
+
+
+class SamplerProcessor(Processor):
+    """Use the profiler interface to fetch the sample data"""
+
+    def __init__(self, config: OpenMetadataWorkflowConfig, metadata: OpenMetadata):
+        super().__init__()
+
+        self.config = config
+        self.metadata = metadata
+
+        self.source_config: DatabaseServiceAutoClassificationPipeline = cast(
+            DatabaseServiceAutoClassificationPipeline,
+            self.config.source.sourceConfig.config,
+        )  # Used to satisfy type checked
+        # We still rely on the orm-processor. We should decouple this in the future
+        self.profiler_config = ProfilerProcessorConfig.model_validate(
+            self.config.processor.model_dump().get("config")
+        )
+
+        self._interface_type: str = config.source.type.lower()
+        self.sampler_class = import_sampler_class(
+            ServiceType.Database, source_type=self._interface_type
+        )
+
+    @property
+    def name(self) -> str:
+        return "Sampler"
+
+    def _run(self, record: ProfilerSourceAndEntity) -> Either[SamplerResponse]:
+        """Fetch the sample data and pass it down the pipeline"""
+
+        try:
+            entity = cast(Table, record.entity)
+            schema_entity, database_entity, _ = get_context_entities(
+                entity=entity, metadata=self.metadata
+            )
+            service_conn_config = self._copy_service_config(
+                self.config, database_entity
+            )
+
+            sampler_interface: SamplerInterface = self.sampler_class.create(
+                service_connection_config=service_conn_config,
+                ometa_client=self.metadata,
+                entity=entity,
+                schema_entity=schema_entity,
+                database_entity=database_entity,
+                table_config=get_config_for_table(entity, self.profiler_config),
+                default_sample_config=SampleConfig(),
+                default_sample_data_count=self.source_config.sampleDataCount,
+            )
+            sample_data = SampleData(
+                data=sampler_interface.generate_sample_data(),
+                store=self.source_config.storeSampleData,
+            )
+
+            return Either(
+                right=SamplerResponse(
+                    table=entity,
+                    sample_data=sample_data,
+                )
+            )
+
+        except Exception as exc:
+            return Either(
+                left=StackTraceError(
+                    name=record.entity.fullyQualifiedName.root,
+                    error=f"Unexpected exception processing entity {record.entity.fullyQualifiedName.root}: {exc}",
+                    stackTrace=traceback.format_exc(),
+                )
+            )
+
+    @classmethod
+    def create(
+        cls,
+        config_dict: dict,
+        metadata: OpenMetadata,
+        pipeline_name: Optional[str] = None,
+    ) -> "Step":
+        config = parse_workflow_config_gracefully(config_dict)
+        return cls(config=config, metadata=metadata)
+
+    def _copy_service_config(
+        self, config: OpenMetadataWorkflowConfig, database: Database
+    ) -> DatabaseConnection:
+        """Make a copy of the service config and update the database name
+
+        Args:
+            database (_type_): a database entity
+
+        Returns:
+            DatabaseService.__config__
+        """
+        config_copy = deepcopy(
+            config.source.serviceConnection.root.config  # type: ignore
+        )
+        if hasattr(
+            config_copy,  # type: ignore
+            "supportsDatabase",
+        ):
+            if hasattr(config_copy, "database"):
+                config_copy.database = database.name.root  # type: ignore
+            if hasattr(config_copy, "catalog"):
+                config_copy.catalog = database.name.root  # type: ignore
+
+        # we know we'll only be working with DatabaseConnection, we cast the type to satisfy type checker
+        config_copy = cast(DatabaseConnection, config_copy)
+
+        return config_copy
+
+    def close(self) -> None:
+        """Nothing to close"""