fix(sql): preserve timezone casts for extraction ops

Apply timezone-aware cast handling in Trino and localize DuckDB timestamp extraction/time operations so timezone casts behave consistently for columns and scalars.

Add regression tests for timezone cast hour/time extraction and Trino SQL timezone-function generation.

Fixes #11965

Refs #11527

Refs #11211

Refs #11879

Signed-off-by: Mridankan Mandal <iib2024017@iiita.ac.in>

Author: RedZapdos123

ibis/backends/duckdb/tests/test_client.py

@@ -201,6 +201,38 @@ def test_to_other_sql(con, snapshot):
     snapshot.assert_match(sql, "out.sql")
 
 
+def test_timezone_cast_extracts_and_time():
+    con = ibis.duckdb.connect()
+    t = ibis.memtable({"x": ["2023-01-02"]})
+    expr = t.select(
+        ams_hour=t.x.cast("timestamp('Europe/Amsterdam')").hour(),
+        utc_hour=t.x.cast("timestamp('UTC')").hour(),
+        ams_time=t.x.cast("timestamp('Europe/Amsterdam')").time(),
+        utc_time=t.x.cast("timestamp('UTC')").time(),
+    )
+
+    result = con.execute(expr)
+
+    assert result.ams_hour.iat[0] == 1
+    assert result.utc_hour.iat[0] == 0
+    assert str(result.ams_time.iat[0]) == "01:00:00"
+    assert str(result.utc_time.iat[0]) == "00:00:00"
+
+
+def test_to_trino_sql_timezone_cast_uses_timezone_functions():
+    t = ibis.memtable({"x": ["2023-01-02"]})
+    expr = t.select(
+        casted=t.x.cast("timestamp('Europe/Paris')"),
+        hour=t.x.cast("timestamp('Europe/Paris')").hour(),
+        time=t.x.cast("timestamp('Europe/Paris')").time(),
+    )
+
+    sql = ibis.to_sql(expr, dialect="trino")
+
+    assert "AT_TIMEZONE(" in sql
+    assert "WITH_TIMEZONE(" in sql
+
+
 def test_insert_preserves_column_case(con):
     name1 = ibis.util.guid()
     name2 = ibis.util.guid()

ibis/backends/sql/compilers/duckdb.py

@@ -392,12 +392,41 @@ def visit_CountDistinctStar(self, op, *, where, arg):
         )
         return self.agg.count(sge.Distinct(expressions=[row]), where=where)
 
+    def _localize_timestamp_for_extract(self, op, *, arg):
+        if op.arg.dtype.is_timestamp() and (timezone := op.arg.dtype.timezone) is not None:
+            return self.f.timezone(timezone, arg)
+        return arg
+
+    def visit_Time(self, op, *, arg):
+        arg = self._localize_timestamp_for_extract(op, arg=arg)
+        return super().visit_Time(op, arg=arg)
+
+    def visit_ExtractEpochSeconds(self, op, *, arg):
+        if op.arg.dtype.is_timestamp() and op.arg.dtype.timezone is not None:
+            return self.f.epoch(arg)
+        return super().visit_ExtractEpochSeconds(op, arg=arg)
+
+    def visit_ExtractHour(self, op, *, arg):
+        return self.f.extract("hour", self._localize_timestamp_for_extract(op, arg=arg))
+
+    def visit_ExtractMinute(self, op, *, arg):
+        return self.f.extract(
+            "minute", self._localize_timestamp_for_extract(op, arg=arg)
+        )
+
+    def visit_ExtractSecond(self, op, *, arg):
+        return self.f.extract(
+            "second", self._localize_timestamp_for_extract(op, arg=arg)
+        )
+
     def visit_ExtractMillisecond(self, op, *, arg):
+        arg = self._localize_timestamp_for_extract(op, arg=arg)
         return self.f.mod(self.f.extract("ms", arg), 1_000)
 
     # DuckDB extracts subminute microseconds and milliseconds
     # so we have to finesse it a little bit
     def visit_ExtractMicrosecond(self, op, *, arg):
+        arg = self._localize_timestamp_for_extract(op, arg=arg)
         return self.f.mod(self.f.extract("us", arg), 1_000_000)
 
     def visit_TimestampFromUNIX(self, op, *, arg, unit):

ibis/backends/sql/compilers/trino.py

@@ -585,9 +585,21 @@ def visit_Cast(self, op, *, arg, to):
             if from_.is_integer():
                 return self.f.from_unixtime(arg, tz)
             else:
-                return self.f.from_unixtime_nanos(
+                out = self.f.from_unixtime_nanos(
                     self.cast(arg, dt.Decimal(38, 9)) * 1_000_000_000
                 )
+                return self.f.at_timezone(out, tz)
+
+        if to.is_timestamp() and (timezone := to.timezone) is not None:
+            if from_.is_string() or from_.is_date():
+                arg = self.cast(arg, dt.Timestamp(scale=to.scale))
+                from_ = dt.Timestamp(scale=to.scale)
+
+            if from_.is_timestamp():
+                if from_.timezone is None:
+                    arg = self.f.with_timezone(arg, "UTC")
+                return self.f.at_timezone(arg, timezone)
+
         return super().visit_Cast(op, arg=arg, to=to)
 
     def visit_CountDistinctStar(self, op, *, arg, where):