feat: native Spark add_months scalar function

CLAUDE.md

@@ -251,6 +251,7 @@ CteBaseNode (base, ToQuery(SqlDialect))
 | `test/sql/data_dependent_optimizers.test` | `lpts_enable_data_dependent_optimizers` |
 | `test/sql/tpch.test` | All 22 TPC-H queries under `lpts_check` |
 | `test/sql/ducklake.test` | DuckLake scans |
+| `test/sql/time_travel.test` | Snapshot pinning: Spark/Delta `VERSION`/`TIMESTAMP AS OF` ↔ DuckDB `AT (...)` |
 | `test/sql/explain_format_sql.test` | `EXPLAIN (FORMAT SQL)` |
 | `test/sql/print_ast.test` | AST `ToString()` output |
 | `test/sql/check_mode.test` | `lpts_check` round-trip semantics (canonical example) |

CMakeLists.txt

@@ -23,7 +23,8 @@ set(EXTENSION_SOURCES
     src/lpts_ast_renderer.cpp
     src/lpts_ast_builder.cpp
     src/lpts_ast_flattener.cpp
-    src/dialect_function_map.cpp)
+    src/dialect_function_map.cpp
+    src/spark_scalar_functions.cpp)
 
 build_static_extension(${TARGET_NAME} ${EXTENSION_SOURCES})
 build_loadable_extension(${TARGET_NAME} " " ${EXTENSION_SOURCES})

src/cte_nodes.cpp

@@ -384,9 +384,17 @@ string GetNode::ToQuery(SqlDialect dialect) {
 		get_str << VecToSeparatedList(RenderGetSelectColumns(column_names, column_is_expression, dialect));
 	}
 	get_str << " FROM ";
+	string base_table_name;
+	string snapshot_suffix;
+	const bool unqualified_snapshot =
+	    catalog.empty() && TrySplitDialectSnapshotSuffix(table_name, dialect, base_table_name, snapshot_suffix);
 	if (!catalog.empty()) {
 		// Fully-qualified: catalog.schema.table (DuckDB / Spark dialect)
 		get_str << DialectQualifiedTableName(catalog, schema, table_name, dialect);
+	} else if (unqualified_snapshot) {
+		// A pinned-snapshot scan rendered unqualified: the qualifier is dialect-specific and must not be
+		// mistaken for a table-function argument list by the `_tf` aliasing below.
+		get_str << base_table_name << snapshot_suffix;
 	} else {
 		// A TABLE-argument function: the child CTE is the function's argument, not a lateral input.
 		const size_t table_arg_pos = table_name.find("%LPTS_TABLE_ARG%");

src/include/lpts_helpers.hpp

@@ -25,6 +25,12 @@ string VecToQuotedIdentifierList(const vector<string> &input_list, const string
 /// Quote a table name, preserving a DuckDB AT (...) snapshot suffix if present.
 string QuoteTableWithOptionalSuffix(const string &table_name);
 
+/// Split a table name that carries a pinned snapshot (`name AT (<PARAM> => <value>)`, the DuckDB
+/// spelling LPTS uses internally) into `base_name` and the `dialect`-rendered snapshot qualifier
+/// (e.g. ` VERSION AS OF 366` for Spark). Returns false and leaves the outputs untouched when
+/// `table_name` carries no snapshot. Throws when `dialect` has no verified time-travel syntax.
+bool TrySplitDialectSnapshotSuffix(const string &table_name, SqlDialect dialect, string &base_name, string &suffix);
+
 /// Build catalog.schema.table with each identifier quoted when needed.
 string QualifiedTableName(const string &catalog, const string &schema, const string &table_name);
 

src/include/spark_scalar_functions.hpp

@@ -0,0 +1,11 @@
+#pragma once
+
+#include "duckdb/main/extension/extension_loader.hpp"
+
+namespace duckdb {
+
+// Registers Spark-compatible scalar functions that DuckDB lacks natively but that
+// appear in Spark SQL translated through lpts (compile/binding coverage).
+void RegisterSparkScalarFunctions(ExtensionLoader &loader);
+
+} // namespace duckdb

src/lpts_ast_builder.cpp

@@ -99,6 +99,8 @@ class AstBuilder {
 		}
 	};
 
+	using ColumnMap = std::map<MappableColumnBinding, unique_ptr<ColStruct>>;
+
 	// A set of emitted CTE column names with case-insensitive membership. DuckDB resolves identifiers
 	// case-insensitively, so two generated names differing only in case (t1_hello vs t1_HeLlO) would
 	// collide when referenced. Deduping against this set treats them as equal, so the second gets a suffix.
@@ -142,6 +144,20 @@ class AstBuilder {
 		return false;
 	}
 
+	static bool IsSetOperationType(LogicalOperatorType type) {
+		return type == LogicalOperatorType::LOGICAL_UNION || type == LogicalOperatorType::LOGICAL_EXCEPT ||
+		       type == LogicalOperatorType::LOGICAL_INTERSECT;
+	}
+
+	static ColumnMap CloneColumnMap(const ColumnMap &source) {
+		ColumnMap cloned;
+		for (const auto &entry : source) {
+			cloned[entry.first] =
+			    make_uniq<ColStruct>(entry.second->table_index, entry.second->column_name, entry.second->alias);
+		}
+		return cloned;
+	}
+
 	static void AddUniqueBinding(vector<ColumnBinding> &bindings, const ColumnBinding &binding) {
 		if (!HasBinding(bindings, binding)) {
 			bindings.push_back(binding);
@@ -462,7 +478,7 @@ class AstBuilder {
 
 	/// Global map: ColumnBinding → ColStruct.
 	/// Populated bottom-up; each operator registers its output columns here.
-	std::map<MappableColumnBinding, unique_ptr<ColStruct>> column_map;
+	ColumnMap column_map;
 
 	/// Maps DELIM_GET table_index → source column names (from the outer/left CTE).
 	/// Populated by PreregisterDelimGetColumns before the right subtree is traversed.
@@ -491,6 +507,23 @@ class AstBuilder {
 		    context, (unsigned long long)binding.table_index, (unsigned long long)binding.column_index);
 	}
 
+	bool TryResolveProjectionBinding(const LogicalProjection &proj, const ColumnBinding &binding,
+	                                 ColumnBinding &resolved) const {
+		if (proj.children.empty()) {
+			return false;
+		}
+		const auto child_bindings = proj.children[0]->GetColumnBindings();
+		resolved = binding;
+		if ((!HasBinding(child_bindings, resolved) && IsSetOperationType(proj.children[0]->type)) ||
+		    column_map.find(MappableColumnBinding(resolved)) == column_map.end()) {
+			if (resolved.column_index >= child_bindings.size()) {
+				return false;
+			}
+			resolved = child_bindings[resolved.column_index];
+		}
+		return true;
+	}
+
 	void RegisterChildBindingFallbacks(Expression &expr, const vector<ColumnBinding> &child_bindings) {
 		if (expr.type == ExpressionType::BOUND_COLUMN_REF) {
 			auto &bcr = expr.Cast<BoundColumnRefExpression>();
@@ -508,6 +541,25 @@ class AstBuilder {
 		});
 	}
 
+	void RegisterChildBindingFallbacks(Expression &expr, const LogicalProjection &proj) {
+		if (expr.type == ExpressionType::BOUND_COLUMN_REF) {
+			auto &bcr = expr.Cast<BoundColumnRefExpression>();
+			ColumnBinding resolved;
+			if (TryResolveProjectionBinding(proj, bcr.binding, resolved) &&
+			    (!(resolved == bcr.binding) ||
+			     column_map.find(MappableColumnBinding(bcr.binding)) == column_map.end())) {
+				auto &src = FindColumnBinding(resolved, "projection fallback");
+				column_map[MappableColumnBinding(bcr.binding)] =
+				    make_uniq<ColStruct>(src->table_index, src->column_name, src->alias);
+			}
+		}
+		ExpressionIterator::EnumerateChildren(expr, [&](unique_ptr<Expression> &child) {
+			if (child) {
+				RegisterChildBindingFallbacks(*child, proj);
+			}
+		});
+	}
+
 	bool EnsureBindingAvailableFrom(LogicalOperator *op, const ColumnBinding &binding) {
 		if (!op) {
 			return false;
@@ -835,12 +887,8 @@ class AstBuilder {
 		}
 		const auto child_bindings = proj.children[0]->GetColumnBindings();
 		auto &bcr = expr->Cast<BoundColumnRefExpression>();
-		resolved = bcr.binding;
-		if (column_map.find(MappableColumnBinding(resolved)) == column_map.end()) {
-			if (resolved.column_index >= child_bindings.size()) {
-				return false;
-			}
-			resolved = child_bindings[resolved.column_index];
+		if (!TryResolveProjectionBinding(proj, bcr.binding, resolved)) {
+			return false;
 		}
 		if (ordinal >= child_bindings.size()) {
 			return false;
@@ -1555,12 +1603,8 @@ class AstBuilder {
 				if (expr->type == ExpressionType::BOUND_COLUMN_REF) {
 					BoundColumnRefExpression &bcr = expr->Cast<BoundColumnRefExpression>();
 					ColumnBinding lookup_binding = bcr.binding;
-					if (column_map.find(MappableColumnBinding(lookup_binding)) == column_map.end() &&
-					    !proj.children.empty()) {
-						auto child_bindings = proj.children[0]->GetColumnBindings();
-						if (lookup_binding.column_index < child_bindings.size()) {
-							lookup_binding = child_bindings[lookup_binding.column_index];
-						}
+					if (TryResolveProjectionBinding(proj, bcr.binding, lookup_binding)) {
+						// lookup_binding already resolved against the set-op child when needed.
 					}
 					const unique_ptr<ColStruct> &desc = FindColumnBinding(lookup_binding, "projection");
 					const string src_name = desc->ToUniqueColumnName();
@@ -1585,7 +1629,7 @@ class AstBuilder {
 					column_map[MappableColumnBinding(new_cb)] = std::move(new_col);
 				} else {
 					if (!proj.children.empty()) {
-						RegisterChildBindingFallbacks(*expr, proj.children[0]->GetColumnBindings());
+						RegisterChildBindingFallbacks(*expr, proj);
 					}
 					string expr_str = ExpressionToAliasedString(expr);
 					expressions.emplace_back(expr_str);
@@ -2085,17 +2129,36 @@ class AstBuilder {
 			vector<string> cte_column_names;
 			const auto &lhs_bindings = op->children[0]->GetColumnBindings();
 			const auto &union_bindings = op->GetColumnBindings();
+			const idx_t physical_column_count =
+			    op->types.empty() ? std::min(lhs_bindings.size(), union_bindings.size()) : op->types.size();
+			if (lhs_bindings.size() < physical_column_count || union_bindings.size() < physical_column_count) {
+				throw InternalException("LPTS UNION: physical output arity exceeds available column bindings");
+			}
 			// Two union output columns can derive from source columns with the same name (e.g.
 			// SELECT t1.a, t2.a ... UNION ...), which would emit a header like (t7_a, t7_a) — DuckDB
 			// resolves later references to the first, silently dropping the second column. Dedup so each
 			// output column gets a distinct generated name.
 			CaseInsensitiveNameSet seen_names;
-			for (size_t i = 0; i < lhs_bindings.size(); ++i) {
+			for (idx_t i = 0; i < physical_column_count; ++i) {
 				const unique_ptr<ColStruct> &lhs_col = FindColumnBinding(lhs_bindings[i], "union lhs");
 				auto new_col = MakeDedupedColumn(table_index, lhs_col->column_name, lhs_col->alias, seen_names, 1);
 				cte_column_names.push_back(new_col->ToUniqueColumnName());
 				column_map[MappableColumnBinding(union_bindings[i])] = std::move(new_col);
 			}
+			if (union_bindings.size() > physical_column_count) {
+				// Some post-optimizer rewrites expose one trailing alias binding for the
+				// physical multiplicity column without widening the UNION types or children.
+				// Preserve that binding by pointing it at the existing last positional output.
+				if (physical_column_count == 0 || union_bindings.size() != physical_column_count + 1) {
+					throw NotImplementedException(
+					    "LPTS_UNSUPPORTED_COLUMN_REF: UNION exposes %llu bindings for %llu physical columns",
+					    (unsigned long long)union_bindings.size(), (unsigned long long)physical_column_count);
+				}
+				const auto &last_output =
+				    FindColumnBinding(union_bindings[physical_column_count - 1], "union trailing alias source");
+				column_map[MappableColumnBinding(union_bindings[physical_column_count])] =
+				    make_uniq<ColStruct>(last_output->table_index, last_output->column_name, last_output->alias);
+			}
 			return make_uniq<AstUnionNode>(set_op.setop_all, std::move(cte_column_names));
 		}
 
@@ -2618,19 +2681,15 @@ class AstBuilder {
 	unique_ptr<AstNode> RecursiveTraversal(unique_ptr<LogicalOperator> &op, bool is_root = false) {
 		// 1. Recurse into children first (post-order).
 		vector<unique_ptr<AstNode>> child_nodes;
-		if ((op->type == LogicalOperatorType::LOGICAL_UNION || op->type == LogicalOperatorType::LOGICAL_EXCEPT ||
-		     op->type == LogicalOperatorType::LOGICAL_INTERSECT) &&
-		    op->children.size() >= 2) {
+		if (IsSetOperationType(op->type) && op->children.size() >= 2) {
 			// Set operations: scope column_map to prevent sibling children from overwriting
 			// each other's entries when subtrees share table indices.
 			child_nodes.push_back(RecursiveTraversal(op->children[0]));
-			// Save column_map after first child; restore before each subsequent child
-			std::map<MappableColumnBinding, unique_ptr<ColStruct>> saved_map;
-			for (auto &entry : column_map) {
-				saved_map[entry.first] =
-				    make_uniq<ColStruct>(entry.second->table_index, entry.second->column_name, entry.second->alias);
-			}
+			// Save the lhs-visible bindings, then replay each sibling under that same scope so one branch
+			// cannot leak its table-index mappings into another or into the parent set-op output.
+			ColumnMap saved_map = CloneColumnMap(column_map);
 			for (size_t ci = 1; ci < op->children.size(); ci++) {
+				column_map = CloneColumnMap(saved_map);
 				child_nodes.push_back(RecursiveTraversal(op->children[ci]));
 			}
 			column_map = std::move(saved_map);

src/lpts_ast_flattener.cpp

@@ -289,8 +289,15 @@ class AstFlattener {
 				}
 			}
 			sql += " FROM ";
+			string inline_base_name;
+			string inline_snapshot_suffix;
 			if (!get.catalog.empty()) {
 				sql += DialectQualifiedTableName(get.catalog, get.schema, get.table_name, dialect);
+			} else if (TrySplitDialectSnapshotSuffix(get.table_name, dialect, inline_base_name,
+			                                         inline_snapshot_suffix)) {
+				// Pinned-snapshot scan rendered unqualified: the dialect-specific qualifier must not be
+				// mistaken for a table-function argument list by the `_tf` aliasing below.
+				sql += inline_base_name + inline_snapshot_suffix;
 			} else {
 				// An in-out (lateral) table function has the delim/correlation source as its AST child:
 				// inline it as the left comma-join input (mirrors GetNode's input_cte_name handling).

src/lpts_extension.cpp

@@ -8,6 +8,7 @@
 #include "lpts_helpers.hpp"
 #include "lpts_debug.hpp"
 #include "lpts_parser.hpp"
+#include "spark_scalar_functions.hpp"
 
 #include "duckdb.hpp"
 #include "duckdb/common/exception.hpp"
@@ -1315,6 +1316,10 @@ static void LptsCheckOptimize(OptimizerExtensionInput &input, unique_ptr<Logical
 //------------------------------------------------------------------------------
 
 static void LoadInternal(ExtensionLoader &loader) {
+	// Register Spark-compatible scalar functions (e.g. add_months) that DuckDB
+	// lacks natively but that appear in Spark SQL translated through lpts.
+	RegisterSparkScalarFunctions(loader);
+
 	// Register the lpts_dialect session setting.
 	// Users can change it with: SET lpts_dialect = 'postgres';
 	DBConfig &config = DBConfig::GetConfig(loader.GetDatabaseInstance());

src/lpts_helpers.cpp

@@ -81,13 +81,67 @@ string DialectVecToQuotedIdentifierList(const vector<string> &input_list, SqlDia
 	return ret_str.str();
 }
 
-string DialectQuoteTableWithOptionalSuffix(const string &table_name, SqlDialect dialect) {
-	static const string at_suffix = " AT (";
-	auto suffix_pos = table_name.find(at_suffix);
+/// Split `name AT (<PARAM> => <value>)` — the DuckDB spelling LPTS uses internally to carry a pinned
+/// snapshot on a table name — into its base name and `AT` parameter/value.
+static bool TrySplitSnapshotSuffix(const string &table_name, string &base_name, string &at_parameter,
+                                   string &at_value) {
+	static const string AT_SUFFIX = " AT (";
+	auto suffix_pos = table_name.find(AT_SUFFIX);
 	if (suffix_pos == string::npos) {
+		return false;
+	}
+	string body = table_name.substr(suffix_pos + AT_SUFFIX.size());
+	if (body.empty() || body.back() != ')') {
+		return false;
+	}
+	body.pop_back();
+	auto arrow_pos = body.find("=>");
+	if (arrow_pos == string::npos) {
+		return false;
+	}
+	base_name = table_name.substr(0, suffix_pos);
+	at_parameter = TrimCopy(body.substr(0, arrow_pos));
+	at_value = TrimCopy(body.substr(arrow_pos + 2));
+	return !at_parameter.empty() && !at_value.empty();
+}
+
+/// Render a pinned-snapshot qualifier in `dialect`. LPTS carries the pin in DuckDB's spelling
+/// (`AT (VERSION => 366)`); Spark/Delta spells the same pin `VERSION AS OF 366`. A dialect with no
+/// verified time-travel syntax refuses instead of emitting a qualifier the target cannot parse —
+/// dropping it would silently turn a pinned scan into a read of the latest snapshot.
+static string DialectSnapshotSuffix(const string &base_name, const string &at_parameter, const string &at_value,
+                                    SqlDialect dialect) {
+	if (dialect == SqlDialect::DUCKDB) {
+		return " AT (" + at_parameter + " => " + at_value + ")";
+	}
+	string parameter = LowerCopy(at_parameter);
+	if (dialect == SqlDialect::SPARK && (parameter == "version" || parameter == "timestamp")) {
+		return (parameter == "version" ? string(" VERSION AS OF ") : string(" TIMESTAMP AS OF ")) + at_value;
+	}
+	ThrowLptsNotImplemented("LPTS_UNSUPPORTED_TIME_TRAVEL", dialect, "time_travel", at_parameter + " => " + at_value,
+	                        base_name, "no verified time-travel syntax for target dialect");
+}
+
+string DialectQuoteTableWithOptionalSuffix(const string &table_name, SqlDialect dialect) {
+	string base_name;
+	string at_parameter;
+	string at_value;
+	if (!TrySplitSnapshotSuffix(table_name, base_name, at_parameter, at_value)) {
 		return DialectQuoteIdent(table_name, dialect);
 	}
-	return DialectQuoteIdent(table_name.substr(0, suffix_pos), dialect) + table_name.substr(suffix_pos);
+	return DialectQuoteIdent(base_name, dialect) + DialectSnapshotSuffix(base_name, at_parameter, at_value, dialect);
+}
+
+bool TrySplitDialectSnapshotSuffix(const string &table_name, SqlDialect dialect, string &base_name, string &suffix) {
+	string at_parameter;
+	string at_value;
+	string split_base;
+	if (!TrySplitSnapshotSuffix(table_name, split_base, at_parameter, at_value)) {
+		return false;
+	}
+	suffix = DialectSnapshotSuffix(split_base, at_parameter, at_value, dialect);
+	base_name = split_base;
+	return true;
 }
 
 string DialectQualifiedTableName(const string &catalog, const string &schema, const string &table_name,