Exact K-Nearest Neighbors
Deprecated C++ edition: This chapter belongs to the 2024 BusTub course. It is frozen to benchmark snapshot
b979953
and is not kept compatible with newer BusTub versions. New course development follows the
Rust course.
This chapter turns an ordinary table scan into exact k-nearest-neighbor search. First you will implement general-purpose sort and limit executors. Then you will replace that pair with a bounded Top-N executor.
Complete Vector Expressions and Storage first. You will likely modify these private BusTub assignment files:
src/execution/sort_executor.cpp (KEEP PRIVATE)
src/execution/topn_executor.cpp (KEEP PRIVATE)
src/execution/limit_executor.cpp (KEEP PRIVATE)
src/include/execution/executors/sort_executor.h (KEEP PRIVATE)
src/include/execution/executors/topn_executor.h (KEEP PRIVATE)
src/include/execution/executors/limit_executor.h (KEEP PRIVATE)
src/optimizer/sort_limit_as_topn.cpp (KEEP PRIVATE)
These files overlap with CMU’s Database Systems assignments. KEEP PRIVATE means that you must add these paths to your
solution repository’s .gitignore and must not commit or publish them. The starter already tracks placeholder versions,
so adding them to .gitignore inside the starter clone is not enough to hide your changes. Check the staged diff before
publishing; other course files are not part of this restriction.
The Query
CREATE TABLE t1(v1 VECTOR(3), v2 integer);
SELECT v1 FROM t1 ORDER BY ARRAY [1.0, 1.0, 1.0] <-> v1 LIMIT 3;
This is an exact k-nearest-neighbor query with k = 3. [1.0, 1.0, 1.0] is the query vector, <-> computes its L2
distance to each stored v1, ORDER BY ranks rows from smallest distance to largest, and LIMIT 3 keeps the three
nearest vectors. Because this chapter has not introduced an approximate index, the query computes every row’s distance.
Before the Top-N rewrite, run the following statement in bustub-shell:
EXPLAIN (o) SELECT v1 FROM t1 ORDER BY ARRAY [1.0, 1.0, 1.0] <-> v1 LIMIT 3;
Its plan has this shape:
Limit { limit=3 }
Sort { order_bys=[("Default", "l2_dist([1.000000,1.000000,1.000000], #0.0)")] }
Projection { exprs=["#0.0"] }
SeqScan { table=t1 }
#0.0 means column 0 from child 0. Evaluate each order-by expression against the child tuple and the child’s output
schema.
Checkpoint 1: Sort and Limit
The sort executor is a pipeline breaker: Init consumes and stores every (Tuple, RID) from its child, then sorts the
stored entries. Next emits them one at a time. Keep the RID paired with its tuple throughout the sort.
Course rules:
- The required vector query has one non-null distance key in ascending or default order. Broader SQL sorting semantics are outside this course’s scope.
- Preserve any order among complete ties; the vector reference allows tied rows to appear in either order.
The limit executor initializes its child and forwards at most limit entries. It must handle limit = 0 and a child with
fewer rows without pulling or emitting an extra tuple.
From bustub-vectordb/build, run:
make -j8 sqllogictest
./bin/bustub-sqllogictest ../test/sql/vector.02-naive-knn.slt --verbose
The vector file exercises all three distance functions. Compare its exact-query rows with the first reference output.
Sort + Limit Reference
<main>:1
CREATE TABLE t1(v1 VECTOR(3), v2 integer);
----
Table created with id = 24
<main>:4
INSERT INTO t1 VALUES (ARRAY [-1.0, 1.0, 1.0], -1), (ARRAY [-2.0, 1.0, 1.0], -2), (ARRAY [-3.0, 1.0, 1.0], -3), (ARRAY [-4.0, 1.0, 1.0], -4), (ARRAY [1.0, 1.0, 1.0], 1), (ARRAY [2.0, 1.0, 1.0], 2), (ARRAY [3.0, 1.0, 1.0], 3), (ARRAY [4.0, 1.0, 1.0], 4);
----
0
<main>:7
EXPLAIN (o) SELECT v1 FROM t1 ORDER BY ARRAY [1.0, 1.0, 1.0] <-> v1 LIMIT 3;
----
=== OPTIMIZER ===
Limit { limit=3 }
Sort { order_bys=[("Default", "l2_dist([1.000000,1.000000,1.000000], #0.0)")] }
Projection { exprs=["#0.0"] }
SeqScan { table=t1 }
<main>:10
EXPLAIN (o) SELECT * FROM (SELECT v1, ARRAY [0.5, 1.0, 1.0] <-> v1 as distance FROM t1) ORDER BY distance LIMIT 3;
----
=== OPTIMIZER ===
Limit { limit=3 }
Sort { order_bys=[("Default", "#0.1")] }
Projection { exprs=["#0.0", "l2_dist([0.500000,1.000000,1.000000], #0.0)"] }
SeqScan { table=t1 }
<main>:13
EXPLAIN (o) SELECT * FROM (SELECT v1, ARRAY [0.5, 1.0, 1.0] <=> v1 as distance FROM t1) ORDER BY distance LIMIT 3;
----
=== OPTIMIZER ===
Limit { limit=3 }
Sort { order_bys=[("Default", "#0.1")] }
Projection { exprs=["#0.0", "cosine_similarity([0.500000,1.000000,1.000000], #0.0)"] }
SeqScan { table=t1 }
<main>:16
EXPLAIN (o) SELECT * FROM (SELECT v1, inner_product(ARRAY [0.5, 1.0, 1.0], v1) as distance FROM t1) ORDER BY distance LIMIT 3;
----
=== OPTIMIZER ===
Limit { limit=3 }
Sort { order_bys=[("Default", "#0.1")] }
Projection { exprs=["#0.0", "inner_product([0.500000,1.000000,1.000000], #0.0)"] }
SeqScan { table=t1 }
<main>:19
SELECT v1 FROM t1 ORDER BY ARRAY [1.0, 1.0, 1.0] <-> v1 LIMIT 3;
----
[1,1,1]
[2,1,1]
[-1,1,1]
<main>:22
SELECT * FROM (SELECT v1, ARRAY [0.5, 1.0, 1.0] <-> v1 as distance FROM t1) ORDER BY distance LIMIT 3;
----
[1,1,1] 0.500000
[-1,1,1] 1.500000
[2,1,1] 1.500000
<main>:25
SELECT * FROM (SELECT v1, ARRAY [0.5, 1.0, 1.0] <=> v1 as distance FROM t1) ORDER BY distance LIMIT 3;
----
[1,1,1] 0.037750
[2,1,1] 0.183503
[3,1,1] 0.296474
<main>:28
SELECT * FROM (SELECT v1, inner_product(ARRAY [0.5, 1.0, 1.0], v1) as distance FROM t1) ORDER BY distance LIMIT 3;
----
[4,1,1] -4.000000
[3,1,1] -3.500000
[2,1,1] -3.000000
Checkpoint 2: Bounded Top-N
Sorting all n rows costs O(n log n) and stores all n entries. For LIMIT k, a max-heap can retain only the best k
entries in O(n log k) time and O(k) space.
First implement OptimizeSortLimitAsTopN. It should replace only a Limit whose direct child is a Sort, copy the sort’s
order-by list and the limit into a TopNPlanNode, and preserve the sort’s child.
Then implement TopNExecutor:
- initialize the child;
- evaluate the same full ordering used by
SortExecutor; - keep at most
kbest(Tuple, RID)entries in a max-heap, with the worst retained entry at the top; and - emit the retained entries in final best-to-worst order.
Popping a max-heap directly produces the worst retained row first. Reverse that sequence, or use another equivalent
method, before Next begins emitting. Keep the retained container bounded to k entries.
Prediction: If the input distances are 4, 1, 3, 2 and k = 2, which values remain after each input? The final output
must be 1, 2, even though the heap’s top is 2.
Run:
./bin/bustub-sqllogictest ../test/sql/vector.02-naive-knn.slt --verbose
The EXPLAIN output should now contain TopN instead of Limit over Sort, and its query rows should match the exact
checkpoint apart from allowed tie ordering.
Top-N Reference
<main>:1
CREATE TABLE t1(v1 VECTOR(3), v2 integer);
----
Table created with id = 24
<main>:4
INSERT INTO t1 VALUES (ARRAY [-1.0, 1.0, 1.0], -1), (ARRAY [-2.0, 1.0, 1.0], -2), (ARRAY [-3.0, 1.0, 1.0], -3), (ARRAY [-4.0, 1.0, 1.0], -4), (ARRAY [1.0, 1.0, 1.0], 1), (ARRAY [2.0, 1.0, 1.0], 2), (ARRAY [3.0, 1.0, 1.0], 3), (ARRAY [4.0, 1.0, 1.0], 4);
----
0
<main>:7
EXPLAIN (o) SELECT v1 FROM t1 ORDER BY ARRAY [1.0, 1.0, 1.0] <-> v1 LIMIT 3;
----
=== OPTIMIZER ===
TopN { n=3, order_bys=[("Default", "l2_dist([1.000000,1.000000,1.000000], #0.0)")]}
Projection { exprs=["#0.0"] }
SeqScan { table=t1 }
<main>:10
EXPLAIN (o) SELECT * FROM (SELECT v1, ARRAY [0.5, 1.0, 1.0] <-> v1 as distance FROM t1) ORDER BY distance LIMIT 3;
----
=== OPTIMIZER ===
TopN { n=3, order_bys=[("Default", "#0.1")]}
Projection { exprs=["#0.0", "l2_dist([0.500000,1.000000,1.000000], #0.0)"] }
SeqScan { table=t1 }
<main>:13
EXPLAIN (o) SELECT * FROM (SELECT v1, ARRAY [0.5, 1.0, 1.0] <=> v1 as distance FROM t1) ORDER BY distance LIMIT 3;
----
=== OPTIMIZER ===
TopN { n=3, order_bys=[("Default", "#0.1")]}
Projection { exprs=["#0.0", "cosine_similarity([0.500000,1.000000,1.000000], #0.0)"] }
SeqScan { table=t1 }
<main>:16
EXPLAIN (o) SELECT * FROM (SELECT v1, inner_product(ARRAY [0.5, 1.0, 1.0], v1) as distance FROM t1) ORDER BY distance LIMIT 3;
----
=== OPTIMIZER ===
TopN { n=3, order_bys=[("Default", "#0.1")]}
Projection { exprs=["#0.0", "inner_product([0.500000,1.000000,1.000000], #0.0)"] }
SeqScan { table=t1 }
<main>:19
SELECT v1 FROM t1 ORDER BY ARRAY [1.0, 1.0, 1.0] <-> v1 LIMIT 3;
----
[1,1,1]
[2,1,1]
[3,1,1]
<main>:22
SELECT * FROM (SELECT v1, ARRAY [0.5, 1.0, 1.0] <-> v1 as distance FROM t1) ORDER BY distance LIMIT 3;
----
[1,1,1] 0.500000
[2,1,1] 1.500000
[-1,1,1] 1.500000
<main>:25
SELECT * FROM (SELECT v1, ARRAY [0.5, 1.0, 1.0] <=> v1 as distance FROM t1) ORDER BY distance LIMIT 3;
----
[1,1,1] 0.037750
[2,1,1] 0.183503
[3,1,1] 0.296474
<main>:28
SELECT * FROM (SELECT v1, inner_product(ARRAY [0.5, 1.0, 1.0], v1) as distance FROM t1) ORDER BY distance LIMIT 3;
----
[4,1,1] -4.000000
[3,1,1] -3.500000
[2,1,1] -3.000000
You are done when you can explain why changing the Top-N heap from a max-heap to a min-heap would retain the wrong end of the ordering, and how the optimizer rewrite preserves the original plan’s result.
Related lecture: Query Planning & Optimization (CMU Intro to Database Systems)
Optional Extension
Extend vector construction to accept mixed integer and decimal array literals, or a cast such as
'[1.0, 1.0, 1.0]'::VECTOR(3).
Your feedback is greatly appreciated. Join our Discord community.
Found an issue? Open an issue or pull request at github.com/skyzh/write-you-a-vector-db.
write-you-a-vector-db © 2024-2026 by Alex Chi Z. All Rights Reserved.