Pushdown uuid functions

Signed-off-by: Aykut Bozkurt <aykut.bozkurt@snowflake.com>

Author: sfc-gh-abozkurt

duckdb_pglake/src/duckdb_pglake_extension.cpp

@@ -110,6 +110,123 @@ inline void AtanhPG(DataChunk &args, ExpressionState &state, Vector &result)
 }
 
 
+/*
+  * Extract timestamp from UUID, mimicking Postgres behavior.
+ * Postgres supports UUID v1 and v7 (above 18), returning NULL for other versions.
+ * DuckDB's uuid_extract_timestamp only supports v7 and throws an error for others.
+ */
+inline void UUIDExtractTimestampPG(DataChunk &args, ExpressionState &state, Vector &result)
+{
+	D_ASSERT(args.ColumnCount() == 2);
+	auto &input_vector = args.data[0];
+	auto &version_vector = args.data[1];
+	auto count = args.size();
+
+	// Prepare input in unified format
+	UnifiedVectorFormat vdata;
+	input_vector.ToUnifiedFormat(count, vdata);
+
+	UnifiedVectorFormat vdata_version;
+	version_vector.ToUnifiedFormat(count, vdata_version);
+
+	auto input_data = UnifiedVectorFormat::GetData<hugeint_t>(vdata);
+	auto pg_version_data = UnifiedVectorFormat::GetData<int32_t>(vdata_version);
+	auto result_data = FlatVector::GetData<timestamp_t>(result);
+	auto &result_validity = FlatVector::Validity(result);
+
+	for (idx_t i = 0; i < count; i++) {
+		auto idx = vdata.sel->get_index(i);
+		auto version_idx = vdata_version.sel->get_index(i);
+
+		// Propagate input NULL
+		if (!vdata.validity.RowIsValid(idx)) {
+			result_validity.SetInvalid(i);
+			continue;
+		}
+
+		auto uuid_val = input_data[idx];
+		auto pg_version = pg_version_data[version_idx];
+
+		// Check if RFC 4122 variant (bits 10xxxxxx in the variant field)
+		// The variant field is in byte 8 (counting from 0)
+		uint8_t variant_byte = static_cast<uint8_t>((uuid_val.lower >> 56) & 0xFF);
+		if ((variant_byte & 0xc0) != 0x80) {
+			// Not RFC 4122 variant, return NULL
+			result_validity.SetInvalid(i);
+			continue;
+		}
+
+		// Extract version (first 4 bits of byte 6)
+		// DuckDB stores UUID with XOR flip on the upper 64 bits, so we need to undo it
+		uint64_t unsigned_upper = static_cast<uint64_t>(uuid_val.upper) ^ (uint64_t(1) << 63);
+		uint8_t version = (static_cast<uint8_t>((unsigned_upper) >> 8) & 0xf0) >> 4;
+
+		if (version == 1) {
+			// UUID v1: Extract timestamp from time_low, time_mid, and time_hi_and_version fields
+			// Mimic PostgreSQL's uuid_extract_timestamp implementation
+
+			// Extract individual bytes from the UUID (bytes 0-7 are in upper)
+			// We already have unsigned_upper computed above
+			uint8_t data[8];
+			for (int i = 0; i < 8; i++) {
+				data[i] = static_cast<uint8_t>((unsigned_upper >> (56 - i * 8)) & 0xFF);
+			}
+
+			// Extract timestamp following PostgreSQL's exact logic
+			// See: src/backend/utils/adt/uuid.c:uuid_extract_timestamp()
+			uint64_t tms = ((uint64_t) data[0] << 24)
+				+ ((uint64_t) data[1] << 16)
+				+ ((uint64_t) data[2] << 8)
+				+ ((uint64_t) data[3])
+				+ ((uint64_t) data[4] << 40)
+				+ ((uint64_t) data[5] << 32)
+				+ (((uint64_t) data[6] & 0xf) << 56)
+				+ ((uint64_t) data[7] << 48);
+
+			// Convert 100-ns intervals to microseconds
+			int64_t timestamp_us = static_cast<int64_t>(tms / 10);
+
+			// Adjust from UUID epoch (1582-10-15) to Postgres epoch (2000-01-01)
+			// This matches PostgreSQL's calculation exactly
+			constexpr int64_t POSTGRES_EPOCH_JDATE = 2451545; // date2j(2000, 1, 1)
+			constexpr int64_t UUIDV1_EPOCH_JDATE = 2299161;   // date2j(1582, 10, 15)
+			constexpr int64_t SECS_PER_DAY = 86400;
+			constexpr int64_t USECS_PER_SEC = 1000000;
+			constexpr int64_t UUID_TO_PG_EPOCH_US = 
+				(POSTGRES_EPOCH_JDATE - UUIDV1_EPOCH_JDATE) * SECS_PER_DAY * USECS_PER_SEC;
+
+			timestamp_us -= UUID_TO_PG_EPOCH_US;
+
+			// Convert from Postgres epoch (2000-01-01) to Unix epoch (1970-01-01)
+			// Unix epoch is 946684800 seconds (30 years) before Postgres epoch
+			// So we ADD this offset to convert from Postgres timestamp to Unix timestamp
+			constexpr int64_t PG_TO_UNIX_EPOCH_US = 946684800LL * USECS_PER_SEC;
+			timestamp_us += PG_TO_UNIX_EPOCH_US;
+
+			result_data[i] = timestamp_t{timestamp_us};
+		}
+		// UUID v7 is supported in Postgres 18 and above
+		else if (version == 7 && pg_version >= 180000) {
+			// UUID v7: Extract timestamp from first 48 bits (Unix timestamp in milliseconds)
+			int64_t upper = uuid_val.upper;
+			// Flip the top byte to handle signed representation
+			upper ^= NumericLimits<int64_t>::Minimum();
+			int64_t unix_ts_milli = upper >> 16;
+
+			// Convert milliseconds to microseconds
+			constexpr int64_t kMilliToMicro = 1000;
+			int64_t unix_ts_us = kMilliToMicro * unix_ts_milli;
+
+			result_data[i] = timestamp_t{unix_ts_us};
+		}
+		else {
+			// Not a timestamp-containing UUID version, return NULL
+			result_validity.SetInvalid(i);
+		}
+	}
+}
+
+
 /*
  * InitcapPG implements the Postgres initcap(text) function for the
  * C collation.
@@ -350,6 +467,9 @@ static void LoadInternal(ExtensionLoader &loader) {
 	substr.AddFunction(ScalarFunction({LogicalType::VARCHAR, LogicalType::BIGINT, LogicalType::BIGINT}, LogicalType::VARCHAR, SubstringPG));
 	loader.RegisterFunction(substr);
 
+	auto uuid_extract_timestamp_pg = ScalarFunction("uuid_extract_timestamp_pg", {LogicalType::UUID, LogicalType::INTEGER}, LogicalType::TIMESTAMP_TZ, UUIDExtractTimestampPG);
+	loader.RegisterFunction(uuid_extract_timestamp_pg);
+
 	PgLakeUtilityFunctions::RegisterFunctions(loader);
 	PgLakeFileSystemFunctions::RegisterFunctions(loader);
 

pg_lake_engine/pg_lake_engine--3.2--3.3.sql

@@ -10,3 +10,9 @@ CREATE FUNCTION __lake__internal__nsp__.initcap_pg(text)
  LANGUAGE C
  IMMUTABLE PARALLEL SAFE STRICT
 AS 'MODULE_PATHNAME', $function$pg_lake_internal_dummy_function$function$;
+
+CREATE FUNCTION __lake__internal__nsp__.uuid_extract_timestamp_pg(uuid, integer)
+ RETURNS timestamp with time zone
+ LANGUAGE C
+ IMMUTABLE PARALLEL SAFE STRICT
+AS 'MODULE_PATHNAME', $function$pg_lake_internal_dummy_function$function$;

pg_lake_engine/src/pgduck/rewrite_query.c

@@ -24,6 +24,7 @@
 #include "catalog/pg_operator.h"
 #include "catalog/pg_operator_d.h"
 #include "catalog/pg_proc.h"
+#include "catalog/pg_type.h"
 #include "common/hashfn.h"
 #include "pg_lake/extensions/pg_lake_spatial.h"
 #include "pg_lake/extensions/postgis.h"
@@ -146,6 +147,9 @@ static ExpressionRewriter GetOperatorRewriter(Oid functionId);
 static char *GetOperatorRewriteFunctionName(Oid functionId);
 
 static Node *RewriteFuncExprBtrim(Node *node, void *context);
+#if PG_VERSION_NUM >= 170000
+static Node *RewriteFuncExprUuidExtractTimestamp(Node *node, void *context);
+#endif
 static Node *RewriteFunctionCallExpression(Oid functionId, Node *node, void *context);
 static Node *RewriteFuncExprCast(Node *node, void *context);
 static Node *RewriteFuncExprExtract(Node *node, void *context);
@@ -350,6 +354,13 @@ static FunctionCallRewriteRuleByName BuiltinFunctionCallRewriteRulesByName[] =
 		"pg_catalog", "decode", RewriteFuncExprDecode, 0
 	},
 
+#if PG_VERSION_NUM >= 170000
+	/* uuid functions */
+	{
+		"pg_catalog", "uuid_extract_timestamp", RewriteFuncExprUuidExtractTimestamp, 0
+	},
+#endif
+
 };
 
 static OperatorRewriteRuleByName BuiltinOperatorRewriteRulesByName[] =
@@ -1689,6 +1700,45 @@ RewriteFuncExprDecode(Node *node, void *context)
 }
 
 
+#if PG_VERSION_NUM >= 170000
+/*
+ * RewriteFuncExprUuidExtractTimestamp rewrites uuid_extract_timestamp(..) function calls
+ * to use the custom uuid_extract_timestamp_pg(..) function that mimics Postgres behavior.
+ *
+ * Postgres returns NULL for UUID versions that don't contain timestamps (v1 and v7 (>=pg18 ) only).
+ * DuckDB's uuid_extract_timestamp throws an error for non-v7 UUIDs.
+ */
+static Node *
+RewriteFuncExprUuidExtractTimestamp(Node *node, void *context)
+{
+	FuncExpr   *funcExpr = castNode(FuncExpr, node);
+	int			argCount = list_length(funcExpr->args);
+
+	/* uuid_extract_timestamp should have exactly 1 argument */
+	if (argCount != 1)
+		return node;
+
+	Node	   *firstArg = linitial(funcExpr->args);
+
+	if (exprType(firstArg) != UUIDOID)
+		return node;
+
+	/* add the second argument, the Postgres version number */
+	funcExpr->args = list_make2(firstArg, MakeIntConst(PG_VERSION_NUM));
+	argCount = 2;
+
+	/* Rewrite to uuid_extract_timestamp_pg(uuid) */
+	List	   *funcName = list_make2(makeString(PG_LAKE_INTERNAL_NSP),
+									  makeString("uuid_extract_timestamp_pg"));
+	Oid			argTypes[] = {UUIDOID, INT4OID};
+
+	funcExpr->funcid = LookupFuncName(funcName, argCount, argTypes, false);
+
+	return (Node *) funcExpr;
+}
+#endif
+
+
 /*
  * RewriteFuncExprMapExtract rewrites map_type.extract(..) function calls
  * into remote map_extract(map,key)[1] expressions

pg_lake_engine/src/pgduck/shippable_builtin_functions.c

@@ -444,6 +444,14 @@ static const PGDuckShippableFunction ShippableBuiltinProcs[] =
 	{"encode", 'f', 2, {"bytea", "text"}, IsEncodeShippable},
 	{"decode", 'f', 2, {"text", "text"}, IsDecodeShippable},
 
+#if PG_VERSION_NUM >= 170000
+	{"uuid_extract_version", 'f', 1, {"uuid"}, NULL},
+	{"uuid_extract_timestamp", 'f', 1, {"uuid"}, NULL},
+#endif
+#if PG_VERSION_NUM >= 180000
+	{"uuidv7", 'f', 0, {}, NULL},
+#endif
+
 	/* trim() */
 };
 

pg_lake_spatial/tests/pytests/test_internal_schema.py

@@ -21,6 +21,6 @@ def test_internal_schema(
     """,
         pg_conn,
     )[0][0]
-    assert result == 64
+    assert result == 65
 
     pg_conn.rollback()