openmetadata-ingestion 1.8.1.0__py3-none-any.whl → 1.8.2.0__py3-none-any.whl

metadata/sampler/sqlalchemy/databricks/sampler.py

@@ -11,9 +11,12 @@
 """
 Helper module to handle data sampling for the profiler
 """
+from sqlalchemy import Column, text
+
 from metadata.ingestion.source.database.databricks.connection import (
     get_connection as databricks_get_connection,
 )
+from metadata.profiler.orm.types.custom_array import CustomArray
 from metadata.sampler.sqlalchemy.sampler import SQASampler
 
 
@@ -28,3 +31,27 @@ class DatabricksSamplerInterface(SQASampler):
         client = super().get_client()
         self.set_catalog(client)
         return client
+
+    def _handle_array_column(self, column: Column) -> bool:
+        """Check if a column is an array type"""
+        return isinstance(column.type, CustomArray)
+
+    def _get_slice_expression(self, column: Column):
+        """Generate SQL expression to slice array elements at query level
+
+        Args:
+            column_name: Name of the column
+            max_elements: Maximum number of elements to extract
+
+        Returns:
+            SQL expression string for array slicing
+        """
+        max_elements = self._get_max_array_elements()
+        return text(
+            f"""
+        CASE 
+            WHEN `{column.name}` IS NULL THEN NULL
+            ELSE slice(`{column.name}`, 1, {max_elements})
+        END AS `{column._label}`
+        """
+        )

metadata/sampler/sqlalchemy/sampler.py

@@ -41,6 +41,9 @@ logger = profiler_interface_registry_logger()
 
 RANDOM_LABEL = "random"
 
+# Default maximum number of elements to extract from array columns to prevent OOM
+DEFAULT_MAX_ARRAY_ELEMENTS = 10
+
 
 def _object_value_for_elem(self, elem):
     """
@@ -103,6 +106,36 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         """
         return selectable
 
+    def _get_max_array_elements(self) -> int:
+        """Get the maximum number of array elements from config or use default"""
+        if (
+            self.sample_config
+            and hasattr(self.sample_config, "maxArrayElements")
+            and self.sample_config.maxArrayElements
+        ):
+            return self.sample_config.maxArrayElements
+        return DEFAULT_MAX_ARRAY_ELEMENTS
+
+    def _handle_array_column(self, column: Column) -> bool:
+        """Check if a column is an array type"""
+        # Implement this method in the child classes
+        return False
+
+    def _process_array_value(self, value):
+        """Process array values to convert numpy arrays to Python lists"""
+        import numpy as np  # pylint: disable=import-outside-toplevel
+
+        if isinstance(value, np.ndarray):
+            return value.tolist()
+        return value
+
+    def _get_slice_expression(self, column: Column):
+        """Generate SQL expression to slice array elements at query level
+        By default, we return the column as is.
+        Child classes can override this method to return a different expression.
+        """
+        return column
+
     def _base_sample_query(self, column: Optional[Column], label=None):
         """Base query for sampling
 
@@ -215,16 +248,49 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
             ]
 
         with self.get_client_context() as client:
+
+            # Handle array columns with special query modification
+            max_elements = self._get_max_array_elements()
+            select_columns = []
+            has_array_columns = False
+
+            for col in sqa_columns:
+                if self._handle_array_column(col):
+                    slice_expression = self._get_slice_expression(col)
+                    select_columns.append(slice_expression)
+                    logger.debug(
+                        f"Limiting array column {col.name} to {max_elements} elements to prevent OOM"
+                    )
+                    has_array_columns = True
+                else:
+                    select_columns.append(col)
+
+            # Create query with modified columns
             sqa_sample = (
-                client.query(*sqa_columns)
+                client.query(*select_columns)
                 .select_from(ds)
                 .limit(self.sample_limit)
                 .all()
             )
 
+        # Process array columns manually if we used text() expressions
+        processed_rows = []
+        if has_array_columns:
+            for row in sqa_sample:
+                processed_row = []
+                for i, col in enumerate(sqa_columns):
+                    value = row[i]
+                    if self._handle_array_column(col):
+                        processed_value = self._process_array_value(value)
+                        processed_row.append(processed_value)
+                    else:
+                        processed_row.append(value)
+                processed_rows.append(processed_row)
+        else:
+            processed_rows = [list(row) for row in sqa_sample]
         return TableData(
             columns=[column.name for column in sqa_columns],
-            rows=[list(row) for row in sqa_sample],
+            rows=processed_rows,
         )
 
     def _fetch_sample_data_from_user_query(self) -> TableData:

metadata/sampler/sqlalchemy/unitycatalog/sampler.py

@@ -13,9 +13,13 @@
 Interfaces with database for all database engine
 supporting sqlalchemy abstraction layer
 """
+
+from sqlalchemy import Column, text
+
 from metadata.ingestion.source.database.databricks.connection import (
     get_connection as databricks_get_connection,
 )
+from metadata.profiler.orm.types.custom_array import CustomArray
 from metadata.sampler.sqlalchemy.sampler import SQASampler
 
 
@@ -30,3 +34,27 @@ class UnityCatalogSamplerInterface(SQASampler):
         client = super().get_client()
         self.set_catalog(client)
         return client
+
+    def _handle_array_column(self, column: Column) -> bool:
+        """Check if a column is an array type"""
+        return isinstance(column.type, CustomArray)
+
+    def _get_slice_expression(self, column: Column):
+        """Generate SQL expression to slice array elements at query level
+
+        Args:
+            column_name: Name of the column
+            max_elements: Maximum number of elements to extract
+
+        Returns:
+            SQL expression string for array slicing
+        """
+        max_elements = self._get_max_array_elements()
+        return text(
+            f"""
+        CASE 
+            WHEN `{column.name}` IS NULL THEN NULL
+            ELSE slice(`{column.name}`, 1, {max_elements})
+        END AS `{column._label}`
+        """
+        )

metadata/utils/filters.py

@@ -164,6 +164,21 @@ def filter_by_dashboard(
     return _filter(dashboard_filter_pattern, dashboard_name)
 
 
+def filter_by_stored_procedure(
+    stored_procedure_filter_pattern: Optional[FilterPattern], stored_procedure_name: str
+) -> bool:
+    """
+    Return True if the stored procedure needs to be filtered, False otherwise
+
+    Include takes precedence over exclude
+
+    :param stored_procedure_filter_pattern: Model defining stored procedure filtering logic
+    :param stored_procedure_name: stored procedure name
+    :return: True for filtering, False otherwise
+    """
+    return _filter(stored_procedure_filter_pattern, stored_procedure_name)
+
+
 def filter_by_fqn(fqn_filter_pattern: Optional[FilterPattern], fqn: str) -> bool:
     """
     Return True if the FQN needs to be filtered, False otherwise