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

ART Index: Support compound key scans

Author: mach-kernel

src/execution/index/art/art.cpp

@@ -1,5 +1,8 @@
 #include "duckdb/execution/index/art/art.hpp"
 
+#include "duckdb/common/assert.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 +42,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 +156,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 +717,30 @@ 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>();
+
+	if (scan_state.values.size() != types.size()) {
+		return false;
+	}
+
+	for (idx_t i = 0; i < scan_state.values.size(); ++i) {
+		D_ASSERT(scan_state.values[i].type().InternalType() == types[i]);
+	}
+
+	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,134 @@ 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;
 	}
 
+	// ...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);
+			}
+		});
+
+		if (referenced_columns.size() != 1 || *referenced_columns.begin() != indexed_columns[i]) {
+			return false;
+		}
+	}
+
 	// Resolve bound column references in the index_expr against the current input projection
-	column_t updated_index_column;
+	vector<column_t> index_column_to_proj_pos;
+	index_column_to_proj_pos.resize(indexed_columns.size(), std::numeric_limits<idx_t>::max());
+
 	bool found_index_column_in_input = false;
 
-	// 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;
+	// 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;
+			}
 		}
 	}
 
-	// If found, update the bound column ref within index_expr
-	if (found_index_column_in_input) {
+	if (!found_index_column_in_input) {
+		return false;
+	}
+
+	for (auto col : index_column_to_proj_pos) {
+		if (col == std::numeric_limits<idx_t>::max()) {
+			return false;
+		}
+	}
+
+	// Update the bound column refs within all index_exprs
+	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 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;
+			// 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;
+				}
 			}
 		});
 	}
 
-	// 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;
-		}
-	}
+	// Reuse the index <-> projection mappings from index expr rebinding
+	vector<vector<unique_ptr<Expression>>> index_filters;
 
-	// No filter matches the ART column.
-	if (!storage_index.IsValid()) {
-		return false;
+	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]);
+
+			index_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()) {
+	// Index filters must:
+	// - Match ART column count 1:1
+	// - Match filter expression set 1:1 (there may be filters on non-indexed columns, bail out if so)
+	if (index_filters.size() != indexed_columns.size() || filter_set.filters.size() != index_filters.size() ||
+	    index_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 (index_filters.size() > 1) {
+		auto scan_state = art.TryInitializeCompoundKeyScan(index_exprs, index_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 : index_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 +638,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.

test/sql/index/art/scan/test_art_composite_key_scan.test

@@ -0,0 +1,25 @@
+# name: test/sql/index/art/scan/test_art_composite_key_scan.test
+# description: Test Index Scan push down against views with reordered projections (issue #17290)
+# group: [scan]
+
+statement ok
+create or replace table test as (
+    select
+        cast(unnest(range(1000)) as varchar) as x,
+        cast(unnest(range(2000,3000)) as varchar) as y,
+        cast(unnest(range(3000,4000)) as varchar) as z
+);
+
+# test simple permutation of initial table projection
+statement ok
+create index test_composite_key on test(x, y, z);
+
+query II
+explain analyze select * from test where x = '525' and y = '2525' and z = '3525';
+----
+analyzed_plan	<REGEX>:.*Index Scan.*
+
+query III
+select * from test where x = '525' and y = '2525' and z = '3525';
+----
+525	2525	3525