From 8a3c0e632d820dc2b80a82b3cea0fe980f8bef10 Mon Sep 17 00:00:00 2001 From: Patrick Ribbsaeter <233187023+patrickswedish@users.noreply.github.com> Date: Mon, 3 Aug 2026 01:30:09 +0200 Subject: [PATCH] fix: support Spark Connect expressions --- datacompy/comparator/array.py | 14 ++-- datacompy/comparator/boolean.py | 34 +++++--- datacompy/comparator/numeric.py | 48 +++++++---- datacompy/comparator/string.py | 37 ++++---- datacompy/comparator/utility.py | 32 +++++++ datacompy/spark.py | 112 +++++++++++++++++-------- tests/comparator/test_utility_spark.py | 18 +++- tests/test_spark.py | 12 +++ 8 files changed, 218 insertions(+), 89 deletions(-) diff --git a/datacompy/comparator/array.py b/datacompy/comparator/array.py index b94c937..ed3dab5 100644 --- a/datacompy/comparator/array.py +++ b/datacompy/comparator/array.py @@ -29,12 +29,13 @@ POLARS_ARRAY_TYPE = ["List", "Array"] try: import pyspark as ps - import pyspark.sql.functions as psf - from datacompy.comparator.utility import get_spark_column_dtypes + from datacompy.comparator.utility import ( + get_spark_column_dtypes, + get_spark_functions, + ) except ImportError: ps = None - psf = None try: import snowflake.snowpark as sp @@ -154,8 +155,11 @@ class SparkArrayLikeComparator(BaseComparator): """ base_dtype, compare_dtype = get_spark_column_dtypes(dataframe, col1, col2) if base_dtype.startswith("array") and compare_dtype.startswith("array"): - when_clause = psf.col(col1).eqNullSafe(psf.col(col2)) - return psf.when(when_clause, psf.lit(True)).otherwise(psf.lit(False)) + functions = get_spark_functions(dataframe) + when_clause = functions.col(col1).eqNullSafe(functions.col(col2)) + return functions.when(when_clause, functions.lit(True)).otherwise( + functions.lit(False) + ) else: return None diff --git a/datacompy/comparator/boolean.py b/datacompy/comparator/boolean.py index f6fe8cb..37f9347 100644 --- a/datacompy/comparator/boolean.py +++ b/datacompy/comparator/boolean.py @@ -32,13 +32,14 @@ STRING_POLARS_TYPES = (pl.String, pl.Categorical, pl.Enum) # Optional Spark dependencies try: import pyspark as ps - import pyspark.sql.functions as psf from datacompy.comparator.numeric import NUMERIC_PYSPARK_TYPES - from datacompy.comparator.utility import get_spark_column_dtypes + from datacompy.comparator.utility import ( + get_spark_column_dtypes, + get_spark_functions, + ) except ImportError: ps = None - psf = None NUMERIC_PYSPARK_TYPES = None # Optional Snowflake dependencies @@ -171,7 +172,9 @@ class SparkBooleanComparator(BaseComparator): """Comparator for Boolean columns in PySpark.""" @staticmethod - def _boolean_equals_numeric(boolean_col: str, numeric_col: str) -> "ps.sql.Column": + def _boolean_equals_numeric( + boolean_col: str, numeric_col: str, functions: Any + ) -> "ps.sql.Column": """Compare a Boolean column against a numeric column's 1/0 equivalents. The numeric side is compared against integer literals rather than both @@ -193,12 +196,16 @@ class SparkBooleanComparator(BaseComparator): A Column of booleans. Two nulls are treated as equal; a null against a value is not. """ - boolean = psf.col(boolean_col) - numeric = psf.col(numeric_col) + boolean = functions.col(boolean_col) + numeric = functions.col(numeric_col) both_null = boolean.isNull() & numeric.isNull() values_equal = ( - boolean.eqNullSafe(psf.lit(True)) & numeric.eqNullSafe(psf.lit(1)) - ) | (boolean.eqNullSafe(psf.lit(False)) & numeric.eqNullSafe(psf.lit(0))) + boolean.eqNullSafe(functions.lit(True)) + & numeric.eqNullSafe(functions.lit(1)) + ) | ( + boolean.eqNullSafe(functions.lit(False)) + & numeric.eqNullSafe(functions.lit(0)) + ) return both_null | values_equal def compare( @@ -261,17 +268,20 @@ class SparkBooleanComparator(BaseComparator): compare_numeric_type = any( compare_dtype.startswith(t) for t in NUMERIC_PYSPARK_TYPES ) + functions = get_spark_functions(dataframe) if base_boolean_type and compare_boolean_type: - when_clause = psf.col(col1).eqNullSafe(psf.col(col2)) + when_clause = functions.col(col1).eqNullSafe(functions.col(col2)) elif base_boolean_type and compare_numeric_type: - when_clause = self._boolean_equals_numeric(col1, col2) + when_clause = self._boolean_equals_numeric(col1, col2, functions) elif compare_boolean_type and base_numeric_type: - when_clause = self._boolean_equals_numeric(col2, col1) + when_clause = self._boolean_equals_numeric(col2, col1, functions) else: return None - return psf.when(when_clause, psf.lit(True)).otherwise(psf.lit(False)) + return functions.when(when_clause, functions.lit(True)).otherwise( + functions.lit(False) + ) class SnowflakeBooleanComparator(BaseComparator): diff --git a/datacompy/comparator/numeric.py b/datacompy/comparator/numeric.py index c32f08d..a033ef8 100644 --- a/datacompy/comparator/numeric.py +++ b/datacompy/comparator/numeric.py @@ -48,9 +48,11 @@ def decimal_comparator(): # Optional Spark dependencies try: import pyspark as ps - import pyspark.sql.functions as psf - from datacompy.comparator.utility import get_spark_column_dtypes + from datacompy.comparator.utility import ( + get_spark_column_dtypes, + get_spark_functions, + ) NUMERIC_PYSPARK_TYPES = [ "tinyint", @@ -64,7 +66,6 @@ try: SPARK_INTEGER_TYPES = {"tinyint", "smallint", "int", "bigint"} except ImportError: ps = None - psf = None NUMERIC_PYSPARK_TYPES = None SPARK_INTEGER_TYPES = None @@ -290,40 +291,53 @@ class SparkNumericComparator(BaseComparator): compare_dtype.startswith(t) for t in NUMERIC_PYSPARK_TYPES ) if (base_numeric_type) and (compare_numeric_type): + functions = get_spark_functions(dataframe) # Cast integer types to double to avoid ANSI-mode overflow on subtraction # and to prevent isnan() being called on non-float types. col1_expr = ( - psf.col(col1).cast("double") + functions.col(col1).cast("double") if base_dtype in SPARK_INTEGER_TYPES - else psf.col(col1) + else functions.col(col1) ) col2_expr = ( - psf.col(col2).cast("double") + functions.col(col2).cast("double") if compare_dtype in SPARK_INTEGER_TYPES - else psf.col(col2) + else functions.col(col2) ) col1_can_be_nan = base_dtype in {"float", "double"} col2_can_be_nan = compare_dtype in {"float", "double"} try: - nan_col1 = psf.isnan(col1_expr) if col1_can_be_nan else psf.lit(False) - nan_col2 = psf.isnan(col2_expr) if col2_can_be_nan else psf.lit(False) + nan_col1 = ( + functions.isnan(col1_expr) + if col1_can_be_nan + else functions.lit(False) + ) + nan_col2 = ( + functions.isnan(col2_expr) + if col2_can_be_nan + else functions.lit(False) + ) return ( - psf.when(nan_col1 & nan_col2, psf.lit(True)) # NaN == NaN + functions.when( + nan_col1 & nan_col2, functions.lit(True) + ) # NaN == NaN .when( - nan_col1 | nan_col2, psf.lit(False) + nan_col1 | nan_col2, functions.lit(False) ) # NaN != real, real != NaN .when( - psf.col(col1).eqNullSafe(psf.col(col2)), psf.lit(True) + functions.col(col1).eqNullSafe(functions.col(col2)), + functions.lit(True), ) # NULL-safe equality .when( - psf.abs(col1_expr - col2_expr) - <= psf.lit(atol) + (psf.lit(rtol) * psf.abs(col2_expr)), - psf.lit(True), # within tolerance + functions.abs(col1_expr - col2_expr) + <= functions.lit(atol) + + (functions.lit(rtol) * functions.abs(col2_expr)), + functions.lit(True), # within tolerance ) - .otherwise(psf.lit(False)) + .otherwise(functions.lit(False)) ) except Exception: - return psf.lit(False) + return functions.lit(False) else: return None diff --git a/datacompy/comparator/string.py b/datacompy/comparator/string.py index 23ab2ae..58bac93 100644 --- a/datacompy/comparator/string.py +++ b/datacompy/comparator/string.py @@ -27,15 +27,16 @@ LOG = logging.getLogger(__name__) # Initialize optional dependencies ps = None -psf = None sp = None spf = None try: import pyspark as ps - import pyspark.sql.functions as psf - from datacompy.comparator.utility import get_spark_column_dtypes + from datacompy.comparator.utility import ( + get_spark_column_dtypes, + get_spark_functions, + ) PYSPARK_STRING_TYPE = {"string", "char", "varchar"} PYSPARK_DATE_TYPE = {"date", "timestamp"} @@ -307,23 +308,24 @@ class SparkStringComparator(BaseComparator): ) or ((base_date_type) and (compare_date_type)) # date/date compare. ): + functions = get_spark_functions(dataframe) try: if base_date_type and compare_date_type: # Both are date/timestamp: compare directly, no string conversion needed. - col1_expr = psf.col(col1) - col2_expr = psf.col(col2) + col1_expr = functions.col(col1) + col2_expr = functions.col(col2) elif base_date_type and compare_string_type: # Cast the string to match the date column's type using TRY_CAST so # malformed strings return NULL instead of throwing in ANSI mode. - col1_expr = psf.col(col1) - col2_expr = psf.expr(f"TRY_CAST(`{col2}` AS {base_dtype})") + col1_expr = functions.col(col1) + col2_expr = functions.expr(f"TRY_CAST(`{col2}` AS {base_dtype})") elif base_string_type and compare_date_type: - col1_expr = psf.expr(f"TRY_CAST(`{col1}` AS {compare_dtype})") - col2_expr = psf.col(col2) + col1_expr = functions.expr(f"TRY_CAST(`{col1}` AS {compare_dtype})") + col2_expr = functions.col(col2) else: # Both strings: compare as-is with optional normalisation. - col1_expr = psf.col(col1) - col2_expr = psf.col(col2) + col1_expr = functions.col(col1) + col2_expr = functions.col(col2) col1_expr = spark_normalize_string_column( col1_expr, ignore_space, ignore_case ) @@ -331,11 +333,11 @@ class SparkStringComparator(BaseComparator): col2_expr, ignore_space, ignore_case ) - return psf.when( - col1_expr.eqNullSafe(col2_expr), psf.lit(True) - ).otherwise(psf.lit(False)) + return functions.when( + col1_expr.eqNullSafe(col2_expr), functions.lit(True) + ).otherwise(functions.lit(False)) except Exception: - return psf.lit(False) + return functions.lit(False) else: return None @@ -504,10 +506,11 @@ def spark_normalize_string_column( pyspark.sql.Column The normalized column """ + functions = get_spark_functions(column) if ignore_spaces: - column = psf.trim(column) + column = functions.trim(column) if ignore_case: - column = psf.upper(column) + column = functions.upper(column) return column diff --git a/datacompy/comparator/utility.py b/datacompy/comparator/utility.py index fd70eb2..e9645b3 100644 --- a/datacompy/comparator/utility.py +++ b/datacompy/comparator/utility.py @@ -15,6 +15,8 @@ """Utility and helper functions for data comparison.""" +from typing import Any + # Optional dependencies initialization ps = None sp = None @@ -30,6 +32,36 @@ except ImportError: pass +def get_spark_functions(spark_object: object) -> Any: + """Return the functions module matching a classic or Connect Spark object.""" + from pyspark.sql.connect.column import Column as ConnectColumn + from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame + + if isinstance(spark_object, ConnectColumn | ConnectDataFrame): + from pyspark.sql.connect import functions as connect_functions + + return connect_functions + + from pyspark.sql import functions as classic_functions + + return classic_functions + + +def get_spark_window(spark_object: object) -> Any: + """Return the Window class matching a classic or Connect Spark object.""" + from pyspark.sql.connect.column import Column as ConnectColumn + from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame + + if isinstance(spark_object, ConnectColumn | ConnectDataFrame): + from pyspark.sql.connect.window import Window as ConnectWindow + + return ConnectWindow + + from pyspark.sql import Window as ClassicWindow + + return ClassicWindow + + def get_spark_column_dtypes( dataframe: "ps.sql.DataFrame", col_1: str, col_2: str ) -> tuple[str, str]: diff --git a/datacompy/spark.py b/datacompy/spark.py index 78c0f0f..b90b482 100644 --- a/datacompy/spark.py +++ b/datacompy/spark.py @@ -27,9 +27,7 @@ from typing import Any, Dict, List, Tuple import pandas as pd import pyspark.sql -import pyspark.sql.functions as F from ordered_set import OrderedSet -from pyspark.sql import Window from pyspark.sql.connect.dataframe import DataFrame from datacompy.base import ( @@ -46,7 +44,11 @@ from datacompy.comparator import ( SparkStringComparator, ) from datacompy.comparator.base import BaseComparator -from datacompy.comparator.utility import get_spark_column_dtypes +from datacompy.comparator.utility import ( + get_spark_column_dtypes, + get_spark_functions, + get_spark_window, +) LOG = logging.getLogger(__name__) @@ -223,6 +225,7 @@ class SparkSQLCompare(BaseCompare): def hide_sensitive_columns(self, sensitive_columns: List[str]) -> None: """Hides sensitive columns of df1 or df2 if applicable in the compare.""" + functions = get_spark_functions(self.df1) # Don't allow hiding columns again before first revealing if self.sensitive_columns: raise ValueError( @@ -249,7 +252,9 @@ class SparkSQLCompare(BaseCompare): continue # Maintains column ordering for the hide select_cols = [ - F.lit("*******").alias(col) if col in cols_to_hide else F.col(col) + functions.lit("*******").alias(col) + if col in cols_to_hide + else functions.col(col) for col in df.columns ] setattr(self, df_name, df.select(select_cols)) @@ -262,7 +267,9 @@ class SparkSQLCompare(BaseCompare): if not cols_to_hide: # skip if empty return select_cols = [ - F.lit("*******").alias(col) if col in cols_to_hide else F.col(col) + functions.lit("*******").alias(col) + if col in cols_to_hide + else functions.col(col) for col in self.intersect_rows.columns ] # avoid orphaned cached data (incl. unmasked sensitive vals) in executor memory @@ -383,14 +390,15 @@ class SparkSQLCompare(BaseCompare): df1 = self.df1 df2 = self.df2 + functions = get_spark_functions(df1) temp_join_columns = deepcopy(self.join_columns) if self._any_dupes: LOG.debug("Duplicate rows found, deduping by order of remaining fields") # setting internal index LOG.info("Adding internal index to dataframes") - df1 = df1.withColumn("__index", F.monotonically_increasing_id()) - df2 = df2.withColumn("__index", F.monotonically_increasing_id()) + df1 = df1.withColumn("__index", functions.monotonically_increasing_id()) + df2 = df2.withColumn("__index", functions.monotonically_increasing_id()) # Create order column for uniqueness of match order_column = temp_column_name(df1, df2) @@ -419,12 +427,12 @@ class SparkSQLCompare(BaseCompare): next(dtype for name, dtype in df1.dtypes if name == column) == "string" ): - df1 = df1.withColumn(column, F.trim(F.col(column))) + df1 = df1.withColumn(column, functions.trim(functions.col(column))) if ( next(dtype for name, dtype in df2.dtypes if name == column) == "string" ): - df2 = df2.withColumn(column, F.trim(F.col(column))) + df2 = df2.withColumn(column, functions.trim(functions.col(column))) df1_non_join_columns = OrderedSet(df1.columns) - OrderedSet(temp_join_columns) df2_non_join_columns = OrderedSet(df2.columns) - OrderedSet(temp_join_columns) @@ -437,8 +445,8 @@ class SparkSQLCompare(BaseCompare): ) # generate merge indicator - df1 = df1.withColumn("_merge_left", F.lit(True)) - df2 = df2.withColumn("_merge_right", F.lit(True)) + df1 = df1.withColumn("_merge_left", functions.lit(True)) + df2 = df2.withColumn("_merge_right", functions.lit(True)) df1 = df1.withColumnsRenamed( {c: f"{c}_{self.df1_name}" for c in temp_join_columns} @@ -465,18 +473,20 @@ class SparkSQLCompare(BaseCompare): + on ) - outer_join = outer_join.withColumn("_merge", F.lit(None)) # initialize col + outer_join = outer_join.withColumn( + "_merge", functions.lit(None) + ) # initialize col # process merge indicator outer_join = outer_join.withColumn( "_merge", - F.when( + functions.when( (outer_join["_merge_left"] == True) # noqa: E712 - & (F.isnull(outer_join["_merge_right"])), + & (functions.isnull(outer_join["_merge_right"])), "left_only", ) .when( - (F.isnull(outer_join["_merge_left"])) + (functions.isnull(outer_join["_merge_left"])) & (outer_join["_merge_right"] == True), # noqa: E712 "right_only", ) @@ -544,6 +554,7 @@ class SparkSQLCompare(BaseCompare): otherwise. """ LOG.debug("Comparing intersection") + functions = get_spark_functions(self.intersect_rows) max_diff: float null_diff: int exprs = {} @@ -572,7 +583,9 @@ class SparkSQLCompare(BaseCompare): if exprs: agg_exprs = [ - F.sum(F.when(F.col(c) == True, 1).otherwise(0)).alias(f"{c}_count") # noqa: E712 + functions.sum( + functions.when(functions.col(c) == True, 1).otherwise(0) # noqa: E712 + ).alias(f"{c}_count") for c in exprs ] match_counts = self.intersect_rows.agg(*agg_exprs).first() @@ -742,15 +755,18 @@ class SparkSQLCompare(BaseCompare): "pertinent" columns, for rows that don't match on the provided column. """ + functions = get_spark_functions(self.intersect_rows) if not self.only_join_columns() and column not in self.join_columns: row_cnt = self.intersect_rows.count() col_match = self.intersect_rows.select(f"{column}_match") match_cnt = col_match.where( - F.col(f"{column}_match") == True # noqa: E712 + functions.col(f"{column}_match") == True # noqa: E712 ).count() sample_count = min(sample_count, row_cnt - match_cnt) sample = ( - self.intersect_rows.where(F.col(f"{column}_match") == False) # noqa: E712 + self.intersect_rows.where( + functions.col(f"{column}_match") == False # noqa: E712 + ) .drop(f"{column}_match") .limit(sample_count) ) @@ -808,6 +824,7 @@ class SparkSQLCompare(BaseCompare): pyspark.sql.DataFrame All rows of the intersection dataframe, containing any columns, that don't match. """ + functions = get_spark_functions(self.intersect_rows) match_list = [] return_list = [] if self.only_join_columns(): @@ -822,7 +839,9 @@ class SparkSQLCompare(BaseCompare): if match_cols: agg_exprs = [ - F.sum(F.when(F.col(c) == False, 1).otherwise(0)).alias(f"{c}_count") # noqa: E712 + functions.sum( + functions.when(functions.col(c) == False, 1).otherwise(0) # noqa: E712 + ).alias(f"{c}_count") for c in match_cols ] mismatch_counts = self.intersect_rows.agg(*agg_exprs).first() @@ -857,8 +876,8 @@ class SparkSQLCompare(BaseCompare): return to_return mm_rows = self.intersect_rows.withColumn( - "match_array", F.array(match_list) - ).where(F.array_contains("match_array", False)) + "match_array", functions.array(match_list) + ).where(functions.array_contains("match_array", False)) mm_rows = mm_rows.withColumnsRenamed( {f"{c}_{self.df1_name}": c for c in self.join_columns} @@ -986,7 +1005,7 @@ def columns_equal( ) return compare - compare = F.lit(False) + compare = get_spark_functions(dataframe).lit(False) return compare @@ -1042,6 +1061,7 @@ def calculate_max_diff( float max diff """ + functions = get_spark_functions(dataframe) dtype1, dtype2 = get_spark_column_dtypes(dataframe, col_1, col_2) if dtype1.startswith("array") and dtype2.startswith("array"): LOG.warning( @@ -1050,13 +1070,13 @@ def calculate_max_diff( return 0 diff = dataframe.select( ( - F.expr(f"TRY_CAST(`{col_1}` AS DOUBLE)") - - F.expr(f"TRY_CAST(`{col_2}` AS DOUBLE)") + functions.expr(f"TRY_CAST(`{col_1}` AS DOUBLE)") + - functions.expr(f"TRY_CAST(`{col_2}` AS DOUBLE)") ).alias("diff") ) - abs_diff = diff.select(F.abs(F.col("diff")).alias("abs_diff")) + abs_diff = diff.select(functions.abs(functions.col("diff")).alias("abs_diff")) max_diff: float = ( - abs_diff.where(F.isnan(F.col("abs_diff")) == False) # noqa: E712 + abs_diff.where(functions.isnan(functions.col("abs_diff")) == False) # noqa: E712 .agg({"abs_diff": "max"}) .collect()[0][0] ) @@ -1086,23 +1106,32 @@ def calculate_null_diff( int null diff """ + functions = get_spark_functions(dataframe) nulls_df = dataframe.withColumn( "col_1_null", - F.when(F.col(col_1).isNull() == True, F.lit(True)).otherwise( # noqa: E712 - F.lit(False) - ), + functions.when( + functions.col(col_1).isNull(), + functions.lit(True), + ).otherwise(functions.lit(False)), ) nulls_df = nulls_df.withColumn( "col_2_null", - F.when(F.col(col_2).isNull() == True, F.lit(True)).otherwise( # noqa: E712 - F.lit(False) - ), + functions.when( + functions.col(col_2).isNull(), + functions.lit(True), + ).otherwise(functions.lit(False)), ).select(["col_1_null", "col_2_null"]) # (not a and b) or (a and not b) null_diff = nulls_df.where( - ((F.col("col_1_null") == False) & (F.col("col_2_null") == True)) # noqa: E712 - | ((F.col("col_1_null") == True) & (F.col("col_2_null") == False)) # noqa: E712 + ( + (functions.col("col_1_null") == False) # noqa: E712 + & (functions.col("col_2_null") == True) # noqa: E712 + ) + | ( + (functions.col("col_1_null") == True) # noqa: E712 + & (functions.col("col_2_null") == False) # noqa: E712 + ) ).count() if pd.isna(null_diff) or pd.isnull(null_diff) or null_diff is None: @@ -1133,6 +1162,8 @@ def _generate_id_within_group( pyspark.sql.DataFrame Original dataframe with the ID column that's unique in each group """ + functions = get_spark_functions(dataframe) + window = get_spark_window(dataframe) default_value = "DATACOMPY_NULL" null_cols = [f"any(isnull({c}))" for c in join_columns] default_cols = [ @@ -1148,12 +1179,17 @@ def _generate_id_within_group( return ( dataframe.select( - *(F.col(c).cast("string").alias(c) for c in [*join_columns, "__index"]) + *( + functions.col(c).cast("string").alias(c) + for c in [*join_columns, "__index"] + ) ) .fillna(default_value) .withColumn( order_column_name, - F.row_number().over(Window.orderBy("__index").partitionBy(join_columns)) + functions.row_number().over( + window.orderBy("__index").partitionBy(join_columns) + ) - 1, ) .select(["__index", order_column_name]) @@ -1163,7 +1199,9 @@ def _generate_id_within_group( dataframe.select([*join_columns, "__index"]) .withColumn( order_column_name, - F.row_number().over(Window.orderBy("__index").partitionBy(join_columns)) + functions.row_number().over( + window.orderBy("__index").partitionBy(join_columns) + ) - 1, ) .select(["__index", order_column_name]) diff --git a/tests/comparator/test_utility_spark.py b/tests/comparator/test_utility_spark.py index 302f5aa..7680e2d 100644 --- a/tests/comparator/test_utility_spark.py +++ b/tests/comparator/test_utility_spark.py @@ -20,7 +20,11 @@ import pytest pytest.importorskip("pyspark") -from datacompy.comparator.utility import get_spark_column_dtypes +from datacompy.comparator.utility import ( + get_spark_column_dtypes, + get_spark_functions, + get_spark_window, +) from pyspark.sql.types import ( DateType, DecimalType, @@ -85,3 +89,15 @@ def test_get_spark_column_dtypes_case_insensitive(spark_session): dtype1, dtype2 = get_spark_column_dtypes(df, "num", "str") assert dtype1 == "bigint" assert dtype2 == "string" + + +def test_get_spark_helpers_for_connect_objects(): + from pyspark.sql.connect import functions as connect_functions + from pyspark.sql.connect.column import Column as ConnectColumn + from pyspark.sql.connect.window import Window as ConnectWindow + + column = connect_functions.col("value") + + assert isinstance(column, ConnectColumn) + assert get_spark_functions(column) is connect_functions + assert get_spark_window(column) is ConnectWindow diff --git a/tests/test_spark.py b/tests/test_spark.py index eb62036..40f1a73 100644 --- a/tests/test_spark.py +++ b/tests/test_spark.py @@ -73,6 +73,18 @@ NULL|NULL|True""" assert_series_equal(expect_out, actual_out, check_names=False) +def test_columns_equal_builds_connect_expression_without_classic_context(): + from pyspark.sql.connect.column import Column as ConnectColumn + from pyspark.sql.connect.dataframe import DataFrame as ConnectDataFrame + + dataframe = mock.MagicMock(spec=ConnectDataFrame) + dataframe.dtypes = [("left", "string"), ("right", "string")] + + result = columns_equal(dataframe, "left", "right") + + assert isinstance(result, ConnectColumn) + + def test_numeric_columns_equal_rel(spark_session): data = """a|b|expected 1|1|True -- 2.55.0.windows.3