Fix #24886: CSV OOM Issue (#24905)

* fix: OOM issue when processing large CSV files by allowing a maximum of 5 chunks

* Revert "fix: OOM issue when processing large CSV files by allowing a maximum of 5 chunks"

This reverts commit aa8dda1.

* fix: use iterator to avoid OOM issues for large files

* fix: parquet and datalake tests

* fix: optimized the file reading code for all the support file types

* fix: s3fs version and json file reading

* test: added and fixed file reading tests

* fix: generate close on read first chunk

* refact: remove iterator instance check fromt tests

* refactor(datalake): Convert DataFrame API from list-based to generator-based

  Refactors the pandas profiler and datalake readers to use lazy
  generator-based
  DataFrame iteration instead of loading entire datasets into memory. This    improves memory efficiency for large files and enables streaming processing.

  ## Core API Changes

  ### DataFrame Readers (dataframes now returns callable)
  - base.py: Updated read_first_chunk() to handle callable dataframes
  - parquet.py: All storage backends (S3, GCS, Azure, Local) return callables
  - dsv.py: CSV/TSV readers return callables
  - json.py: JSON/JSONL readers return callables
  - avro.py: All storage backends return callables (fixed S3 reader)
  - mf4.py: Fixed empty case to return callable

  ### Profiler Interface
  - profiler_interface.py: Added _type_casted_dataset() wrapper that applies
    type casting lazily; dataset is now a callable generator factory
  - runner.py: Updated PandasRunner to work with callable dataset

  ### Metrics (updated to iterate over dataset)
  - All static metrics (count, min, max, sum, mean, stddev, etc.)
  - All window metrics (median, first_quartile, third_quartile)
  - Hybrid metrics (histogram, cardinality_distribution)

  ### Data Quality Validators
  - pandas_validator_mixin.py: Updated to iterate over dataframes
  - base_test_handler.py: Handle callable dataframes
  - Table validators: tableColumnCount*, tableColumnNameToExist,
    tableColumnToMatchSet, tableCustomSQLQuery

  ### Utilities
  - datalake_utils.py: fetch_dataframe() and fetch_dataframe_first_chunk()
    now call the dataframes callable before iterating
  - pandas_mixin.py: Updated partitioning logic
  - sampler.py: Updated DatalakeSampler for generator-based dataset

  ## Test Updates
  - test_profiler_interface.py, test_profiler.py, test_sample.py,
    test_custom_metrics.py, test_datalake_metrics.py: Mock get_dataframes()
  - test_*_reader.py: Use list(result.dataframes()) pattern
  - test_parquet_azure_reader.py: Consume generator before mock assertions

* style: ran python linting

* refactor(datalake): addressed issues after initial review

* refactor: add logging when error

* fix: sdk validator

* style: ran java style check

* fix(observability): failing DQ tests

* fix: tests related to callable dataframes

* fix: failing tests

* fix: _stream_json_lines - use file_obj.readline()

* fix: pymssql version

* fix: ijson version

* fix: gitar typehinting comment

* fix: removed fetch_dataframe method

* fix: setup.py libraries to fix timeout

---------

Co-authored-by: TeddyCr <teddy.crepineau@gmail.com>

Author: keshavmohta09

ingestion/setup.py

@@ -22,15 +22,17 @@
 VERSIONS = {
     "airflow": "apache-airflow==3.1.5",
     "adlfs": "adlfs>=2023.1.0",
+    "aiobotocore": "aiobotocore~=2.26.0",
     "avro": "avro>=1.11.4,<1.12",
-    "boto3": "boto3>=1.20,<2.0",  # No need to add botocore separately. It's a dep from boto3
+    "boto3": "boto3~=1.41.5",
     "geoalchemy2": "GeoAlchemy2~=0.12",
     "google-cloud-monitoring": "google-cloud-monitoring>=2.0.0",
     "google-cloud-storage": "google-cloud-storage>=1.43.0",
-    "gcsfs": "gcsfs>=2023.1.0",
+    "gcsfs": "gcsfs~=2023.12.1",
     "great-expectations": "great-expectations~=0.18.0",
     "great-expectations-1xx": "great-expectations~=1.0",
     "grpc-tools": "grpcio-tools>=1.47.2",
+    "ijson": "ijson~=3.4",
     "msal": "msal~=1.2",
     "neo4j": "neo4j~=5.3",
     "pandas": "pandas~=2.0.3",
@@ -56,6 +58,7 @@
     "mongo": "pymongo~=4.3",
     "redshift": "sqlalchemy-redshift==0.8.12",
     "snowflake": "snowflake-sqlalchemy~=1.4",
+    "snowflake-connector": "snowflake-connector-python~=3.18.0",
     "elasticsearch8": "elasticsearch8~=8.9.0",
     "giturlparse": "giturlparse",
     "validators": "validators~=0.22.0",
@@ -66,7 +69,8 @@
     "pydoris": "pydoris-custom>=1.0.2,<1.5",
     "pyiceberg": "pyiceberg==0.5.1",
     "google-cloud-bigtable": "google-cloud-bigtable>=2.0.0",
-    "pyathena": "pyathena~=3.0",
+    "pyathena": "pyathena~=3.25.0",
+    "s3fs": "s3fs~=2023.12.1",
     "sqlalchemy-bigquery": "sqlalchemy-bigquery~=1.15.0",
     "presidio-analyzer": "presidio-analyzer==2.2.358",
     "asammdf": "asammdf~=7.4.5",
@@ -80,12 +84,14 @@
         VERSIONS["asammdf"],
         VERSIONS["avro"],
         VERSIONS["boto3"],
+        VERSIONS["ijson"],
         VERSIONS["pandas"],
         VERSIONS["pyarrow"],
         VERSIONS["numpy"],
         # python-snappy does not work well on 3.11 https://github.com/aio-libs/aiokafka/discussions/931
         # Using this as an alternative
         "cramjam~=2.7",
+        "fastavro>=1.2.0",
     },
     "hive": {
         "pure-transport==0.2.0",
@@ -167,7 +173,7 @@
     "collate-data-diff>=0.11.9",
     "jaraco.functools<4.2.0",  # above 4.2 breaks the build
     # TODO: Remove one once we have updated datadiff version
-    "snowflake-connector-python>=3.13.1,<4.0.0",
+    VERSIONS["snowflake-connector"],
     "mysql-connector-python>=8.0.29;python_version<'3.9'",
     "mysql-connector-python>=9.1;python_version>='3.9'",
     "httpx~=0.28.0",
@@ -233,15 +239,19 @@
         VERSIONS["azure-storage-blob"],
         VERSIONS["azure-identity"],
         VERSIONS["adlfs"],
+        VERSIONS["aiobotocore"],
         *COMMONS["datalake"],
     },
     "datalake-gcs": {
         VERSIONS["google-cloud-monitoring"],
         VERSIONS["google-cloud-storage"],
         VERSIONS["gcsfs"],
+        VERSIONS["aiobotocore"],
         *COMMONS["datalake"],
     },
     "datalake-s3": {
+        VERSIONS["s3fs"],
+        VERSIONS["aiobotocore"],
         *COMMONS["datalake"],
     },
     "deltalake": {
@@ -252,7 +262,7 @@
     "deltalake-storage": {"deltalake>=0.19.0,<0.20"},
     "deltalake-spark": {"delta-spark>=3.0.0,<4.0.0", "pyspark==3.5.6"},
     "domo": {VERSIONS["pydomo"]},
-    "doris": {"pydoris==1.0.2"},
+    "doris": {VERSIONS["pydoris"]},
     "druid": {"pydruid>=0.6.5"},
     "dynamodb": {VERSIONS["boto3"]},
     "elasticsearch": {
@@ -335,7 +345,7 @@
     },
     "qliksense": {"websocket-client~=1.6.1"},
     "presto": {*COMMONS["hive"], DATA_DIFF["presto"]},
-    "pymssql": {"pymssql~=2.2.0"},
+    "pymssql": {"pymssql~=2.3.9"},
     "quicksight": {VERSIONS["boto3"]},
     "redash": {VERSIONS["packaging"]},
     "redpanda": {*COMMONS["kafka"]},

ingestion/src/metadata/data_quality/builders/validator_builder.py

@@ -28,7 +28,9 @@
 from metadata.utils.importer import import_test_case_class
 
 if TYPE_CHECKING:
-    from pandas import DataFrame
+    from metadata.data_quality.interface.pandas.pandas_test_suite_interface import (
+        PandasRunner,
+    )
 
 
 class TestCaseImporter:
@@ -51,15 +53,15 @@ class ValidatorBuilder(TestCaseImporter):
 
     def __init__(
         self,
-        runner: Union[QueryRunner, "DataFrame"],
+        runner: Union[QueryRunner, "PandasRunner"],
         test_case: TestCase,
         source_type: SourceType,
         entity_type: str,
     ) -> None:
         """Builder object for SQA validators. This builder is used to create a validator object
 
         Args:
-            runner (QueryRunner): The runner object
+            runner (Union[QueryRunner, PandasRunner]): The runner object
             test_case (TestCase): The test case object
             source_type (SourceType): The source type
             entity_type (str): one of COLUMN or TABLE -- fetched from the test definition

ingestion/src/metadata/data_quality/validations/base_test_handler.py

@@ -43,12 +43,11 @@
 from metadata.generated.schema.tests.dimensionResult import DimensionResult
 from metadata.generated.schema.tests.testCase import TestCase, TestCaseParameterValue
 from metadata.generated.schema.type.basic import Timestamp
-from metadata.profiler.processor.runner import QueryRunner
+from metadata.profiler.processor.runner import PandasRunner, QueryRunner
 from metadata.utils.logger import test_suite_logger
 from metadata.utils.sqa_like_column import SQALikeColumn
 
 if TYPE_CHECKING:
-    from pandas import DataFrame
     from sqlalchemy import Column
 
 logger = test_suite_logger()
@@ -105,7 +104,7 @@ class BaseTestValidator(ABC):
 
     def __init__(
         self,
-        runner: Union[QueryRunner, List["DataFrame"]],
+        runner: Union[QueryRunner, PandasRunner],
         test_case: TestCase,
         execution_date: Timestamp,
     ) -> None:
@@ -148,10 +147,7 @@ def run_validation(self) -> TestCaseResult:
             )
             logger.debug(f"Dimension columns: {self.test_case.dimensionColumns}")
 
-            # Validate dimension columns exist in the target table
             if not self.are_dimension_columns_valid():
-                # Don't abort the main test, just skip dimensional validation
-                # The main test result is still valid
                 return test_result
 
             try:

ingestion/src/metadata/data_quality/validations/column/pandas/columnValueLengthsToBeBetween.py

@@ -89,7 +89,7 @@ def _execute_dimensional_validation(
         dimension_results = []
 
         try:
-            dfs = self.runner if isinstance(self.runner, list) else [self.runner]
+            dfs = self.runner
             min_impl = Metrics.MIN_LENGTH(column).get_pandas_computation()
             max_impl = Metrics.MAX_LENGTH(column).get_pandas_computation()
             row_count_impl = Metrics.ROW_COUNT().get_pandas_computation()

ingestion/src/metadata/data_quality/validations/column/pandas/columnValueMaxToBeBetween.py

@@ -87,7 +87,7 @@ def _execute_dimensional_validation(
         dimension_results = []
 
         try:
-            dfs = self.runner if isinstance(self.runner, list) else [self.runner]
+            dfs = self.runner
             max_impl = Metrics.MAX(column).get_pandas_computation()
 
             dimension_aggregates = defaultdict(

ingestion/src/metadata/data_quality/validations/column/pandas/columnValueMeanToBeBetween.py

@@ -91,7 +91,7 @@ def _execute_dimensional_validation(
         dimension_results = []
 
         try:
-            dfs = self.runner if isinstance(self.runner, list) else [self.runner]
+            dfs = self.runner
             mean_impl = Metrics.MEAN(column).get_pandas_computation()
 
             dimension_aggregates = defaultdict(

ingestion/src/metadata/data_quality/validations/column/pandas/columnValueMedianToBeBetween.py

@@ -92,7 +92,7 @@ def _execute_dimensional_validation(
         dimension_results = []
 
         try:
-            dfs = self.runner if isinstance(self.runner, list) else [self.runner]
+            dfs = self.runner
             median_impl = Metrics.MEDIAN(column).get_pandas_computation()
 
             dimension_aggregates = defaultdict(

ingestion/src/metadata/data_quality/validations/column/pandas/columnValueMinToBeBetween.py

@@ -88,7 +88,7 @@ def _execute_dimensional_validation(
         dimension_results = []
 
         try:
-            dfs = self.runner if isinstance(self.runner, list) else [self.runner]
+            dfs = self.runner
             min_impl = Metrics.MIN(column).get_pandas_computation()
 
             dimension_aggregates = defaultdict(

ingestion/src/metadata/data_quality/validations/column/pandas/columnValueStdDevToBeBetween.py

@@ -94,7 +94,7 @@ def _execute_dimensional_validation(
         dimension_results = []
 
         try:
-            dfs = self.runner if isinstance(self.runner, list) else [self.runner]
+            dfs = self.runner
             stddev_impl = Metrics.STDDEV(column).get_pandas_computation()
             row_count_impl = Metrics.ROW_COUNT().get_pandas_computation()
 

ingestion/src/metadata/data_quality/validations/column/pandas/columnValuesMissingCount.py

@@ -89,7 +89,7 @@ def _execute_dimensional_validation(
         dimension_results = []
 
         try:
-            dfs = self.runner if isinstance(self.runner, list) else [self.runner]
+            dfs = self.runner
 
             metric_expressions = {
                 Metrics.NULL_MISSING_COUNT.name: Metrics.NULL_MISSING_COUNT(