Merge pull request #7 from spiceai/mach/compound-art-scan

ART Index: Support compound key scans

Squashed commit of the following:

commit a22a430
Author: David Stancu <david@spice.ai>
Date:   Wed Oct 29 16:37:30 2025 -0400

    simplify filter expression storage index bindings (just reuse the ones we made earlier), fix single-ref-per-expr predicate to correctly walk expr tree and yank refs (allowing nesting in fns, etc)

commit 9c8c1ed
Author: David Stancu <david@spice.ai>
Date:   Wed Oct 29 15:11:23 2025 -0400

    copy index expressions before rewriting column refs

commit aff2c98
Author: David Stancu <david@spice.ai>
Date:   Wed Oct 29 14:36:33 2025 -0400

    table scan:

    rebind projected columns in ALL index exprs
    do not bail out early if more than one index expr
    hook up composite key scan

commit bfc6f02
Author: David Stancu <david@spice.ai>
Date:   Wed Oct 29 14:35:09 2025 -0400

    make specialized compound key scan state for eq compares, specialized scan using ARTKey::Concat

Author: mach-kernel

src/execution/index/art/art.cpp

@@ -1,5 +1,7 @@
 #include "duckdb/execution/index/art/art.hpp"
 
+#include "duckdb/common/helper.hpp"
+#include "duckdb/common/typedefs.hpp"
 #include "duckdb/common/types/conflict_manager.hpp"
 #include "duckdb/common/unordered_map.hpp"
 #include "duckdb/common/vector_operations/vector_operations.hpp"
@@ -39,6 +41,17 @@ struct ARTIndexScanState : public IndexScanState {
 	set<row_t> row_ids;
 };
 
+struct ARTIndexCompoundKeyScanState : public IndexScanState {
+	//! The predicates to scan.
+	//! A single predicate for each constituent key in a compound index.
+	vector<Value> values;
+	//! The expressions over the scan predicates.
+	vector<ExpressionType> expressions;
+	bool checked = false;
+	//! All scanned row IDs.
+	set<row_t> row_ids;
+};
+
 //===--------------------------------------------------------------------===//
 // ART
 //===--------------------------------------------------------------------===//
@@ -142,6 +155,34 @@ static unique_ptr<IndexScanState> InitializeScanTwoPredicates(const Value &low_v
 	return std::move(result);
 }
 
+// Build compound scan state by building individual index scans and collecting their exprs/values
+unique_ptr<IndexScanState> ART::TryInitializeCompoundKeyScan(const vector<unique_ptr<Expression>> &index_exprs,
+                                                             vector<vector<unique_ptr<Expression>>> &exprs) {
+	auto compound_scan_state = make_uniq<ARTIndexCompoundKeyScanState>();
+
+	for (idx_t i = 0; i < index_exprs.size(); ++i) {
+		auto index_expr = &index_exprs[i];
+		auto filter_exprs = &exprs[i];
+
+		for (const auto &filter_expr : *filter_exprs) {
+			auto single_scan = ART::TryInitializeScan(**index_expr, *filter_expr);
+			if (!single_scan) {
+				return nullptr;
+			}
+
+			auto single_scan_concrete = single_scan->Cast<ARTIndexScanState>();
+			if (single_scan_concrete.expressions[0] != ExpressionType::COMPARE_EQUAL) {
+				return nullptr;
+			}
+
+			compound_scan_state->values.push_back(single_scan_concrete.values[0]);
+			compound_scan_state->expressions.push_back(single_scan_concrete.expressions[0]);
+		}
+	}
+
+	return compound_scan_state;
+}
+
 unique_ptr<IndexScanState> ART::TryInitializeScan(const Expression &expr, const Expression &filter_expr) {
 	Value low_value, high_value, equal_value;
 	ExpressionType low_comparison_type = ExpressionType::INVALID, high_comparison_type = ExpressionType::INVALID;
@@ -675,6 +716,23 @@ bool ART::SearchCloseRange(ARTKey &lower_bound, ARTKey &upper_bound, bool left_e
 	return it.Scan(upper_bound, max_count, row_ids, right_equal);
 }
 
+bool ART::CompoundKeyScan(IndexScanState &state, const idx_t max_count, set<row_t> &row_ids) {
+	auto &scan_state = state.Cast<ARTIndexCompoundKeyScanState>();
+	D_ASSERT(scan_state.values.size() == types.size());
+
+	ArenaAllocator arena_allocator(Allocator::Get(db));
+
+	// Make a compound key from the collected state values
+	auto compound_key = ARTKey::CreateKey(arena_allocator, types[0], scan_state.values[0]);
+	for (idx_t i = 1; i < scan_state.values.size(); ++i) {
+		auto part_key = ARTKey::CreateKey(arena_allocator, types[i], scan_state.values[i]);
+		compound_key.Concat(arena_allocator, part_key);
+	}
+
+	lock_guard<mutex> l(lock);
+	return SearchEqual(compound_key, max_count, row_ids);
+}
+
 bool ART::Scan(IndexScanState &state, const idx_t max_count, set<row_t> &row_ids) {
 	auto &scan_state = state.Cast<ARTIndexScanState>();
 	D_ASSERT(scan_state.values[0].type().InternalType() == types[0]);

src/function/table/table_scan.cpp

@@ -30,7 +30,9 @@
 #include "duckdb/planner/filter/conjunction_filter.hpp"
 #include "duckdb/common/types/value_map.hpp"
 #include "duckdb/main/settings.hpp"
+#include <limits>
 #include <list>
+#include <utility>
 
 namespace duckdb {
 
@@ -486,86 +488,123 @@ vector<unique_ptr<Expression>> ExtractFilterExpressions(const ColumnDefinition &
 
 bool TryScanIndex(ART &art, const ColumnList &column_list, TableFunctionInitInput &input, TableFilterSet &filter_set,
                   idx_t max_count, set<row_t> &row_ids) {
-	// FIXME: No support for index scans on compound ARTs.
-	// See note above on multi-filter support.
-	if (art.unbound_expressions.size() > 1) {
-		return false;
+	vector<unique_ptr<Expression>> index_exprs;
+	for (const auto &expr : art.unbound_expressions) {
+		index_exprs.push_back(expr->Copy());
 	}
 
-	auto index_expr = art.unbound_expressions[0]->Copy();
 	auto &indexed_columns = art.GetColumnIds();
 
-	// NOTE: We do not push down multi-column filters, e.g., 42 = a + b.
-	if (indexed_columns.size() != 1) {
+	// Allow composite ART scans
+	if (indexed_columns.size() != index_exprs.size()) {
 		return false;
 	}
 
-	// Resolve bound column references in the index_expr against the current input projection
-	column_t updated_index_column;
-	bool found_index_column_in_input = false;
+	// ...only if each expression has a single column reference, and
+	// that single column reference positionally matches that of the index (this is guaranteed? paranoid?)
+	// NOTE: We do not push down multi-column filters, e.g., 42 = a + b.
+	for (idx_t i = 0; i < index_exprs.size(); ++i) {
+		unordered_set<column_t> referenced_columns;
+		auto expr = &index_exprs[i];
+
+		// Walk the expr in case of nesting (e.g. function)
+		ExpressionIterator::EnumerateExpression(*expr, [&](Expression &child_expr) {
+			if (child_expr.GetExpressionClass() == ExpressionClass::BOUND_COLUMN_REF) {
+				auto &col_ref = child_expr.Cast<BoundColumnRefExpression>();
+				referenced_columns.insert(col_ref.binding.column_index);
+			}
+		});
 
-	// Find the indexed column amongst the input columns
-	for (idx_t i = 0; i < input.column_ids.size(); ++i) {
-		if (input.column_ids[i] == indexed_columns[0]) {
-			updated_index_column = i;
-			found_index_column_in_input = true;
-			break;
+		if (referenced_columns.size() != 1 || *referenced_columns.begin() != indexed_columns[i]) {
+			return false;
 		}
 	}
 
-	// If found, update the bound column ref within index_expr
-	if (found_index_column_in_input) {
-		ExpressionIterator::EnumerateExpression(index_expr, [&](Expression &expr) {
-			if (expr.GetExpressionClass() != ExpressionClass::BOUND_COLUMN_REF) {
-				return;
-			}
+	// Resolve bound column references in the index_expr against the current input projection
+	vector<column_t> index_column_to_proj_pos;
+	index_column_to_proj_pos.resize(indexed_columns.size(), std::numeric_limits<idx_t>::max());
 
-			auto &bound_column_ref_expr = expr.Cast<BoundColumnRefExpression>();
+	bool found_index_column_in_input = false;
 
-			// If the bound column references the index column, use updated_index_column
-			if (bound_column_ref_expr.binding.column_index == indexed_columns[0]) {
-				bound_column_ref_expr.binding.column_index = updated_index_column;
+	// Associate indexed columns to input columns
+	for (idx_t i = 0; i < indexed_columns.size(); ++i) {
+		for (idx_t j = 0; j < input.column_ids.size(); ++j) {
+			if (indexed_columns[i] == input.column_ids[j]) {
+				index_column_to_proj_pos.at(i) = j;
+				found_index_column_in_input = true;
 			}
-		});
+		}
 	}
 
-	// Get ART column.
-	auto &col = column_list.GetColumn(LogicalIndex(indexed_columns[0]));
-
-	// The indexes of the filters match input.column_indexes, which are: i -> column_index.
-	// Try to find a filter on the ART column.
-	optional_idx storage_index;
-	for (idx_t i = 0; i < input.column_indexes.size(); i++) {
-		if (input.column_indexes[i].ToLogical() == col.Logical()) {
-			storage_index = i;
-			break;
+	// If found, update the bound column refs within all index_exprs
+	if (found_index_column_in_input) {
+		for (auto &index_expr : index_exprs) {
+			ExpressionIterator::EnumerateExpression(index_expr, [&](Expression &expr) {
+				if (expr.GetExpressionClass() != ExpressionClass::BOUND_COLUMN_REF) {
+					return;
+				}
+
+				auto &bound_column_ref_expr = expr.Cast<BoundColumnRefExpression>();
+
+				// If the bound column references an indexed column, update it
+				for (idx_t i = 0; i < indexed_columns.size(); ++i) {
+					if (bound_column_ref_expr.binding.column_index == indexed_columns[i]) {
+						bound_column_ref_expr.binding.column_index = index_column_to_proj_pos[i];
+						break;
+					}
+				}
+			});
 		}
 	}
 
-	// No filter matches the ART column.
-	if (!storage_index.IsValid()) {
-		return false;
+	// The indexes of the filters match input.column_indexes, which are: i -> column_index.
+	// Reuse the index <-> projection mappings from index expr rebinding
+	vector<vector<unique_ptr<Expression>>> filters;
+
+	for (idx_t i = 0; i < index_column_to_proj_pos.size(); ++i) {
+		auto column_def = &column_list.GetColumn(LogicalIndex(indexed_columns[i]));
+		auto maybe_filter = filter_set.filters.find(index_column_to_proj_pos[i]);
+		if (maybe_filter != filter_set.filters.end()) {
+			auto filter = &maybe_filter->second;
+			auto filter_expressions = ExtractFilterExpressions(*column_def, *filter, index_column_to_proj_pos[i]);
+
+			filters.push_back(std::move(filter_expressions));
+		}
 	}
 
-	// Try to find a matching filter for the column.
-	auto filter = filter_set.filters.find(storage_index.GetIndex());
-	if (filter == filter_set.filters.end()) {
+	// Filters must match ART columns 1:1
+	if (filters.size() != indexed_columns.size() || filters.empty()) {
 		return false;
 	}
 
-	auto expressions = ExtractFilterExpressions(col, filter->second, storage_index.GetIndex());
-	for (const auto &filter_expr : expressions) {
-		auto scan_state = art.TryInitializeScan(*index_expr, *filter_expr);
+	// Do a compound scan if we have filter exprs bound for several columns
+	if (filters.size() > 1) {
+		auto scan_state = art.TryInitializeCompoundKeyScan(index_exprs, filters);
 		if (!scan_state) {
 			return false;
 		}
 
-		// Check if we can use an index scan, and already retrieve the matching row ids.
-		if (!art.Scan(*scan_state, max_count, row_ids)) {
+		if (!art.CompoundKeyScan(*scan_state, max_count, row_ids)) {
 			row_ids.clear();
 			return false;
 		}
 	}
+	// Original single column index scan
+	else {
+		for (const auto &filter_expr : filters[0]) {
+			auto scan_state = art.TryInitializeScan(*index_exprs[0], *filter_expr);
+			if (!scan_state) {
+				return false;
+			}
+
+			// Check if we can use an index scan, and already retrieve the matching row ids.
+			if (!art.Scan(*scan_state, max_count, row_ids)) {
+				row_ids.clear();
+				return false;
+			}
+		}
+	}
+
 	return true;
 }
 
@@ -588,9 +627,9 @@ unique_ptr<GlobalTableFunctionState> TableScanInitGlobal(ClientContext &context,
 	//		1.2. Find + scan one ART for b = 24.
 	//		1.3. Return the intersecting row IDs.
 	// 2. (Reorder and) scan a single ART with a compound key of (a, b).
-	if (filter_set.filters.size() != 1) {
-		return DuckTableScanInitGlobal(context, input, storage, bind_data);
-	}
+	// if (filter_set.filters.size() != 1) {
+	// 	return DuckTableScanInitGlobal(context, input, storage, bind_data);
+	// }
 
 	// The checkpoint lock ensures that we do not checkpoint while scanning this table.
 	auto &transaction = DuckTransaction::Get(context, storage.db);

src/include/duckdb/execution/index/art/art.hpp

@@ -11,6 +11,7 @@
 #include "duckdb/execution/index/bound_index.hpp"
 #include "duckdb/execution/index/art/node.hpp"
 #include "duckdb/common/array.hpp"
+#include "duckdb/planner/expression.hpp"
 
 namespace duckdb {
 
@@ -24,6 +25,7 @@ class ARTKeySection;
 class FixedSizeAllocator;
 
 struct ARTIndexScanState;
+struct ARTIndexCompoundKeyScanState;
 
 class ART : public BoundIndex {
 public:
@@ -70,10 +72,20 @@ class ART : public BoundIndex {
 public:
 	//! Try to initialize a scan on the ART with the given expression and filter.
 	unique_ptr<IndexScanState> TryInitializeScan(const Expression &expr, const Expression &filter_expr);
+
+	//! Try to initialize a compound key scan on the ART, using the given index expr -> filter expr mappings.
+	//! Supports equality comparisons only.
+	unique_ptr<IndexScanState> TryInitializeCompoundKeyScan(const vector<unique_ptr<Expression>> &index_exprs,
+	                                                        vector<vector<unique_ptr<Expression>>> &exprs);
+
 	//! Perform a lookup on the ART, fetching up to max_count row IDs.
 	//! If all row IDs were fetched, it return true, else false.
 	bool Scan(IndexScanState &state, idx_t max_count, set<row_t> &row_ids);
 
+	//! Like `ART::Scan`, but uses `ARTIndexCompoundKeyScanState` to concatenate multiple
+	//! values for equality comparisons only.
+	bool CompoundKeyScan(IndexScanState &state, idx_t max_count, set<row_t> &row_ids);
+
 	//! Appends data to the locked index.
 	ErrorData Append(IndexLock &l, DataChunk &chunk, Vector &row_ids) override;
 	//! Appends data to the locked index and verifies constraint violations.