openmetadata-ingestion 1.8.7.1__py3-none-any.whl → 1.8.8.0__py3-none-any.whl

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

metadata/sampler/sqlalchemy/sampler.py

@@ -13,7 +13,6 @@ Helper module to handle data sampling
 for the profiler
 """
 import hashlib
-from contextlib import contextmanager
 from typing import List, Optional, Union, cast
 
 from sqlalchemy import Column, inspect, text
@@ -75,7 +74,7 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         self._table = self.build_table_orm(
             self.entity, self.service_connection_config, self.ometa_client
         )
-        self._active_sessions = set()
+        self.session_factory = create_and_bind_thread_safe_session(self.connection)
 
     @property
     def raw_dataset(self):
@@ -83,21 +82,7 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
 
     def get_client(self):
         """Build the SQA Client"""
-        session_factory = create_and_bind_thread_safe_session(self.connection)
-        session = session_factory()
-        self._active_sessions.add(session)
-        return session
-
-    @contextmanager
-    def get_client_context(self):
-        """Get client as context manager for proper cleanup"""
-        session = self.get_client()
-        try:
-            yield session
-        finally:
-            if session in self._active_sessions:
-                self._active_sessions.remove(session)
-            session.close()
+        return self.session_factory()
 
     def set_tablesample(self, selectable: Table):
         """Set the tablesample for the table. To be implemented by the child SQA sampler class
@@ -147,17 +132,17 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         """
         # only sample the column if we are computing a column metric to limit the amount of data scaned
         selectable = self.set_tablesample(self.raw_dataset.__table__)
-        client = self.get_client()
 
-        entity = selectable if column is None else selectable.c.get(column.key)
-        if label is not None:
-            query = client.query(entity, label)
-        else:
-            query = client.query(entity)
+        with self.session_factory() as client:
+            entity = selectable if column is None else selectable.c.get(column.key)
+            if label is not None:
+                query = client.query(entity, label)
+            else:
+                query = client.query(entity)
 
-        if self.partition_details:
-            query = self.get_partitioned_query(query)
-        return query
+            if self.partition_details:
+                query = self.get_partitioned_query(query)
+            return query
 
     def get_sampler_table_name(self) -> str:
         """Get the base name of the SQA table for sampling.
@@ -172,34 +157,34 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
 
     def get_sample_query(self, *, column=None) -> Query:
         """get query for sample data"""
-        client = self.get_client()
-        if self.sample_config.profileSampleType == ProfileSampleType.PERCENTAGE:
-            rnd = self._base_sample_query(
+        with self.session_factory() as client:
+            if self.sample_config.profileSampleType == ProfileSampleType.PERCENTAGE:
+                rnd = self._base_sample_query(
+                    column,
+                    (ModuloFn(RandomNumFn(), 100)).label(RANDOM_LABEL),
+                ).cte(f"{self.get_sampler_table_name()}_rnd")
+                session_query = client.query(rnd)
+                return session_query.where(
+                    rnd.c.random <= self.sample_config.profileSample
+                ).cte(f"{self.get_sampler_table_name()}_sample")
+
+            table_query = client.query(self.raw_dataset)
+            if self.partition_details:
+                table_query = self.get_partitioned_query(table_query)
+            session_query = self._base_sample_query(
                 column,
-                (ModuloFn(RandomNumFn(), 100)).label(RANDOM_LABEL),
-            ).cte(f"{self.get_sampler_table_name()}_rnd")
-            session_query = client.query(rnd)
-            return session_query.where(
-                rnd.c.random <= self.sample_config.profileSample
-            ).cte(f"{self.get_sampler_table_name()}_sample")
-
-        table_query = client.query(self.raw_dataset)
-        if self.partition_details:
-            table_query = self.get_partitioned_query(table_query)
-        session_query = self._base_sample_query(
-            column,
-            (ModuloFn(RandomNumFn(), table_query.count())).label(RANDOM_LABEL)
-            if self.sample_config.randomizedSample
-            else None,
-        )
-        query = (
-            session_query.order_by(RANDOM_LABEL)
-            if self.sample_config.randomizedSample
-            else session_query
-        )
-        return query.limit(self.sample_config.profileSample).cte(
-            f"{self.get_sampler_table_name()}_rnd"
-        )
+                (ModuloFn(RandomNumFn(), table_query.count())).label(RANDOM_LABEL)
+                if self.sample_config.randomizedSample
+                else None,
+            )
+            query = (
+                session_query.order_by(RANDOM_LABEL)
+                if self.sample_config.randomizedSample
+                else session_query
+            )
+            return query.limit(self.sample_config.profileSample).cte(
+                f"{self.get_sampler_table_name()}_rnd"
+            )
 
     def get_dataset(self, column=None, **__) -> Union[DeclarativeMeta, AliasedClass]:
         """
@@ -247,7 +232,7 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
                 if col.name != RANDOM_LABEL and col.name in names
             ]
 
-        with self.get_client_context() as client:
+        with self.session_factory() as client:
 
             # Handle array columns with special query modification
             max_elements = self._get_max_array_elements()
@@ -298,7 +283,7 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         if not is_safe_sql_query(self.sample_query):
             raise RuntimeError(f"SQL expression is not safe\n\n{self.sample_query}")
 
-        with self.get_client_context() as client:
+        with self.session_factory() as client:
             rnd = client.execute(f"{self.sample_query}")
         try:
             columns = [col.name for col in rnd.cursor.description]
@@ -316,11 +301,11 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
 
         stmt = text(f"{self.sample_query}")
         stmt = stmt.columns(*list(inspect(self.raw_dataset).c))
-        return (
-            self.get_client()
-            .query(stmt.subquery())
-            .cte(f"{self.get_sampler_table_name()}_user_sampled")
-        )
+
+        with self.session_factory() as client:
+            return client.query(stmt.subquery()).cte(
+                f"{self.get_sampler_table_name()}_user_sampled"
+            )
 
     def _partitioned_table(self) -> Query:
         """Return the Query object for partitioned tables"""
@@ -337,7 +322,9 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
         )
         if query is not None:
             return query.filter(partition_filter)
-        return self.get_client().query(self.raw_dataset).filter(partition_filter)
+
+        with self.session_factory() as client:
+            return client.query(self.raw_dataset).filter(partition_filter)
 
     def get_columns(self):
         """get columns from entity"""
@@ -345,6 +332,15 @@ class SQASampler(SamplerInterface, SQAInterfaceMixin):
 
     def close(self):
         """Close the connection"""
-        for session in self._active_sessions:
-            session.close()
-        self.connection.pool.dispose()
+        try:
+            self.connection.pool.dispose()
+        except Exception as e:
+            logger.warning(f"Error disposing connection pool: {e}")
+
+    def __del__(self):
+        """Destructor to ensure cleanup when object is garbage collected"""
+        try:
+            self.close()
+        except Exception:
+            # Ignore errors during cleanup in destructor
+            pass

metadata/sampler/sqlalchemy/snowflake/sampler.py

@@ -91,5 +91,6 @@ class SnowflakeSampler(SQASampler):
         rnd = self._base_sample_query(column).cte(
             f"{self.get_sampler_table_name()}_rnd"
         )
-        query = self.get_client().query(rnd)
+        with self.session_factory() as client:
+            query = client.query(rnd)
         return query.cte(f"{self.get_sampler_table_name()}_sample")