open-metadata--openmetadata/ingestion/tests/unit/readers/test_parquet_reader.py

#  Copyright 2025 Collate
#  Licensed under the Collate Community License, Version 1.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
#  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.

"""
Tests for ParquetDataFrameReader S3, GCS, and Local
"""

import tempfile
import unittest
from unittest.mock import MagicMock, Mock, patch

import pandas as pd

from metadata.generated.schema.entity.services.connections.database.datalake.gcsConfig import (
    GCSConfig,
)
from metadata.generated.schema.entity.services.connections.database.datalake.s3Config import (
    S3Config,
)
from metadata.generated.schema.entity.services.connections.database.datalakeConnection import (
    LocalConfig,
)
from metadata.generated.schema.security.credentials.awsCredentials import AWSCredentials
from metadata.readers.dataframe.base import MAX_FILE_SIZE_FOR_PREVIEW
from metadata.readers.dataframe.parquet import ParquetDataFrameReader


class TestParquetReader(unittest.TestCase):
    def test_local_small_parquet_file(self):
        df = pd.DataFrame({"id": [1, 2, 3], "name": ["Alice", "Bob", "Charlie"]})

        with tempfile.NamedTemporaryFile(suffix=".parquet", delete=False) as tmp:
            df.to_parquet(tmp.name)
            tmp_path = tmp.name

        try:
            config = LocalConfig()
            reader = ParquetDataFrameReader(config, None)

            result = reader._read(key=tmp_path, bucket_name="")

            self.assertIsNotNone(result.dataframes)
            dataframes = result.dataframes()
            self.assertIsNotNone(dataframes)
            chunks = list(dataframes)
            self.assertTrue(len(chunks) > 0)
            total_rows = sum(len(chunk) for chunk in chunks)
            self.assertEqual(total_rows, 3)
        finally:
            import os

            os.unlink(tmp_path)  # noqa: PTH108

    @patch("pyarrow.parquet.ParquetFile")
    @patch("os.path.getsize")
    def test_local_large_parquet_file_chunking(self, mock_getsize, mock_parquet_file):
        mock_getsize.return_value = MAX_FILE_SIZE_FOR_PREVIEW + 1000

        mock_pf = Mock()
        mock_parquet_file.return_value = mock_pf

        batch1_data = pd.DataFrame({"id": [1, 2], "name": ["Alice", "Bob"]})
        batch2_data = pd.DataFrame({"id": [3], "name": ["Charlie"]})

        mock_batch1 = Mock()
        mock_batch1.to_pandas.return_value = batch1_data
        mock_batch2 = Mock()
        mock_batch2.to_pandas.return_value = batch2_data

        mock_pf.iter_batches = Mock(return_value=iter([mock_batch1, mock_batch2]))

        config = LocalConfig()
        reader = ParquetDataFrameReader(config, None)

        result = reader._read(key="test.parquet", bucket_name="")

        self.assertIsNotNone(result.dataframes)
        dataframes = result.dataframes()
        self.assertIsNotNone(dataframes)
        chunks = list(dataframes)
        self.assertTrue(len(chunks) > 0)

    def _create_s3_reader(self):
        """Helper to create an S3 ParquetDataFrameReader with a mocked session."""
        from collections import namedtuple

        config = S3Config(
            securityConfig=AWSCredentials(awsAccessKeyId="test", awsSecretAccessKey="test", awsRegion="us-east-1")
        )
        mock_client = Mock()
        mock_session = Mock()

        FrozenCreds = namedtuple("FrozenCreds", ["access_key", "secret_key", "token"])
        mock_session.get_credentials.return_value.get_frozen_credentials.return_value = FrozenCreds(
            access_key="test", secret_key="test", token=None
        )

        reader = ParquetDataFrameReader(config, mock_client, session=mock_session)
        return reader, mock_client

    @patch("s3fs.S3FileSystem")
    @patch("pyarrow.parquet.ParquetFile")
    def test_s3_small_parquet_file(self, mock_parquet_file_cls, mock_s3fs):
        """Test S3 parquet reading uses credentials extracted from boto3 client."""
        reader, _ = self._create_s3_reader()

        mock_fs = MagicMock()
        mock_s3fs.return_value = mock_fs
        mock_fs.info.return_value = {"size": 1000}

        mock_df = pd.DataFrame({"id": [1, 2, 3], "name": ["Alice", "Bob", "Charlie"]})
        mock_table = Mock()
        mock_table.to_pandas.return_value = mock_df
        mock_pf = Mock()
        mock_pf.read.return_value = mock_table
        mock_parquet_file_cls.return_value = mock_pf

        result = reader._read(key="test.parquet", bucket_name="test-bucket")

        self.assertIsNotNone(result.dataframes)
        chunks = list(result.dataframes())
        total_rows = sum(len(chunk) for chunk in chunks)
        self.assertEqual(total_rows, 3)

        mock_s3fs.assert_called_once_with(
            key="test",
            secret="test",
            token=None,
            client_kwargs={"region_name": "us-east-1"},
        )
        mock_fs.open.assert_called_once_with("test-bucket/test.parquet")

    @patch("s3fs.S3FileSystem")
    @patch("pyarrow.parquet.ParquetFile")
    def test_s3_large_parquet_file_chunking(self, mock_parquet_file_cls, mock_s3fs):
        """Test S3 large parquet file triggers batched reading."""
        reader, _ = self._create_s3_reader()

        mock_fs = MagicMock()
        mock_s3fs.return_value = mock_fs
        mock_fs.info.return_value = {"size": MAX_FILE_SIZE_FOR_PREVIEW + 1000}

        mock_pf = Mock()
        mock_parquet_file_cls.return_value = mock_pf
        batch_data = pd.DataFrame({"id": [1], "name": ["Test"]})
        mock_batch = Mock()
        mock_batch.to_pandas.return_value = batch_data
        mock_pf.iter_batches = Mock(return_value=iter([mock_batch]))

        result = reader._read(key="test.parquet", bucket_name="test-bucket")

        self.assertIsNotNone(result.dataframes)
        chunks = list(result.dataframes())
        self.assertTrue(len(chunks) > 0)

    @patch("s3fs.S3FileSystem")
    def test_s3_file_size_error_falls_back_to_chunking(self, mock_s3fs):
        """Test that file size check failure falls back to chunked reading."""
        reader, _ = self._create_s3_reader()

        mock_fs = MagicMock()
        mock_s3fs.return_value = mock_fs
        mock_fs.info.side_effect = Exception("HeadObject failed")

        result = reader._read(key="test.parquet", bucket_name="test-bucket")

        self.assertIsNotNone(result.dataframes)

    @patch("gcsfs.GCSFileSystem")
    @patch("pyarrow.parquet.ParquetFile")
    def test_gcs_small_parquet_file(self, mock_parquet_file, mock_gcsfs):
        mock_gcs = Mock()
        mock_gcsfs.return_value = mock_gcs

        mock_gcs.info.return_value = {"size": 1000}

        mock_pf = Mock()
        mock_parquet_file.return_value = mock_pf

        mock_table = Mock()
        mock_df = pd.DataFrame({"id": [1], "name": ["Test"]})
        mock_table.to_pandas.return_value = mock_df
        mock_pf.read.return_value = mock_table

        config = GCSConfig()
        reader = ParquetDataFrameReader(config, None)

        result = reader._read(key="test.parquet", bucket_name="test-bucket")

        self.assertIsNotNone(result.dataframes)
        dataframes = result.dataframes()
        self.assertIsNotNone(dataframes)
        chunks = list(dataframes)
        self.assertTrue(len(chunks) > 0)

    @patch("gcsfs.GCSFileSystem")
    @patch("pyarrow.parquet.ParquetFile")
    def test_gcs_large_parquet_file_chunking(self, mock_parquet_file, mock_gcsfs):
        mock_gcs = Mock()
        mock_gcsfs.return_value = mock_gcs

        mock_gcs.info.return_value = {"size": MAX_FILE_SIZE_FOR_PREVIEW + 1000}

        mock_pf = Mock()
        mock_parquet_file.return_value = mock_pf

        batch_data = pd.DataFrame({"id": [1], "name": ["Test"]})
        mock_batch = Mock()
        mock_batch.to_pandas.return_value = batch_data

        mock_pf.iter_batches = Mock(return_value=iter([mock_batch]))

        config = GCSConfig()
        reader = ParquetDataFrameReader(config, None)

        result = reader._read(key="test.parquet", bucket_name="test-bucket")

        self.assertIsNotNone(result.dataframes)
        dataframes = result.dataframes()
        self.assertIsNotNone(dataframes)
        chunks = list(dataframes)
        self.assertTrue(len(chunks) > 0)

    def test_should_use_chunking_logic(self):
        config = LocalConfig()
        reader = ParquetDataFrameReader(config, None)

        self.assertTrue(reader._should_use_chunking(MAX_FILE_SIZE_FOR_PREVIEW + 1))
        self.assertFalse(reader._should_use_chunking(MAX_FILE_SIZE_FOR_PREVIEW - 1))
        self.assertTrue(reader._should_use_chunking(0))


if __name__ == "__main__":
    unittest.main()