diff --git a/src/databricks/labs/dqx/profiler/profiler.py b/src/databricks/labs/dqx/profiler/profiler.py index 447b5549b..1cd604610 100644 --- a/src/databricks/labs/dqx/profiler/profiler.py +++ b/src/databricks/labs/dqx/profiler/profiler.py @@ -681,19 +681,29 @@ def _round_datetime(value: datetime.datetime, direction: str) -> datetime.dateti """ Rounds a datetime value to midnight based on the specified direction. + - "down" → truncate to midnight (00:00:00). + - "up" → return the next midnight unless value is already midnight. + - Raises ValueError for invalid direction. + :param value: The datetime value to round. :param direction: The rounding direction ("up" or "down"). :return: The rounded datetime value. + :raises ValueError: If direction is not 'up' or 'down'. """ + midnight = value.replace(hour=0, minute=0, second=0, microsecond=0) + if direction == "down": - return value.replace(hour=0, minute=0, second=0, microsecond=0) + return midnight + if direction == "up": + if midnight == value: + return value try: - return value.replace(hour=0, minute=0, second=0, microsecond=0) + datetime.timedelta(days=1) + return midnight + datetime.timedelta(days=1) except OverflowError: logger.warning("Rounding datetime up caused overflow; returning datetime.max instead.") return datetime.datetime.max - return value + raise ValueError(f"Invalid rounding direction: {direction}. Use 'up' or 'down'.") @staticmethod def _round_float(value: float, direction: str) -> float: diff --git a/tests/integration/test_profiler.py b/tests/integration/test_profiler.py index 733ef05c8..c842dccc4 100644 --- a/tests/integration/test_profiler.py +++ b/tests/integration/test_profiler.py @@ -104,6 +104,103 @@ def test_profiler(spark, ws): assert rules == expected_rules +def test_profiler_rounding_midnight_behavior(spark, ws): + inp_schema = T.StructType( + [ + T.StructField("t1", T.IntegerType()), + T.StructField("d1", T.DecimalType(10, 2)), + T.StructField("t2", T.StringType()), + T.StructField( + "s1", + T.StructType( + [ + T.StructField("ns1", T.TimestampType()), + T.StructField( + "s2", + T.StructType([T.StructField("ns2", T.StringType()), T.StructField("ns3", T.DateType())]), + ), + ] + ), + ), + T.StructField("b1", T.ByteType()), + ] + ) + inp_df = spark.createDataFrame( + [ + [ + 1, + Decimal("1.23"), + " test ", + { + "ns1": datetime.fromisoformat("2023-01-08T00:00:00+00:00"), + "s2": {"ns2": "test", "ns3": date.fromisoformat("2023-01-08")}, + }, + 0, + ], + [ + 2, + Decimal("2.41"), + "test2", + { + "ns1": datetime.fromisoformat("2023-01-07T10:00:11+00:00"), + "s2": {"ns2": "test2", "ns3": date.fromisoformat("2023-01-07")}, + }, + 1, + ], + [ + 3, + Decimal("333323.0"), + None, + { + "ns1": datetime.fromisoformat("2023-01-06T10:00:11+00:00"), + "s2": {"ns2": "test", "ns3": date.fromisoformat("2023-01-06")}, + }, + 0, + ], + ], + schema=inp_schema, + ) + + profiler = DQProfiler(ws) + stats, rules = profiler.profile(inp_df, options={"sample_fraction": None}) + + expected_rules = [ + DQProfile(name="is_not_null", column="t1", description=None, parameters=None), + DQProfile( + name="min_max", column="t1", description="Real min/max values were used", parameters={"min": 1, "max": 3} + ), + DQProfile(name='is_not_null', column='d1', description=None, parameters=None), + DQProfile( + name='min_max', + column='d1', + description='Real min/max values were used', + parameters={'max': Decimal('333323.00'), 'min': Decimal('1.23')}, + ), + DQProfile(name='is_not_null_or_empty', column='t2', description=None, parameters={'trim_strings': True}), + DQProfile(name="is_not_null", column="s1.ns1", description=None, parameters=None), + DQProfile( + name="min_max", + column="s1.ns1", + description="Real min/max values were used", + parameters={ + "min": datetime(2023, 1, 6, 0, 0, tzinfo=timezone.utc), + "max": datetime(2023, 1, 8, 0, 0, tzinfo=timezone.utc), + }, + ), + DQProfile(name="is_not_null", column="s1.s2.ns2", description=None, parameters=None), + DQProfile(name="is_not_null", column="s1.s2.ns3", description=None, parameters=None), + DQProfile( + name="min_max", + column="s1.s2.ns3", + description="Real min/max values were used", + parameters={"min": date(2023, 1, 6), "max": date(2023, 1, 8)}, + ), + DQProfile(name="is_not_null", column="b1", description=None, parameters=None), + ] + assert len(stats.keys()) > 0 + assert rules == expected_rules + + def test_profiler_non_default_profile_options(spark, ws): inp_schema = T.StructType( [