diff --git a/pydeequ/checks.py b/pydeequ/checks.py index 87280c3..4a70ca4 100644 --- a/pydeequ/checks.py +++ b/pydeequ/checks.py @@ -727,81 +727,76 @@ def isPositive(self, column, assertion=None, hint=None): self._Check = self._Check.isPositive(column, assertion_func, hint) return self + def _column_comparison(self, columnA, columnB, operator, description, assertion, hint): + """Build a ``columnA columnB`` predicate and apply it via ``satisfies``. + + ``columnB`` may be a column name or a SQL literal/expression. + """ + assertion_func = ( + ScalaFunction1(self._spark_session.sparkContext._gateway, assertion) + if assertion + else getattr(self._Check, "satisfies$default$3")() + ) + hint = self._jvm.scala.Option.apply(hint) + column_condition = f"`{columnA}` {operator} {columnB}" + constraint_name = f"{columnA} is {description} {columnB}" + self._Check = self._Check.satisfies( + column_condition, + constraint_name, + assertion_func, + hint, + self._jvm.scala.collection.Seq.empty(), + self._jvm.scala.Option.apply(None), + ) + return self + def isLessThan(self, columnA, columnB, assertion=None, hint=None): """ Asserts that, in each row, the value of columnA is less than the value of columnB :param str columnA: Column in DataFrame to run the assertion on. - :param str columnB: Column in DataFrame to run the assertion on. + :param str columnB: Column in DataFrame to compare against, or a SQL literal/expression. :param lambda assertion: A function that accepts an int or float parameter. :param str hint: A hint that states why a constraint could have failed. :return: isLessThan self : A Check object that checks the assertion on the columns. """ - assertion_func = ( - ScalaFunction1(self._spark_session.sparkContext._gateway, assertion) - if assertion - else getattr(self._Check, "isLessThan$default$3")() - ) - hint = self._jvm.scala.Option.apply(hint) - self._Check = self._Check.isLessThan(columnA, columnB, assertion_func, hint) - return self + return self._column_comparison(columnA, columnB, "<", "less than", assertion, hint) def isLessThanOrEqualTo(self, columnA, columnB, assertion=None, hint=None): """ Asserts that, in each row, the value of columnA is less than or equal to the value of columnB. :param str columnA: Column in DataFrame to run the assertion on. - :param str columnB: Column in DataFrame to run the assertion on. + :param str columnB: Column in DataFrame to compare against, or a SQL literal/expression. :param lambda assertion: A function that accepts an int or float parameter. :param str hint: A hint that states why a constraint could have failed. :return: isLessThanOrEqualTo self (isLessThanOrEqualTo): A Check object that checks the assertion on the columns. """ - assertion_func = ( - ScalaFunction1(self._spark_session.sparkContext._gateway, assertion) - if assertion - else getattr(self._Check, "isLessThanOrEqualTo$default$3")() - ) - hint = self._jvm.scala.Option.apply(hint) - self._Check = self._Check.isLessThanOrEqualTo(columnA, columnB, assertion_func, hint) - return self + return self._column_comparison(columnA, columnB, "<=", "less than or equal to", assertion, hint) def isGreaterThan(self, columnA, columnB, assertion=None, hint=None): """ Asserts that, in each row, the value of columnA is greater than the value of columnB :param str columnA: Column in DataFrame to run the assertion on. - :param str columnB: Column in DataFrame to run the assertion on. + :param str columnB: Column in DataFrame to compare against, or a SQL literal/expression. :param lambda assertion: A function that accepts an int or float parameter. :param str hint: A hint that states why a constraint could have failed. :return: isGreaterThan self: A Check object that runs the assertion on the columns. """ - assertion_func = ( - ScalaFunction1(self._spark_session.sparkContext._gateway, assertion) - if assertion - else getattr(self._Check, "isGreaterThan$default$3")() - ) - hint = self._jvm.scala.Option.apply(hint) - self._Check = self._Check.isGreaterThan(columnA, columnB, assertion_func, hint) - return self + return self._column_comparison(columnA, columnB, ">", "greater than", assertion, hint) def isGreaterThanOrEqualTo(self, columnA, columnB, assertion=None, hint=None): """ Asserts that, in each row, the value of columnA is greather than or equal to the value of columnB :param str columnA: Column in DataFrame to run the assertion on. - :param str columnB: Column in DataFrame to run the assertion on. + :param str columnB: Column in DataFrame to compare against, or a SQL literal/expression. :param lambda assertion: A function that accepts an int or float parameter. :param str hint: A hint that states why a constraint could have failed. :return: isGreaterThanOrEqualTo self: A Check object that runs the assertion on the columns. """ - assertion_func = ( - ScalaFunction1(self._spark_session.sparkContext._gateway, assertion) - if assertion - else getattr(self._Check, "isGreaterThanOrEqualTo$default$3")() - ) - hint = self._jvm.scala.Option.apply(hint) - self._Check = self._Check.isGreaterThanOrEqualTo(columnA, columnB, assertion_func, hint) - return self + return self._column_comparison(columnA, columnB, ">=", "greater than or equal to", assertion, hint) def isContainedIn(self, column, allowed_values, assertion=None, hint=None): """ diff --git a/tests/test_checks.py b/tests/test_checks.py index 878257b..3ed000d 100644 --- a/tests/test_checks.py +++ b/tests/test_checks.py @@ -1020,6 +1020,43 @@ def test_fail_isGreaterThanOrEqualTo(self): self.isGreaterThanOrEqualTo("h", "f", lambda x: x == 1), [Row(constraint_status="Failure")] ) + def test_comparator_against_literal(self): + # column "b" holds 1, 2, 3 + self.assertEqual( + self.isGreaterThanOrEqualTo("b", "1", hint="Cluster should have at least one element"), + [Row(constraint_status="Success")], + ) + self.assertEqual( + self.isLessThanOrEqualTo("b", "10", hint="b never exceeds 10"), + [Row(constraint_status="Success")], + ) + self.assertEqual( + self.isGreaterThan("b", "0"), + [Row(constraint_status="Success")], + ) + self.assertEqual( + self.isLessThan("b", "100"), + [Row(constraint_status="Success")], + ) + + def test_fail_comparator_against_literal(self): + # Column "b" holds values 1, 2, 3 -> not all are >= 3. + self.assertEqual( + self.isGreaterThanOrEqualTo("b", "3"), [Row(constraint_status="Failure")] + ) + + def test_comparator_column_name_with_space(self): + df = self.df.withColumnRenamed("b", "my col") + check = Check(self.spark, CheckLevel.Warning, "test spaced column").isGreaterThanOrEqualTo( + "my col", "1", hint="values in 'my col' are at least 1" + ) + result = VerificationSuite(self.spark).onData(df).addCheck(check).run() + result_df = VerificationResult.checkResultsAsDataFrame(self.spark, result) + self.assertEqual( + result_df.select("constraint_status").collect(), + [Row(constraint_status="Success")], + ) + def test_where(self): self.assertEqual( self.where(lambda x: x == 2.0, "boolean='true'", "column 'boolean' has two values true"),