diff --git a/.bumpversion.cfg b/.bumpversion.cfg index efdb1af..7245553 100644 --- a/.bumpversion.cfg +++ b/.bumpversion.cfg @@ -1,5 +1,5 @@ [bumpversion] -current_version = 0.14.1 +current_version = 0.15.0 commit = False tag = False parse = (?P\d+)\.(?P\d+)\.(?P\d+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 363c5d1..270220e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -17,6 +17,20 @@ and this project adheres to [semantic versioning](https://semver.org/spec/v2.0.0 ### Removed +## [0.15.0] - 2025-08-21 + +### Added +- Added functions `check_year_range` & `assert_same_distinct_value` + to `helpers/pyspark.py`. + +### Changed + +### Deprecated + +### Fixed + +### Removed + ## [0.14.1] - 2025-08-13 ### Added @@ -900,6 +914,8 @@ and this project adheres to [semantic versioning](https://semver.org/spec/v2.0.0 > due to bugs in the GitHub Action `deploy_pypi.yaml`, which deploys to PyPI > and GitHub Releases. +- rdsa-utils v0.15.0: [GitHub Release](https://github.com/ONSdigital/rdsa-utils/releases/tag/v0.15.0) | + [PyPI](https://pypi.org/project/rdsa-utils/0.15.0/) - rdsa-utils v0.14.1: [GitHub Release](https://github.com/ONSdigital/rdsa-utils/releases/tag/v0.14.1) | [PyPI](https://pypi.org/project/rdsa-utils/0.14.1/) - rdsa-utils v0.14.0: [GitHub Release](https://github.com/ONSdigital/rdsa-utils/releases/tag/v0.14.0) | diff --git a/rdsa_utils/__init__.py b/rdsa_utils/__init__.py index f075dd3..9da2f8f 100644 --- a/rdsa_utils/__init__.py +++ b/rdsa_utils/__init__.py @@ -1 +1 @@ -__version__ = "0.14.1" +__version__ = "0.15.0" diff --git a/rdsa_utils/helpers/pyspark.py b/rdsa_utils/helpers/pyspark.py index 96d1a9e..73bddbc 100644 --- a/rdsa_utils/helpers/pyspark.py +++ b/rdsa_utils/helpers/pyspark.py @@ -13,6 +13,7 @@ Mapping, Optional, Sequence, + Set, Union, ) @@ -1819,3 +1820,208 @@ def has_no_nulls(df: SparkDF, column_name: str) -> bool: else: logger.info(f"Column '{column_name}' contains no null values.") return True + + +def check_year_range( + df: SparkDF, + start_year: int, + end_year: int, + year_col: str, +) -> None: + """Check if a DataFrame contains all years within a given range. + + This function verifies that every year from `start_year` to `end_year` + (inclusive) exists in the specified year column of the DataFrame. If any + years are missing, it raises a ValueError. + + Parameters + ---------- + df + The input DataFrame to be checked. + start_year + The starting year of the range (inclusive). + end_year + The ending year of the range (inclusive). + year_col + The name of the column containing integer year values. + + Raises + ------ + ValueError + - If `start_year` is greater than `end_year`. + - If the specified `year_col` is not found in the DataFrame. + - If one or more years within the specified range are missing from + the DataFrame's `year_col`. + + Returns + ------- + None + The function completes successfully if all years within the specified + range are present in the DataFrame. + + Examples + -------- + >>> from pyspark.sql import SparkSession + >>> spark = SparkSession.builder.appName("YearCheckExample").getOrCreate() + + >>> # Example DataFrame + >>> data = [(2018, "A"), (2019, "B"), (2020, "C"), (2021, "D")] + >>> df = spark.createDataFrame(data, ["year", "data"]) + + >>> # This will pass successfully and print log messages + >>> check_year_range(df, start_year=2019, end_year=2021, year_col="year") + + >>> # This will raise a ValueError because 2022 is missing + >>> try: + ... check_year_range(df, start_year=2018, end_year=2022) + ... except ValueError as e: + ... print(f"ERROR: {e}") + """ + logger.info( + f"Starting year range check for column '{year_col}' " + f"from {start_year} to {end_year}.", + ) + + # --- 1. Input Validation --- + if start_year > end_year: + error_msg = ( + f"start_year ({start_year}) cannot be greater than end_year ({end_year})." + ) + logger.error(error_msg) + raise ValueError(error_msg) + + if year_col not in df.columns: + error_msg = ( + f"Column '{year_col}' not found in the DataFrame. " + f"Available columns: {df.columns}", + ) + logger.error(error_msg) + raise ValueError(error_msg) + + # --- 2. Identify Required and Actual Years --- + required_years = set(range(start_year, end_year + 1)) + logger.info(f"Generated a requirement set of {len(required_years)} years.") + + actual_years_rows = df.select(year_col).distinct().collect() + actual_years = {row[year_col] for row in actual_years_rows} + logger.info(f"Found {len(actual_years)} distinct years in the DataFrame.") + + # --- 3. Compare and Raise Error if Necessary --- + if not required_years.issubset(actual_years): + missing_years: List[int] = sorted(required_years - actual_years) + error_msg = ( + f"DataFrame is missing the following required year(s): {missing_years}" + ) + logger.error(error_msg) + raise ValueError(error_msg) + + logger.info( + f"Validation successful: All years from {start_year} " + f"to {end_year} are present.", + ) + + +def assert_same_distinct_values(df1: SparkDF, df2: SparkDF, col_name: str) -> None: + """Assert that two DataFrames have an identical set of distinct values. + + This function extracts the unique values from the specified column in each + DataFrame and asserts that the two resulting sets are identical. + + Parameters + ---------- + df1 + The first DataFrame for comparison. + df2 + The second DataFrame for comparison. + col_name + The name of the column whose distinct values will be compared. + + Returns + ------- + None + The function completes successfully if the sets are identical. + + Raises + ------ + ValueError + - If `col_name` is not found in either `df1` or `df2`. + - If the sets of distinct values in the specified column are not identical. + + Examples + -------- + >>> from pyspark.sql import SparkSession, types as T + >>> spark = SparkSession.builder.appName("DistinctExample").getOrCreate() + + >>> # --- Create Sample DataFrames --- + >>> schema = T.StructType([T.StructField("category", T.StringType())]) + >>> df_a = spark.createDataFrame([("A",), ("B",), ("A",)], schema) + >>> df_b = spark.createDataFrame([("B",), ("C",)], schema) + >>> df_c = spark.createDataFrame([("B",), ("A",)], schema) + + >>> # --- 1. Success Case: Identical Sets --- + >>> # Assertion passes silently because the distinct values {'A', 'B'} are the same. + >>> assert_same_distinct_values(df_a, df_c, col_name="category") + + >>> # --- 2. Failure Case: Different Sets --- + >>> # This will raise a ValueError with a descriptive message. + >>> try: + ... assert_same_distinct_values(df_a, df_b, col_name="category") + ... except ValueError as e: + ... print(e) + Column 'category' has different distinct values across DataFrames. + Values only in first DataFrame: {'A'} + Values only in second DataFrame: {'C'} + + >>> # --- 3. Failure Case: Column Not Found --- + >>> # This will raise a ValueError because the column does not exist. + >>> try: + ... assert_same_distinct_values(df_a, df_b, col_name="product_id") + ... except ValueError as e: + ... print(e) + Column 'product_id' not found in the first DataFrame. Available: ['category'] + """ + logger.info(f"Asserting same distinct values in column '{col_name}'.") + + # --- 1. Input Validation --- + if col_name not in df1.columns: + error_msg = ( + f"Column '{col_name}' not found in the first DataFrame. " + f"Available: {df1.columns}" + ) + logger.error(error_msg) + raise ValueError(error_msg) + if col_name not in df2.columns: + error_msg = ( + f"Column '{col_name}' not found in the second DataFrame. " + f"Available: {df2.columns}" + ) + logger.error(error_msg) + raise ValueError(error_msg) + + # --- 2. Extract Distinct Values --- + values1: Set[Any] = { + row[col_name] for row in df1.select(col_name).distinct().collect() + } + logger.info(f"Found {len(values1)} distinct values in the first DataFrame.") + + values2: Set[Any] = { + row[col_name] for row in df2.select(col_name).distinct().collect() + } + logger.info(f"Found {len(values2)} distinct values in the second DataFrame.") + + # --- 3. Assert and Raise on Failure --- + if values1 != values2: + values_only_in_df1 = values1 - values2 + values_only_in_df2 = values2 - values1 + + error_msg = ( + f"Column '{col_name}' has different distinct values across DataFrames.\n" + "Values only in first DataFrame: " + f"{values_only_in_df1 if values_only_in_df1 else '{}'}\n" + "Values only in second DataFrame: " + f"{values_only_in_df2 if values_only_in_df2 else '{}'}" + ) + logger.error(error_msg) + raise ValueError(error_msg) + + logger.info(f"Assertion successful: Sets are identical for column '{col_name}'.") diff --git a/tests/helpers/test_pyspark.py b/tests/helpers/test_pyspark.py index 41e9f3c..8ad0f65 100644 --- a/tests/helpers/test_pyspark.py +++ b/tests/helpers/test_pyspark.py @@ -1885,3 +1885,210 @@ def test_empty_dataframe(self, create_spark_df: Callable) -> None: ], ) assert has_no_nulls(df, "id") is True + + +class TestCheckYearRange: + """Tests for check_year_range function.""" + + def test_success_case_logs( + self, + create_spark_df: Callable, + caplog: pytest.LogCaptureFixture, + ) -> None: + """Tests the function passes and emits the correct logs on success.""" + caplog.set_level(logging.INFO) + schema = T.StructType([T.StructField("year", T.IntegerType())]) + input_df = create_spark_df([schema, (2018,), (2019,), (2020,), (2021,)]) + + check_year_range(input_df, start_year=2019, end_year=2020, year_col="year") + + assert ( + "Starting year range check for column 'year' from 2019 to 2020" + in caplog.text + ) + assert ( + "Validation successful: All years from 2019 to 2020 are present" + in caplog.text + ) + + def test_success_with_exact_match_and_duplicates( + self, + create_spark_df: Callable, + ) -> None: + """Tests the function passes when the data contains exactly the required years, with duplicates.""" + schema = T.StructType([T.StructField("year", T.IntegerType())]) + input_df = create_spark_df([schema, (2020,), (2021,), (2021,), (2022,)]) + + check_year_range(input_df, start_year=2020, end_year=2022, year_col="year") + + def test_success_with_custom_column_name(self, create_spark_df: Callable) -> None: + """Tests the function passes when using a non-default column name for the year.""" + schema = T.StructType([T.StructField("fiscal_year", T.IntegerType())]) + input_df = create_spark_df([schema, (2019,), (2020,)]) + + check_year_range( + input_df, + start_year=2019, + end_year=2020, + year_col="fiscal_year", + ) + + def test_raises_error_for_missing_year(self, create_spark_df: Callable) -> None: + """Tests the function raises a ValueError when a year in the middle of the range is missing.""" + schema = T.StructType([T.StructField("year", T.IntegerType())]) + input_df = create_spark_df([schema, (2019,), (2021,)]) + + with pytest.raises(ValueError) as excinfo: + check_year_range(input_df, start_year=2019, end_year=2021, year_col="year") + + assert "missing the following required year(s): [2020]" in str(excinfo.value) + + def test_raises_error_for_multiple_missing_years( + self, + create_spark_df: Callable, + ) -> None: + """Tests the function raises a ValueError listing all missing years.""" + schema = T.StructType([T.StructField("year", T.IntegerType())]) + input_df = create_spark_df([schema, (2020,)]) + + with pytest.raises(ValueError) as excinfo: + check_year_range(input_df, start_year=2019, end_year=2022, year_col="year") + + assert "missing the following required year(s): [2019, 2021, 2022]" in str( + excinfo.value, + ) + + def test_raises_error_for_invalid_year_column( + self, + create_spark_df: Callable, + ) -> None: + """Tests the function raises a ValueError if the specified year column does not exist.""" + schema = T.StructType([T.StructField("year", T.IntegerType())]) + input_df = create_spark_df([schema, (2020,)]) + + with pytest.raises(ValueError) as excinfo: + check_year_range( + input_df, + start_year=2019, + end_year=2020, + year_col="non_existent_col", + ) + + assert "Column 'non_existent_col' not found" in str(excinfo.value) + + def test_raises_error_for_invalid_range(self, create_spark_df: Callable) -> None: + """Tests the function raises a ValueError if start_year is greater than end_year.""" + schema = T.StructType([T.StructField("year", T.IntegerType())]) + input_df = create_spark_df([schema, (2020,)]) + + with pytest.raises(ValueError) as excinfo: + check_year_range(input_df, start_year=2022, end_year=2020, year_col="year") + + assert "start_year (2022) cannot be greater than end_year (2020)" in str( + excinfo.value, + ) + + +class TestAssertSameDistinctValues: + """Tests for assert_same_distinct_values function.""" + + def test_success_identical_sets_logs( + self, + create_spark_df: Callable, + caplog: pytest.LogCaptureFixture, + ) -> None: + """Tests the function passes and emits all expected logs on success.""" + caplog.set_level(logging.INFO) + schema = T.StructType([T.StructField("category", T.StringType())]) + df1 = create_spark_df([schema, ("A",), ("B",), ("A",)]) + df2 = create_spark_df([schema, ("B",), ("A",)]) + + assert_same_distinct_values(df1, df2, col_name="category") + + assert "Asserting same distinct values in column 'category'" in caplog.text + assert "Found 2 distinct values in the first DataFrame" in caplog.text + assert "Found 2 distinct values in the second DataFrame" in caplog.text + assert ( + "Assertion successful: Sets are identical for column 'category'" + in caplog.text + ) + + def test_success_with_null_values(self, create_spark_df: Callable) -> None: + """Tests the function passes when both sets contain null values.""" + schema = T.StructType([T.StructField("id", T.StringType(), True)]) + df1 = create_spark_df([schema, ("1",), (None,)]) + df2 = create_spark_df([schema, (None,), ("1",)]) + + assert_same_distinct_values(df1, df2, col_name="id") + + def test_success_with_empty_dataframes(self, create_spark_df: Callable) -> None: + """Tests the function passes when both DataFrames are empty.""" + schema = T.StructType([T.StructField("id", T.StringType())]) + df1 = create_spark_df([schema]) + df2 = create_spark_df([schema]) + + assert_same_distinct_values(df1, df2, col_name="id") + + def test_raises_error_on_different_sets_logs( + self, + create_spark_df: Callable, + caplog: pytest.LogCaptureFixture, + ) -> None: + """Tests the function emits logs correctly before raising a ValueError.""" + caplog.set_level(logging.INFO) + schema = T.StructType([T.StructField("category", T.StringType())]) + df1 = create_spark_df([schema, ("A",), ("B",)]) + df2 = create_spark_df([schema, ("A",), ("C",)]) + + with pytest.raises(ValueError) as excinfo: + assert_same_distinct_values(df1, df2, col_name="category") + + assert "Values only in first DataFrame: {'B'}" in str(excinfo.value) + assert "Values only in second DataFrame: {'C'}" in str(excinfo.value) + + assert "Asserting same distinct values in column 'category'" in caplog.text + assert "Found 2 distinct values in the first DataFrame" in caplog.text + assert "Found 2 distinct values in the second DataFrame" in caplog.text + assert "Assertion successful" not in caplog.text + + def test_raises_error_when_one_is_subset(self, create_spark_df: Callable) -> None: + """Tests the function raises a ValueError when one set is a subset of the other.""" + schema = T.StructType([T.StructField("category", T.StringType())]) + df1 = create_spark_df([schema, ("A",), ("B",), ("C",)]) + df2 = create_spark_df([schema, ("A",), ("B",)]) + + with pytest.raises(ValueError) as excinfo: + assert_same_distinct_values(df1, df2, col_name="category") + + assert "Values only in first DataFrame: {'C'}" in str(excinfo.value) + assert "Values only in second DataFrame: {}" in str(excinfo.value) + + def test_raises_error_if_col_not_in_first_df( + self, + create_spark_df: Callable, + ) -> None: + """Tests the function raises a ValueError if the column is missing from the first DataFrame.""" + schema1 = T.StructType([T.StructField("col_a", T.StringType())]) + schema2 = T.StructType([T.StructField("col_b", T.StringType())]) + df1 = create_spark_df([schema1, ("val",)]) + df2 = create_spark_df([schema2, ("val",)]) + + with pytest.raises(ValueError) as excinfo: + assert_same_distinct_values(df1, df2, col_name="col_b") + + assert "Column 'col_b' not found in the first DataFrame" in str(excinfo.value) + + def test_raises_error_if_col_not_in_second_df( + self, + create_spark_df: Callable, + ) -> None: + """Tests the function raises a ValueError if the column is missing from the second DataFrame.""" + schema1 = T.StructType([T.StructField("col_a", T.StringType())]) + schema2 = T.StructType([T.StructField("col_b", T.StringType())]) + df1 = create_spark_df([schema1, ("val",)]) + df2 = create_spark_df([schema2, ("val",)]) + + with pytest.raises(ValueError) as excinfo: + assert_same_distinct_values(df1, df2, col_name="col_a") + + assert "Column 'col_a' not found in the second DataFrame" in str(excinfo.value)