# 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. """ Test datalake utils """ import os from unittest import TestCase import pandas as pd from metadata.generated.schema.entity.data.table import Column, DataType from metadata.readers.dataframe.reader_factory import SupportedTypes from metadata.utils.datalake.datalake_utils import ( DataFrameColumnParser, GenericDataFrameColumnParser, ParquetDataFrameColumnParser, get_file_format_type, ) STRUCTURE = { "a": "w", "b": 4, "c": { "d": 2, "e": 4, "f": { "g": 9, "h": {"i": 6}, "n": { "o": 10, "p": 11, }, }, "j": 7, "k": 8, }, } class TestDatalakeUtils(TestCase): """class for datalake utils test""" def test_unique_json_structure(self): """test unique json structure fn""" sample_data = [ {"a": "x", "b": 1, "c": {"d": 2}}, {"a": "y", "b": 2, "c": {"e": 4, "f": {"g": 5, "h": {"i": 6}, "n": 5}}}, {"a": "z", "b": 3, "c": {"j": 7}}, {"a": "w", "b": 4, "c": {"k": 8, "f": {"g": 9, "n": {"o": 10, "p": 11}}}}, ] expected = STRUCTURE actual = GenericDataFrameColumnParser.unique_json_structure(sample_data) self.assertDictEqual(expected, actual) def test_construct_column(self): """test construct column fn""" expected = [ { "dataTypeDisplay": "STRING", "dataType": "STRING", "name": "a", "displayName": "a", }, { "dataTypeDisplay": "INT", "dataType": "INT", "name": "b", "displayName": "b", }, { "dataTypeDisplay": "JSON", "dataType": "JSON", "name": "c", "displayName": "c", "children": [ { "dataTypeDisplay": "INT", "dataType": "INT", "name": "d", "displayName": "d", }, { "dataTypeDisplay": "INT", "dataType": "INT", "name": "e", "displayName": "e", }, { "dataTypeDisplay": "JSON", "dataType": "JSON", "name": "f", "displayName": "f", "children": [ { "dataTypeDisplay": "INT", "dataType": "INT", "name": "g", "displayName": "g", }, { "dataTypeDisplay": "JSON", "dataType": "JSON", "name": "h", "displayName": "h", "children": [ { "dataTypeDisplay": "INT", "dataType": "INT", "name": "i", "displayName": "i", } ], }, { "dataTypeDisplay": "JSON", "dataType": "JSON", "name": "n", "displayName": "n", "children": [ { "dataTypeDisplay": "INT", "dataType": "INT", "name": "o", "displayName": "o", }, { "dataTypeDisplay": "INT", "dataType": "INT", "name": "p", "displayName": "p", }, ], }, ], }, { "dataTypeDisplay": "INT", "dataType": "INT", "name": "j", "displayName": "j", }, { "dataTypeDisplay": "INT", "dataType": "INT", "name": "k", "displayName": "k", }, ], }, ] actual = GenericDataFrameColumnParser.construct_json_column_children(STRUCTURE) for el in zip(expected, actual): self.assertDictEqual(el[0], el[1]) def test_create_column_object(self): """test create column object fn""" formatted_column = GenericDataFrameColumnParser.construct_json_column_children( STRUCTURE ) column = { "dataTypeDisplay": "STRING", "dataType": "STRING", "name": "a", "displayName": "a", "children": formatted_column, } column_obj = Column(**column) assert len(column_obj.children) == 3 class TestParquetDataFrameColumnParser(TestCase): """Test parquet dataframe column parser""" @classmethod def setUpClass(cls) -> None: resources_path = os.path.join( os.path.dirname(os.path.dirname(__file__)), "resources" ) cls.parquet_path = os.path.join(resources_path, "datalake", "example.parquet") cls.df = pd.read_parquet(cls.parquet_path) cls.parquet_parser = ParquetDataFrameColumnParser(cls.df) def test_parser_instantiation(self): """Test the right parser is instantiated from the creator method""" parquet_parser = DataFrameColumnParser.create(self.df, SupportedTypes.PARQUET) self.assertIsInstance(parquet_parser.parser, ParquetDataFrameColumnParser) parquet_types = [ SupportedTypes.PARQUET, SupportedTypes.PARQUET_PQ, SupportedTypes.PARQUET_PQT, SupportedTypes.PARQUET_PARQ, SupportedTypes.PARQUET_SNAPPY, ] other_types = [typ for typ in SupportedTypes if typ not in parquet_types] for other_type in other_types: with self.subTest(other_type=other_type): generic_parser = DataFrameColumnParser.create(self.df, other_type) self.assertIsInstance( generic_parser.parser, GenericDataFrameColumnParser ) def test_shuffle_and_sample_from_parser(self): """test the shuffle and sampling logic from the parser creator method""" parquet_parser = DataFrameColumnParser.create(self.df, SupportedTypes.PARQUET) self.assertEqual(parquet_parser.parser.data_frame.shape, self.df.shape) parquet_parser = DataFrameColumnParser.create( [self.df, self.df], SupportedTypes.PARQUET ) self.assertEqual(parquet_parser.parser.data_frame.shape, self.df.shape) parquet_parser = DataFrameColumnParser.create( [self.df, self.df], SupportedTypes.PARQUET, sample=False ) self.assertEqual( parquet_parser.parser.data_frame.shape, pd.concat([self.df, self.df]).shape ) def test_get_columns(self): """test `get_columns` method of the parquet column parser""" expected = [ Column( dataTypeDisplay="bool", dataType=DataType.BOOLEAN, name="a", displayName="a", ), # type: ignore Column( dataTypeDisplay="int8", dataType=DataType.INT, name="b", displayName="b", ), # type: ignore Column( dataTypeDisplay="int16", dataType=DataType.INT, name="c", displayName="c", ), # type: ignore Column( dataTypeDisplay="int32", dataType=DataType.INT, name="d", displayName="d", ), # type: ignore Column( dataTypeDisplay="int64", dataType=DataType.INT, name="e", displayName="e", ), # type: ignore Column( dataTypeDisplay="uint8", dataType=DataType.UINT, name="f", displayName="f", ), # type: ignore Column( dataTypeDisplay="uint16", dataType=DataType.UINT, name="g", displayName="g", ), # type: ignore Column( dataTypeDisplay="uint32", dataType=DataType.UINT, name="h", displayName="h", ), # type: ignore Column( dataTypeDisplay="uint64", dataType=DataType.UINT, name="i", displayName="i", ), # type: ignore Column( dataTypeDisplay="float", dataType=DataType.FLOAT, name="k", displayName="k", ), # type: ignore Column( dataTypeDisplay="double", dataType=DataType.FLOAT, name="l", displayName="l", ), # type: ignore Column( dataTypeDisplay="time64[us]", dataType=DataType.DATETIME, name="n", displayName="n", ), # type: ignore Column( dataTypeDisplay="timestamp[ns]", dataType=DataType.DATETIME, name="o", displayName="o", ), # type: ignore Column( dataTypeDisplay="date32[day]", dataType=DataType.DATE, name="p", displayName="p", ), # type: ignore Column( dataTypeDisplay="date32[day]", dataType=DataType.DATE, name="q", displayName="q", ), # type: ignore Column( dataTypeDisplay="duration[ns]", dataType=DataType.INT, name="r", displayName="r", ), # type: ignore Column( dataTypeDisplay="binary", dataType=DataType.BINARY, name="t", displayName="t", ), # type: ignore Column( dataTypeDisplay="string", dataType=DataType.STRING, name="u", displayName="u", ), # type: ignore Column( dataTypeDisplay="string", dataType=DataType.STRING, name="v", displayName="v", ), # type: ignore Column( dataTypeDisplay="binary", dataType=DataType.BINARY, name="w", displayName="w", ), # type: ignore Column( dataTypeDisplay="string", dataType=DataType.STRING, name="x", displayName="x", ), # type: ignore Column( dataTypeDisplay="string", dataType=DataType.STRING, name="y", displayName="y", ), # type: ignore Column( dataTypeDisplay="list", dataType=DataType.ARRAY, name="aa", displayName="aa", ), # type: ignore Column( dataTypeDisplay="list", dataType=DataType.ARRAY, name="bb", displayName="bb", ), # type: ignore Column( dataTypeDisplay="struct>>", dataType=DataType.STRUCT, name="dd", displayName="dd", children=[ Column( dataTypeDisplay="int64", dataType=DataType.INT, name="ee", displayName="ee", ), # type: ignore Column( dataTypeDisplay="int64", dataType=DataType.INT, name="ff", displayName="ff", ), # type: ignore Column( dataTypeDisplay="struct>", dataType=DataType.STRUCT, name="gg", displayName="gg", children=[ Column( dataTypeDisplay="struct", dataType=DataType.STRUCT, name="hh", displayName="hh", children=[ Column( dataTypeDisplay="int64", dataType=DataType.INT, name="ii", displayName="ii", ), # type: ignore Column( dataTypeDisplay="int64", dataType=DataType.INT, name="jj", displayName="jj", ), # type: ignore Column( dataTypeDisplay="int64", dataType=DataType.INT, name="kk", displayName="kk", ), # type: ignore ], ), ], ), ], ), # type: ignore ] actual = self.parquet_parser.get_columns() for validation in zip(expected, actual): with self.subTest(validation=validation): expected_col, actual_col = validation self.assertEqual(expected_col.name, actual_col.name) self.assertEqual(expected_col.displayName, actual_col.displayName) self.assertEqual(expected_col.dataType, actual_col.dataType) def _validate_parsed_column(self, expected, actual): """validate parsed column""" self.assertEqual(expected.name, actual.name) self.assertEqual(expected.dataType, actual.dataType) self.assertEqual(expected.displayName, actual.displayName) if expected.children: self.assertEqual(len(expected.children), len(actual.children)) for validation in zip(expected.children, actual.children): with self.subTest(validation=validation): expected_col, actual_col = validation self._validate_parsed_column(expected_col, actual_col) def test_get_file_format_type_csv_gz(self): """test get_file_format_type function for csv.gz files""" # Test csv.gz file detection result = get_file_format_type("data.csv.gz") self.assertEqual(result, SupportedTypes.CSVGZ) # Test regular csv file detection (should still work) result = get_file_format_type("data.csv") self.assertEqual(result, SupportedTypes.CSV) # Test other gzipped files result = get_file_format_type("data.json.gz") self.assertEqual(result, SupportedTypes.JSONGZ) # Test unsupported gzipped format result = get_file_format_type("data.txt.gz") self.assertEqual(result, False) def test_csv_gz_file_format_detection_edge_cases(self): """test edge cases for csv.gz file format detection""" # Test with nested paths result = get_file_format_type("folder/subfolder/data.csv.gz") self.assertEqual(result, SupportedTypes.CSVGZ) # Test with multiple dots result = get_file_format_type("data.backup.csv.gz") self.assertEqual(result, SupportedTypes.CSVGZ) # Test with no extension result = get_file_format_type("data") self.assertEqual(result, False) # Test with just .gz result = get_file_format_type("data.gz") self.assertEqual(result, False) def test_csv_gz_compression_detection(self): """test compression detection for various file types""" # Test csv.gz compression detection test_cases = [ ("data.csv.gz", SupportedTypes.CSVGZ), ("data.csv", SupportedTypes.CSV), ("data.json.gz", SupportedTypes.JSONGZ), ("data.json", SupportedTypes.JSON), ("data.jsonl.gz", SupportedTypes.JSONLGZ), ("data.jsonl", SupportedTypes.JSONL), ("data.parquet", SupportedTypes.PARQUET), ("data.txt.gz", False), # Unsupported ("data.unknown.gz", False), # Unsupported ] for filename, expected in test_cases: with self.subTest(filename=filename): result = get_file_format_type(filename) self.assertEqual(result, expected, f"Failed for {filename}") def test_csv_gz_reader_factory_integration(self): """test that csv.gz is properly integrated with reader factory""" from metadata.readers.dataframe.reader_factory import SupportedTypes # Test that CSVGZ is properly handled try: # Test that the enum value exists self.assertEqual(SupportedTypes.CSVGZ.value, "csv.gz") # Test that it's different from regular CSV self.assertNotEqual(SupportedTypes.CSVGZ, SupportedTypes.CSV) self.assertNotEqual(SupportedTypes.CSVGZ.value, SupportedTypes.CSV.value) except Exception as e: self.fail(f"CSVGZ enum test failed: {e}") def test_csv_gz_supported_types_enum(self): """test that CSVGZ is properly defined in SupportedTypes enum""" # Test that CSVGZ exists in the enum self.assertIn(SupportedTypes.CSVGZ, SupportedTypes) self.assertEqual(SupportedTypes.CSVGZ.value, "csv.gz") # Test that it's different from regular CSV self.assertNotEqual(SupportedTypes.CSVGZ, SupportedTypes.CSV) self.assertNotEqual(SupportedTypes.CSVGZ.value, SupportedTypes.CSV.value) def test_csv_gz_dsv_reader_compression_detection(self): """test that DSV reader properly detects compression for csv.gz files""" from metadata.generated.schema.entity.services.connections.database.datalakeConnection import ( LocalConfig, ) from metadata.readers.dataframe.dsv import DSVDataFrameReader # Create a mock config local_config = LocalConfig() # Create DSV reader reader = DSVDataFrameReader(config_source=local_config, client=None) # Test compression detection logic (this is the same logic used in the dispatch methods) test_cases = [ ("data.csv.gz", "gzip"), ("data.csv", None), ("data.json.gz", "gzip"), ("data.txt.gz", "gzip"), ("data.unknown.gz", "gzip"), ] for filename, expected_compression in test_cases: with self.subTest(filename=filename): # Simulate the compression detection logic from the dispatch methods compression = None if filename.endswith(".gz"): compression = "gzip" self.assertEqual( compression, expected_compression, f"Compression detection failed for {filename}", ) def test_csv_gz_integration_completeness(self): """test that csv.gz support is complete across all components""" # Test that CSVGZ is in the reader factory mapping from metadata.readers.dataframe.reader_factory import ( DF_READER_MAP, SupportedTypes, ) # Check that CSVGZ is mapped to CSVDataFrameReader self.assertIn(SupportedTypes.CSVGZ.value, DF_READER_MAP) # Test that the get_df_reader function includes CSVGZ in DSV handling # This should not raise an exception for CSVGZ try: # Test that CSVGZ is included in the DSV types dsv_types = {SupportedTypes.CSV, SupportedTypes.CSVGZ, SupportedTypes.TSV} self.assertIn(SupportedTypes.CSVGZ, dsv_types) except Exception as e: self.fail(f"CSVGZ integration test failed: {e}")