openmetadata-ingestion 1.7.5.3__py3-none-any.whl → 1.7.6.0__py3-none-any.whl

metadata/ingestion/models/patch_request.py

@@ -342,6 +342,7 @@ def build_patch(
     array_entity_fields: Optional[List] = None,
     remove_change_description: bool = True,
     override_metadata: Optional[bool] = False,
+    skip_on_failure: Optional[bool] = True,
 ) -> Optional[jsonpatch.JsonPatch]:
     """
     Given an Entity type and Source entity and Destination entity,
@@ -352,6 +353,10 @@ def build_patch(
         destination: payload with changes applied to the source.
         allowed_fields: List of field names to filter from source and destination models
         restrict_update_fields: List of field names which will only support add operation
+        array_entity_fields: List of array fields to sort for consistent patching
+        remove_change_description: Whether to remove change description from entities
+        override_metadata: Whether to override existing metadata fields
+        skip_on_failure: Whether to skip the patch operation on failure (default: True)
 
     Returns
         Updated Entity
@@ -421,10 +426,37 @@ def build_patch(
             patch.patch = updated_operations
 
         return patch
-    except Exception:
+    except Exception as exc:
         logger.debug(traceback.format_exc())
-        logger.warning("Couldn't build patch for Entity.")
-        return None
+        if skip_on_failure:
+            entity_info = ""
+            try:
+                if hasattr(source, "fullyQualifiedName"):
+                    entity_info = f" for '{source.fullyQualifiedName.root}'"
+                elif hasattr(source, "name"):
+                    entity_info = f" for '{source.name.root}'"
+            except Exception:
+                pass
+
+            logger.warning(
+                f"Failed to build patch{entity_info}. The patch generation was skipped. "
+                f"Reason: {exc}"
+            )
+            return None
+        else:
+            entity_info = ""
+            try:
+                if hasattr(source, "fullyQualifiedName"):
+                    entity_info = f" for '{source.fullyQualifiedName.root}'"
+                elif hasattr(source, "name"):
+                    entity_info = f" for '{source.name.root}'"
+            except Exception:
+                pass
+
+            raise RuntimeError(
+                f"Failed to build patch{entity_info}. The patch generation failed. "
+                f"Set 'skip_on_failure=True' to skip failed patch operations. Error: {exc}"
+            ) from exc
 
 
 def _get_attribute_name(attr: T) -> str:
@@ -534,7 +566,7 @@ def _remove_change_description(entity: T) -> T:
     We never want to patch that, and we won't have that information
     from the source. It's fully handled in the server.
     """
-    if getattr(entity, "changeDescription"):
+    if hasattr(entity, "changeDescription") and getattr(entity, "changeDescription"):
         entity.changeDescription = None
 
     return entity

metadata/ingestion/ometa/mixins/patch_mixin.py

@@ -124,6 +124,7 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
         restrict_update_fields: Optional[List] = None,
         array_entity_fields: Optional[List] = None,
         override_metadata: Optional[bool] = False,
+        skip_on_failure: Optional[bool] = True,
     ) -> Optional[T]:
         """
         Given an Entity type and Source entity and Destination entity,
@@ -135,6 +136,9 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
             destination: payload with changes applied to the source.
             allowed_fields: List of field names to filter from source and destination models
             restrict_update_fields: List of field names which will only support add operation
+            array_entity_fields: List of array fields to sort for consistent patching
+            override_metadata: Whether to override existing metadata fields
+            skip_on_failure: Whether to skip the patch operation on failure (default: True)
 
         Returns
             Updated Entity
@@ -147,6 +151,7 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
                 restrict_update_fields=restrict_update_fields,
                 array_entity_fields=array_entity_fields,
                 override_metadata=override_metadata,
+                skip_on_failure=skip_on_failure,
             )
 
             if not patch:
@@ -160,9 +165,19 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
 
         except Exception as exc:
             logger.debug(traceback.format_exc())
-            logger.warning(f"Error trying to PATCH {get_log_name(source)}: {exc}")
-
-        return None
+            if skip_on_failure:
+                entity_name = get_log_name(source)
+                logger.warning(
+                    f"Failed to update {entity_name}. The patch operation was skipped. "
+                    f"Reason: {exc}"
+                )
+                return None
+            else:
+                entity_name = get_log_name(source)
+                raise RuntimeError(
+                    f"Failed to update {entity_name}. The patch operation failed. "
+                    f"Set 'skip_on_failure=True' to skip failed patches. Error: {exc}"
+                ) from exc
 
     def patch_description(
         self,
@@ -170,6 +185,7 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
         source: T,
         description: str,
         force: bool = False,
+        skip_on_failure: bool = True,
     ) -> Optional[T]:
         """
         Given an Entity type and ID, JSON PATCH the description.
@@ -180,33 +196,51 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
             description: new description to add
             force: if True, we will patch any existing description. Otherwise, we will maintain
                 the existing data.
+            skip_on_failure: if True, return None on failure instead of raising exception
         Returns
             Updated Entity
         """
-        if isinstance(source, TestCase):
-            instance: Optional[T] = self._fetch_entity_if_exists(
-                entity=entity,
-                entity_id=source.id,
-                fields=["testDefinition", "testSuite"],
-            )
-        else:
-            instance: Optional[T] = self._fetch_entity_if_exists(
-                entity=entity, entity_id=source.id
-            )
+        try:
+            if isinstance(source, TestCase):
+                instance: Optional[T] = self._fetch_entity_if_exists(
+                    entity=entity,
+                    entity_id=source.id,
+                    fields=["testDefinition", "testSuite"],
+                )
+            else:
+                instance: Optional[T] = self._fetch_entity_if_exists(
+                    entity=entity, entity_id=source.id
+                )
 
-        if not instance:
-            return None
+            if not instance:
+                return None
 
-        if instance.description and not force:
-            # If the description is already present and force is not passed,
-            # description will not be overridden
-            return None
+            if instance.description and not force:
+                # If the description is already present and force is not passed,
+                # description will not be overridden
+                return None
 
-        # https://docs.pydantic.dev/latest/usage/exporting_models/#modelcopy
-        destination = source.model_copy(deep=True)
-        destination.description = Markdown(description)
+            # https://docs.pydantic.dev/latest/usage/exporting_models/#modelcopy
+            destination = source.model_copy(deep=True)
+            destination.description = Markdown(description)
 
-        return self.patch(entity=entity, source=source, destination=destination)
+            return self.patch(
+                entity=entity,
+                source=source,
+                destination=destination,
+                skip_on_failure=skip_on_failure,
+            )
+        except Exception as exc:
+            if skip_on_failure:
+                logger.debug(traceback.format_exc())
+                entity_name = get_log_name(source)
+                logger.warning(
+                    f"Failed to patch description for {entity_name}. The patch operation was skipped. "
+                    f"Reason: {exc}"
+                )
+                return None
+            else:
+                raise
 
     def patch_table_constraints(
         self,
@@ -275,6 +309,7 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
         operation: Union[
             PatchOperation.ADD, PatchOperation.REMOVE
         ] = PatchOperation.ADD,
+        skip_on_failure: bool = True,
     ) -> Optional[T]:
         """
         Given an Entity type and ID, JSON PATCH the tag.
@@ -284,29 +319,47 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
             source: Source entity object
             tag_label: TagLabel to add or remove
             operation: Patch Operation to add or remove the tag.
+            skip_on_failure: if True, return None on failure instead of raising exception
         Returns
             Updated Entity
         """
-        instance: Optional[T] = self._fetch_entity_if_exists(
-            entity=entity, entity_id=source.id, fields=["tags"]
-        )
-        if not instance:
-            return None
+        try:
+            instance: Optional[T] = self._fetch_entity_if_exists(
+                entity=entity, entity_id=source.id, fields=["tags"]
+            )
+            if not instance:
+                return None
 
-        # Initialize empty tag list or the last updated tags
-        source.tags = instance.tags or []
-        destination = source.model_copy(deep=True)
+            # Initialize empty tag list or the last updated tags
+            source.tags = instance.tags or []
+            destination = source.model_copy(deep=True)
 
-        tag_fqns = {label.tagFQN.root for label in tag_labels}
+            tag_fqns = {label.tagFQN.root for label in tag_labels}
 
-        if operation == PatchOperation.REMOVE:
-            for tag in destination.tags:
-                if tag.tagFQN.root in tag_fqns:
-                    destination.tags.remove(tag)
-        else:
-            destination.tags.extend(tag_labels)
+            if operation == PatchOperation.REMOVE:
+                for tag in destination.tags:
+                    if tag.tagFQN.root in tag_fqns:
+                        destination.tags.remove(tag)
+            else:
+                destination.tags.extend(tag_labels)
 
-        return self.patch(entity=entity, source=source, destination=destination)
+            return self.patch(
+                entity=entity,
+                source=source,
+                destination=destination,
+                skip_on_failure=skip_on_failure,
+            )
+        except Exception as exc:
+            if skip_on_failure:
+                logger.debug(traceback.format_exc())
+                entity_name = get_log_name(source)
+                logger.warning(
+                    f"Failed to patch tags for {entity_name}. The patch operation was skipped. "
+                    f"Reason: {exc}"
+                )
+                return None
+            else:
+                raise
 
     def patch_tag(
         self,
@@ -316,11 +369,16 @@ class OMetaPatchMixin(OMetaPatchMixinBase):
         operation: Union[
             PatchOperation.ADD, PatchOperation.REMOVE
         ] = PatchOperation.ADD,
+        skip_on_failure: bool = True,
     ) -> Optional[T]:
         """Will be deprecated in 1.3"""
         logger.warning("patch_tag will be deprecated in 1.3. Use `patch_tags` instead.")
         return self.patch_tags(
-            entity=entity, source=source, tag_labels=[tag_label], operation=operation
+            entity=entity,
+            source=source,
+            tag_labels=[tag_label],
+            operation=operation,
+            skip_on_failure=skip_on_failure,
         )
 
     def patch_owner(

metadata/readers/dataframe/base.py

@@ -22,6 +22,8 @@ from metadata.readers.file.config_source_factory import get_reader
 from metadata.readers.models import ConfigSource
 from metadata.utils.logger import ingestion_logger
 
+MAX_FILE_SIZE_FOR_PREVIEW = 50 * 1024 * 1024  # 50MB
+
 logger = ingestion_logger()
 
 

metadata/readers/dataframe/parquet.py

@@ -14,6 +14,8 @@ Generic Delimiter-Separated-Values implementation
 """
 from functools import singledispatchmethod
 
+from pyarrow.parquet import ParquetFile
+
 from metadata.generated.schema.entity.services.connections.database.datalake.azureConfig import (
     AzureConfig,
 )
@@ -26,11 +28,18 @@ from metadata.generated.schema.entity.services.connections.database.datalake.s3C
 from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
     LocalConfig,
 )
-from metadata.readers.dataframe.base import DataFrameReader, FileFormatException
+from metadata.readers.dataframe.base import (
+    MAX_FILE_SIZE_FOR_PREVIEW,
+    DataFrameReader,
+    FileFormatException,
+)
 from metadata.readers.dataframe.common import dataframe_to_chunks
 from metadata.readers.dataframe.models import DatalakeColumnWrapper
 from metadata.readers.file.adls import AZURE_PATH, return_azure_storage_options
 from metadata.readers.models import ConfigSource
+from metadata.utils.logger import ingestion_logger
+
+logger = ingestion_logger()
 
 
 class ParquetDataFrameReader(DataFrameReader):
@@ -39,6 +48,90 @@ class ParquetDataFrameReader(DataFrameReader):
     from any source based on its init client.
     """
 
+    def _read_parquet_in_batches(
+        self, parquet_file: ParquetFile, batch_size: int = 10000
+    ):
+        """
+        Read a large parquet file in batches to avoid memory issues.
+        Includes multiple fallback strategies for older PyArrow versions.
+
+        Args:
+            parquet_file: PyArrow ParquetFile or similar object
+            batch_size: Number of rows to read per batch
+
+        Returns:
+            List of DataFrame chunks
+        """
+        chunks = []
+        batch_count = 0
+
+        try:
+            # Method 1: iter_batches (PyArrow >= 3.0 - preferred)
+            if hasattr(parquet_file, "iter_batches"):
+                logger.info(
+                    "Reading large parquet file in batches to avoid memory issues"
+                )
+                for batch in parquet_file.iter_batches(batch_size=batch_size):
+                    df_batch = batch.to_pandas()
+                    if not df_batch.empty:
+                        chunks.extend(dataframe_to_chunks(df_batch))
+                        batch_count += 1
+
+                logger.info(
+                    f"Successfully processed {batch_count} batches from large parquet file"
+                )
+                return chunks
+
+            # Method 2: Row group reading (PyArrow >= 0.15.0)
+            elif hasattr(parquet_file, "num_row_groups") and hasattr(
+                parquet_file, "read_row_group"
+            ):
+                logger.warning(
+                    "iter_batches not available, using row group reading as fallback"
+                )
+
+                for i in range(parquet_file.num_row_groups):
+                    try:
+                        row_group_table = parquet_file.read_row_group(i)
+                        df_chunk = row_group_table.to_pandas()
+                        if not df_chunk.empty:
+                            # Further chunk if row group is still too large
+                            if len(df_chunk) > batch_size:
+                                chunks.extend(dataframe_to_chunks(df_chunk))
+                            else:
+                                chunks.append(df_chunk)
+                            batch_count += 1
+                    except Exception as row_exc:
+                        logger.warning(f"Failed to read row group {i}: {row_exc}")
+                        continue
+
+                if chunks:
+                    logger.info(
+                        f"Successfully processed {batch_count} row groups from large parquet file"
+                    )
+                    return chunks
+
+            # Method 3: Regular reading (final fallback)
+            logger.warning(
+                "No chunking methods available, falling back to regular reading"
+            )
+            df = parquet_file.read().to_pandas()
+            chunks.extend(dataframe_to_chunks(df))
+
+        except Exception as exc:
+            # If all chunking fails, try regular reading as final fallback
+            logger.warning(
+                f"Batched reading failed: {exc}. Falling back to regular reading - this may cause memory issues for large files"
+            )
+            try:
+                df = parquet_file.read().to_pandas()
+                chunks.extend(dataframe_to_chunks(df))
+            except Exception as fallback_exc:
+                logger.error(f"Failed to read parquet file: {fallback_exc}")
+                raise fallback_exc
+
+        return chunks
+
     @singledispatchmethod
     def _read_parquet_dispatch(
         self, config_source: ConfigSource, key: str, bucket_name: str
@@ -48,24 +141,46 @@ class ParquetDataFrameReader(DataFrameReader):
     @_read_parquet_dispatch.register
     def _(self, _: GCSConfig, key: str, bucket_name: str) -> DatalakeColumnWrapper:
         """
-        Read the CSV file from the gcs bucket and return a dataframe
+        Read the Parquet file from the gcs bucket and return a dataframe
         """
         # pylint: disable=import-outside-toplevel
         from gcsfs import GCSFileSystem
-        from pyarrow.parquet import ParquetFile
 
         gcs = GCSFileSystem()
-        file = gcs.open(f"gs://{bucket_name}/{key}")
-        dataframe_response = (
-            ParquetFile(file).read().to_pandas(split_blocks=True, self_destruct=True)
-        )
-        return dataframe_to_chunks(dataframe_response)
+        file_path = f"gs://{bucket_name}/{key}"
+
+        # Check file size to determine reading strategy
+        try:
+            file_info = gcs.info(file_path)
+            file_size = file_info.get("size", 0)
+
+            file = gcs.open(file_path)
+            parquet_file = ParquetFile(file)
+
+            if self._should_use_chunking(file_size):
+                # Use batched reading for large files
+                return self._read_parquet_in_batches(parquet_file)
+            else:
+                # Use regular reading for smaller files
+                dataframe_response = parquet_file.read().to_pandas(
+                    split_blocks=True, self_destruct=True
+                )
+                return dataframe_to_chunks(dataframe_response)
+
+        except Exception:
+            # Fallback to regular reading if size check fails
+            file = gcs.open(file_path)
+            parquet_file = ParquetFile(file)
+            dataframe_response = parquet_file.read().to_pandas(
+                split_blocks=True, self_destruct=True
+            )
+            return dataframe_to_chunks(dataframe_response)
 
     @_read_parquet_dispatch.register
     def _(self, _: S3Config, key: str, bucket_name: str) -> DatalakeColumnWrapper:
         # pylint: disable=import-outside-toplevel
         from pyarrow.fs import S3FileSystem
-        from pyarrow.parquet import ParquetDataset
+        from pyarrow.parquet import ParquetDataset, ParquetFile
 
         client_kwargs = {
             "endpoint_override": (
@@ -90,13 +205,40 @@ class ParquetDataFrameReader(DataFrameReader):
         s3_fs = S3FileSystem(**client_kwargs)
 
         bucket_uri = f"{bucket_name}/{key}"
-        dataset = ParquetDataset(bucket_uri, filesystem=s3_fs)
 
-        return dataframe_to_chunks(dataset.read_pandas().to_pandas())
+        # Check file size to determine reading strategy
+        try:
+            file_info = s3_fs.get_file_info(bucket_uri)
+            file_size = file_info.size if hasattr(file_info, "size") else 0
+
+            if self._should_use_chunking(file_size):
+                # Use ParquetFile for batched reading of large files
+                logger.info(
+                    f"Large parquet file detected ({file_size} bytes > {MAX_FILE_SIZE_FOR_PREVIEW} bytes). "
+                    f"Using batched reading for file: {bucket_uri}"
+                )
+                parquet_file = ParquetFile(bucket_uri, filesystem=s3_fs)
+                return self._read_parquet_in_batches(parquet_file)
+            else:
+                # Use ParquetDataset for regular reading of smaller files
+                logger.debug(
+                    f"Reading small parquet file ({file_size} bytes): {bucket_uri}"
+                )
+                dataset = ParquetDataset(bucket_uri, filesystem=s3_fs)
+                return dataframe_to_chunks(dataset.read_pandas().to_pandas())
+
+        except Exception as exc:
+            # Fallback to regular reading if size check fails
+            logger.warning(
+                f"Could not determine file size for {bucket_uri}: {exc}. Using regular reading"
+            )
+            dataset = ParquetDataset(bucket_uri, filesystem=s3_fs)
+            return dataframe_to_chunks(dataset.read_pandas().to_pandas())
 
     @_read_parquet_dispatch.register
     def _(self, _: AzureConfig, key: str, bucket_name: str) -> DatalakeColumnWrapper:
         import pandas as pd  # pylint: disable=import-outside-toplevel
+        import pyarrow.fs as fs
 
         storage_options = return_azure_storage_options(self.config_source)
         account_url = AZURE_PATH.format(
@@ -104,8 +246,33 @@ class ParquetDataFrameReader(DataFrameReader):
             account_name=self.config_source.securityConfig.accountName,
             key=key,
         )
-        dataframe = pd.read_parquet(account_url, storage_options=storage_options)
-        return dataframe_to_chunks(dataframe)
+
+        # Check file size to determine reading strategy
+        try:
+            # Try to get file size from Azure
+            azure_fs = fs.SubTreeFileSystem(
+                account_url, fs.AzureFileSystem(**storage_options)
+            )
+            file_info = azure_fs.get_file_info("/")
+            file_size = file_info.size if hasattr(file_info, "size") else 0
+
+            if self._should_use_chunking(file_size):
+                # Use PyArrow ParquetFile for batched reading of large files
+                parquet_file = ParquetFile(
+                    account_url, filesystem=fs.AzureFileSystem(**storage_options)
+                )
+                return self._read_parquet_in_batches(parquet_file)
+            else:
+                # Use pandas for regular reading of smaller files
+                dataframe = pd.read_parquet(
+                    account_url, storage_options=storage_options
+                )
+                return dataframe_to_chunks(dataframe)
+
+        except Exception:
+            # Fallback to regular pandas reading if size check or batching fails
+            dataframe = pd.read_parquet(account_url, storage_options=storage_options)
+            return dataframe_to_chunks(dataframe)
 
     @_read_parquet_dispatch.register
     def _(
@@ -114,10 +281,27 @@ class ParquetDataFrameReader(DataFrameReader):
         key: str,
         bucket_name: str,  # pylint: disable=unused-argument
     ) -> DatalakeColumnWrapper:
+        import os
+
         import pandas as pd  # pylint: disable=import-outside-toplevel
 
-        dataframe = pd.read_parquet(key)
-        return dataframe_to_chunks(dataframe)
+        # Check file size to determine reading strategy
+        try:
+            file_size = os.path.getsize(key)
+
+            if self._should_use_chunking(file_size):
+                # Use PyArrow ParquetFile for batched reading of large files
+                parquet_file = ParquetFile(key)
+                return self._read_parquet_in_batches(parquet_file)
+            else:
+                # Use pandas for regular reading of smaller files
+                dataframe = pd.read_parquet(key)
+                return dataframe_to_chunks(dataframe)
+
+        except Exception:
+            # Fallback to regular pandas reading if size check fails
+            dataframe = pd.read_parquet(key)
+            return dataframe_to_chunks(dataframe)
 
     def _read(self, *, key: str, bucket_name: str, **__) -> DatalakeColumnWrapper:
         return DatalakeColumnWrapper(
@@ -125,3 +309,6 @@ class ParquetDataFrameReader(DataFrameReader):
                 self.config_source, key=key, bucket_name=bucket_name
             )
         )
+
+    def _should_use_chunking(self, file_size: int) -> bool:
+        return file_size > MAX_FILE_SIZE_FOR_PREVIEW or file_size == 0