Truncate TIMESTAMPTZ arithmetically in ICU date_trunc for fixed-length parts

benchmark/micro/timestamp/date_trunc_timestamp.benchmark

@@ -0,0 +1,19 @@
+# name: benchmark/micro/timestamp/date_trunc_timestamp.benchmark
+# description: date_trunc('day') on 10 million naive TIMESTAMP values
+# group: [timestamp]
+
+#
+# Baseline for date_trunc_timestamptz*.benchmark. ns/row = timing_seconds * 100.
+
+name TIMESTAMP date_trunc day
+group timestamp
+subgroup date_trunc
+
+load
+CREATE TABLE ts AS SELECT TIMESTAMP '2024-01-01 00:00:00' + to_seconds(i) AS ts FROM range(10000000) t(i);
+
+run
+SELECT count(DISTINCT date_trunc('day', ts)) FROM ts;
+
+result I
+116

benchmark/micro/timestamp/date_trunc_timestamptz.benchmark

@@ -0,0 +1,23 @@
+# name: benchmark/micro/timestamp/date_trunc_timestamptz.benchmark
+# description: date_trunc('day') on 10 million TIMESTAMPTZ values
+# group: [timestamp]
+
+#
+# ns/row = timing_seconds * 1e9 / 10000000 = timing_seconds * 100.
+# Compare with date_trunc_timestamp.benchmark on the same instants.
+
+name TIMESTAMPTZ date_trunc day (UTC)
+group timestamp
+subgroup date_trunc
+
+require icu
+
+load
+SET TimeZone='UTC';
+CREATE TABLE ts AS SELECT (TIMESTAMP '2024-01-01 00:00:00' + to_seconds(i))::TIMESTAMPTZ AS ts FROM range(10000000) t(i);
+
+run
+SELECT count(DISTINCT date_trunc('day', ts)) FROM ts;
+
+result I
+116

benchmark/micro/timestamp/date_trunc_timestamptz_fixed.benchmark

@@ -0,0 +1,23 @@
+# name: benchmark/micro/timestamp/date_trunc_timestamptz_fixed.benchmark
+# description: date_trunc('day') on 10 million TIMESTAMPTZ values in a fixed-offset zone
+# group: [timestamp]
+
+#
+# ns/row = timing_seconds * 100. Etc/GMT-5 is a constant +05:00 offset with no transitions.
+
+name TIMESTAMPTZ date_trunc day (fixed offset)
+group timestamp
+subgroup date_trunc
+
+require icu
+
+load
+SET TimeZone='UTC';
+CREATE TABLE ts AS SELECT (TIMESTAMP '2024-01-01 00:00:00' + to_seconds(i))::TIMESTAMPTZ AS ts FROM range(10000000) t(i);
+SET TimeZone='Etc/GMT-5';
+
+run
+SELECT count(DISTINCT date_trunc('day', ts)) FROM ts;
+
+result I
+116

benchmark/micro/timestamp/date_trunc_timestamptz_new_york.benchmark

@@ -0,0 +1,23 @@
+# name: benchmark/micro/timestamp/date_trunc_timestamptz_new_york.benchmark
+# description: date_trunc('day') on 10 million TIMESTAMPTZ values across America/New_York transitions
+# group: [timestamp]
+
+#
+# ns/row = timing_seconds * 100. The instants span the 2024 spring-forward transition.
+
+name TIMESTAMPTZ date_trunc day (America/New_York)
+group timestamp
+subgroup date_trunc
+
+require icu
+
+load
+SET TimeZone='UTC';
+CREATE TABLE ts AS SELECT (TIMESTAMP '2024-01-01 00:00:00' + to_seconds(i))::TIMESTAMPTZ AS ts FROM range(10000000) t(i);
+SET TimeZone='America/New_York';
+
+run
+SELECT count(DISTINCT date_trunc('day', ts)) FROM ts;
+
+result I
+117

extension/icu/CMakeLists.txt

@@ -25,7 +25,8 @@ set(ICU_EXTENSION_FILES
     icu-table-range.cpp
     icu-strptime.cpp
     icu-timebucket.cpp
-    icu-timezone.cpp)
+    icu-timezone.cpp
+    icu-zone-offsets.cpp)
 
 build_static_extension(icu ${ICU_EXTENSION_FILES})
 link_threads(icu_extension "")

extension/icu/icu-datetrunc.cpp

@@ -1,11 +1,17 @@
 #include "include/icu-datetrunc.hpp"
 #include "include/icu-datefunc.hpp"
+#include "include/icu-zone-offsets.hpp"
 
+#include "duckdb/common/limits.hpp"
+#include "duckdb/common/types/interval.hpp"
 #include "duckdb/common/types/timestamp.hpp"
 #include "duckdb/common/vector_operations/binary_executor.hpp"
+#include "duckdb/execution/expression_executor.hpp"
 #include "duckdb/planner/expression/bound_function_expression.hpp"
 #include "duckdb/main/extension/extension_loader.hpp"
 
+#include <cstring>
+
 namespace duckdb {
 
 struct ICUDateTrunc : public ICUDateFunc {
@@ -123,6 +129,141 @@ struct ICUDateTrunc : public ICUDateFunc {
 		calendar->set(UCAL_ERA, era);
 	}
 
+	static timestamp_t TruncWithICU(icu::Calendar *calendar, part_trunc_t truncator, timestamp_t input) {
+		auto micros = SetTime(calendar, input);
+		truncator(calendar, micros);
+		return GetTimeUnsafe(calendar, micros);
+	}
+
+	//! The length in µs of the parts that always span the same amount of wall time, or 0.
+	static int64_t FixedPartLength(DatePartSpecifier part) {
+		switch (part) {
+		case DatePartSpecifier::DAY:
+		case DatePartSpecifier::DOW:
+		case DatePartSpecifier::ISODOW:
+		case DatePartSpecifier::DOY:
+		case DatePartSpecifier::JULIAN_DAY:
+			return Interval::MICROS_PER_DAY;
+		case DatePartSpecifier::HOUR:
+			return Interval::MICROS_PER_HOUR;
+		case DatePartSpecifier::MINUTE:
+			return Interval::MICROS_PER_MINUTE;
+		case DatePartSpecifier::SECOND:
+		case DatePartSpecifier::EPOCH:
+			return Interval::MICROS_PER_SEC;
+		case DatePartSpecifier::MILLISECONDS:
+			return Interval::MICROS_PER_MSEC;
+		case DatePartSpecifier::MICROSECONDS:
+			return 1;
+		default:
+			return 0;
+		}
+	}
+
+	//! A constant part of fixed length, truncated with arithmetic on µs instead of ICU calls.
+	struct ArithmeticTruncData : public BindData {
+		ArithmeticTruncData(ClientContext &context, DatePartSpecifier part, int64_t unit_p)
+		    : BindData(context), truncator(TruncationFactory(part)), unit(unit_p) {
+		}
+
+		//! Truncates the instants the arithmetic cannot handle
+		part_trunc_t truncator;
+		int64_t unit;
+		//! Unset below a minute: offsets are whole seconds, so the instant truncates like its wall time.
+		shared_ptr<const ZoneOffsets> offsets;
+
+		unique_ptr<FunctionData> Copy() const override {
+			return make_uniq<ArithmeticTruncData>(*this);
+		}
+	};
+
+	//! The per-execution state of an arithmetic truncation.
+	struct ArithmeticTruncator {
+		explicit ArithmeticTruncator(const ArithmeticTruncData &data) : unit(data.unit), offsets(data.offsets.get()) {
+		}
+
+		const int64_t unit;
+		const ZoneOffsets *const offsets;
+		ZoneOffsets::Cursor instant_cursor;
+		ZoneOffsets::Cursor wall_cursor;
+
+		//! Offsets and units are under a day, so truncation moves an instant by less than three days.
+		static constexpr int64_t MIN_SAFE = NumericLimits<int64_t>::Minimum() + 3 * Interval::MICROS_PER_DAY;
+		static constexpr int64_t MAX_SAFE = NumericLimits<int64_t>::Maximum() - 3 * Interval::MICROS_PER_DAY;
+
+		static int64_t Floor(int64_t value, int64_t unit) {
+			const auto remainder = value % unit;
+			return value - (remainder < 0 ? remainder + unit : remainder);
+		}
+
+		bool TryTrunc(int64_t utc, int64_t &result) {
+			if (utc < MIN_SAFE || utc > MAX_SAFE) {
+				return false;
+			}
+			if (!offsets) {
+				result = Floor(utc, unit);
+				return true;
+			}
+			int64_t offset;
+			if (!offsets->OffsetAtInstant(utc, instant_cursor, offset)) {
+				return false;
+			}
+			const auto floored = Floor(utc + offset, unit);
+			//	Like ICU, keep the instant's offset below an hour (PreserveOffsets) and re-resolve it from an hour up.
+			if (unit >= Interval::MICROS_PER_HOUR && !offsets->OffsetAtWallTime(floored, wall_cursor, offset)) {
+				return false;
+			}
+			result = floored - offset;
+			return true;
+		}
+	};
+
+	template <typename T>
+	static void ArithmeticTruncFunction(DataChunk &args, ExpressionState &state, Vector &result) {
+		auto &info = state.expr.Cast<BoundFunctionExpression>().bind_info->Cast<ArithmeticTruncData>();
+		ArithmeticTruncator truncator(info);
+		CalendarPtr calendar;
+		UnaryExecutor::Execute<T, timestamp_t>(args.data[1], result, args.size(), [&](T input) {
+			if (!Timestamp::IsFinite(input)) {
+				return input;
+			}
+			int64_t truncated;
+			if (truncator.TryTrunc(input.value, truncated)) {
+				return timestamp_t(truncated);
+			}
+			if (!calendar) {
+				calendar.reset(info.calendar->clone());
+			}
+			return TruncWithICU(calendar.get(), info.truncator, input);
+		});
+	}
+
+	template <typename T>
+	static unique_ptr<FunctionData> BindDateTrunc(ClientContext &context, ScalarFunction &bound_function,
+	                                              vector<unique_ptr<Expression>> &arguments) {
+		if (!arguments[0]->IsFoldable()) {
+			return Bind(context, bound_function, arguments);
+		}
+		const auto part_value = ExpressionExecutor::EvaluateScalar(context, *arguments[0]);
+		DatePartSpecifier part;
+		if (part_value.IsNull() || !TryGetDatePartSpecifier(part_value.ToString(), part) || !FixedPartLength(part)) {
+			return Bind(context, bound_function, arguments);
+		}
+
+		auto data = make_uniq<ArithmeticTruncData>(context, part, FixedPartLength(part));
+		if (data->unit >= Interval::MICROS_PER_MINUTE) {
+			//	Only the offsets are modelled, so other calendars stay on ICU.
+			if (std::strcmp(data->calendar->getType(), "gregorian") == 0) {
+				data->offsets = ZoneOffsets::Build(data->calendar->getTimeZone());
+			}
+			if (!data->offsets) {
+				return Bind(context, bound_function, arguments);
+			}
+		}
+		bound_function.SetFunctionCallback(ArithmeticTruncFunction<T>);
+		return std::move(data);
+	}
+
 	template <typename T>
 	static void ICUDateTruncFunction(DataChunk &args, ExpressionState &state, Vector &result) {
 		D_ASSERT(args.ColumnCount() == 2);
@@ -143,9 +284,7 @@ struct ICUDateTrunc : public ICUDateFunc {
 				auto truncator = TruncationFactory(GetDatePartSpecifier(specifier));
 				UnaryExecutor::Execute<T, timestamp_t>(date_arg, result, args.size(), [&](T input) {
 					if (Timestamp::IsFinite(input)) {
-						auto micros = SetTime(calendar.get(), input);
-						truncator(calendar.get(), micros);
-						return GetTimeUnsafe(calendar.get(), micros);
+						return TruncWithICU(calendar.get(), truncator, input);
 					} else {
 						return input;
 					}
@@ -156,9 +295,7 @@ struct ICUDateTrunc : public ICUDateFunc {
 			    part_arg, date_arg, result, args.size(), [&](string_t specifier, T input) {
 				    if (Timestamp::IsFinite(input)) {
 					    auto truncator = TruncationFactory(GetDatePartSpecifier(specifier.GetString()));
-					    auto micros = SetTime(calendar.get(), input);
-					    truncator(calendar.get(), micros);
-					    return GetTimeUnsafe(calendar.get(), micros);
+					    return TruncWithICU(calendar.get(), truncator, input);
 				    } else {
 					    return input;
 				    }
@@ -168,7 +305,8 @@ struct ICUDateTrunc : public ICUDateFunc {
 
 	template <typename TA>
 	static ScalarFunction GetDateTruncFunction(const LogicalTypeId &type) {
-		return ScalarFunction({LogicalType::VARCHAR, type}, LogicalType::TIMESTAMP_TZ, ICUDateTruncFunction<TA>, Bind);
+		return ScalarFunction({LogicalType::VARCHAR, type}, LogicalType::TIMESTAMP_TZ, ICUDateTruncFunction<TA>,
+		                      BindDateTrunc<TA>);
 	}
 
 	static void AddBinaryTimestampFunction(const string &name, ExtensionLoader &loader) {

extension/icu/icu-zone-offsets.cpp

@@ -0,0 +1,76 @@
+#include "include/icu-zone-offsets.hpp"
+
+#include "duckdb/common/types/date.hpp"
+#include "duckdb/common/types/interval.hpp"
+#include "unicode/basictz.h"
+#include "unicode/tzrule.h"
+#include "unicode/tztrans.h"
+
+namespace duckdb {
+
+static int64_t RuleOffset(const icu::TimeZoneRule &rule) {
+	return (int64_t(rule.getRawOffset()) + rule.getDSTSavings()) * Interval::MICROS_PER_MSEC;
+}
+
+static bool IsUnderADay(int64_t offset) {
+	return offset > -Interval::MICROS_PER_DAY && offset < Interval::MICROS_PER_DAY;
+}
+
+shared_ptr<const ZoneOffsets> ZoneOffsets::Build(const icu::TimeZone &zone) {
+	const auto basic_zone = dynamic_cast<const icu::BasicTimeZone *>(&zone);
+	if (!basic_zone) {
+		return nullptr;
+	}
+
+	//	Recurring rules produce transitions forever, so stop at a year well past scan-scale data.
+	const int64_t coverage_end = int64_t(Date::FromDate(2250, 1, 1).days) * Interval::MICROS_PER_DAY;
+	//	Early enough to precede every tzdata transition, late enough for exact µs conversion.
+	UDate search_from = -1.0e15;
+
+	icu::TimeZoneTransition transition;
+	bool has_next = basic_zone->getNextTransition(search_from, false, transition);
+	int64_t offset;
+	if (has_next) {
+		offset = RuleOffset(*transition.getFrom());
+	} else {
+		int32_t raw = 0;
+		int32_t dst = 0;
+		UErrorCode status = U_ZERO_ERROR;
+		basic_zone->getOffset(0, false, raw, dst, status);
+		if (U_FAILURE(status)) {
+			return nullptr;
+		}
+		offset = (int64_t(raw) + dst) * Interval::MICROS_PER_MSEC;
+	}
+	if (!IsUnderADay(offset)) {
+		return nullptr;
+	}
+
+	auto result = make_shared_ptr<ZoneOffsets>();
+	auto &segments = result->segments;
+	segments.push_back(Segment {NumericLimits<int64_t>::Minimum(), NumericLimits<int64_t>::Minimum(), offset});
+	for (; has_next; has_next = basic_zone->getNextTransition(search_from, false, transition)) {
+		search_from = transition.getTime();
+		const auto utc = int64_t(search_from) * Interval::MICROS_PER_MSEC;
+		const auto offset_after = RuleOffset(*transition.getTo());
+		const bool past_coverage = utc >= coverage_end;
+		if (!past_coverage && offset_after == segments.back().offset) {
+			continue;
+		}
+		if (!IsUnderADay(offset_after)) {
+			return nullptr;
+		}
+		//	A fall-back longer than the time since the previous transition would make wall times non-monotonic.
+		if (utc + offset_after <= segments.back().start_local) {
+			return nullptr;
+		}
+		segments.push_back(Segment {utc, utc + offset_after, offset_after});
+		if (past_coverage) {
+			return std::move(result);
+		}
+	}
+	segments.push_back(Segment {NumericLimits<int64_t>::Maximum(), NumericLimits<int64_t>::Maximum(), 0});
+	return std::move(result);
+}
+
+} // namespace duckdb